module PGF.Generate 
         ( generateAll,         generateAllDepth
         , generateFrom,        generateFromDepth
         , generateRandom,      generateRandomDepth
         , generateRandomFrom,  generateRandomFromDepth
         , generateOntology,    generateOntologyDepth
         , prove
         ) where

import PGF.CId
import PGF.Data
import PGF.TypeCheck
import Control.Monad
import Control.Monad.State
import Control.Monad.Identity
import System.Random
import Data.Maybe(isNothing)

------------------------------------------------------------------------------
-- The API

-- | Generates an exhaustive possibly infinite list of
-- abstract syntax expressions.
generateAll :: PGF -> Type -> [Expr]
generateAll pgf ty = generateAllDepth pgf ty Nothing

-- | A variant of 'generateAll' which also takes as argument
-- the upper limit of the depth of the generated expression.
generateAllDepth :: PGF -> Type -> Maybe Int -> [Expr]
generateAllDepth pgf ty dp = generate () pgf ty dp

-- | Generates a list of abstract syntax expressions
-- in a way similar to 'generateAll' but instead of
-- generating all instances of a given type, this
-- function uses a template. 
generateFrom :: PGF -> Expr -> [Expr]
generateFrom pgf ex = generateFromDepth pgf ex Nothing

-- | A variant of 'generateFrom' which also takes as argument
-- the upper limit of the depth of the generated subexpressions.
generateFromDepth :: PGF -> Expr -> Maybe Int -> [Expr]
generateFromDepth pgf e dp = 
  [e | (_,_,e) <- snd $ runTcM (abstract pgf)
                               (generateForMetas (prove dp) e)
                               emptyMetaStore ()]

-- | Generates an infinite list of random abstract syntax expressions.
-- This is usefull for tree bank generation which after that can be used
-- for grammar testing.
generateRandom :: RandomGen g => g -> PGF -> Type -> [Expr]
generateRandom g pgf ty = generateRandomDepth g pgf ty Nothing

-- | A variant of 'generateRandom' which also takes as argument
-- the upper limit of the depth of the generated expression.
generateRandomDepth :: RandomGen g => g -> PGF -> Type -> Maybe Int -> [Expr]
generateRandomDepth g pgf ty dp = restart g (\g -> generate (Identity g) pgf ty dp)

-- | Random generation based on template
generateRandomFrom :: RandomGen g => g -> PGF -> Expr -> [Expr]
generateRandomFrom g pgf e = generateRandomFromDepth g pgf e Nothing

-- | Random generation based on template with a limitation in the depth.
generateRandomFromDepth :: RandomGen g => g -> PGF -> Expr -> Maybe Int -> [Expr]
generateRandomFromDepth g pgf e dp = 
  restart g (\g -> [e | (_,ms,e) <- snd $ runTcM (abstract pgf)
                                                 (generateForMetas (prove dp) e)
                                                 emptyMetaStore (Identity g)])

generateOntology :: RandomGen g => g -> PGF -> Type -> [(Maybe Expr, Type)] -> [Expr]
generateOntology g pgf ty args = generateOntologyDepth g pgf ty args Nothing

generateOntologyDepth :: RandomGen g => g -> PGF -> Type -> [(Maybe Expr, Type)] -> Maybe Int -> [Expr]
generateOntologyDepth g pgf ty args dp =
  restart g (\g -> [e | (_,(Ontology args' _),e) <- snd $ runTcM (abstract pgf)
                                                                 (prove dp emptyScope (TTyp [] ty) >>= checkResolvedMetaStore emptyScope)
                                                                 emptyMetaStore
                                                                 (Ontology args g),
                        all (isNothing . fst) args'])

------------------------------------------------------------------------------
-- The main generation algorithm

generate :: Selector sel => sel -> PGF -> Type -> Maybe Int -> [Expr]
generate sel pgf ty dp =
  [e | (_,ms,e) <- snd $ runTcM (abstract pgf)
                                (prove dp emptyScope (TTyp [] ty) >>= checkResolvedMetaStore emptyScope)
                                emptyMetaStore sel]


prove :: Selector sel => Maybe Int -> Scope -> TType -> TcM sel Expr
prove dp scope (TTyp env1 (DTyp hypos1 cat es1)) = do
  vs1 <- mapM (PGF.TypeCheck.eval env1) es1
  let scope' = exScope scope env1 hypos1
  (fe,TTyp env2 (DTyp hypos2 _ es2)) <- select cat scope' dp
  case dp of
    Just 0 | not (null hypos2) -> mzero
    _                          -> return ()
  (env2,args) <- mkEnv scope' env2 hypos2
  vs2 <- mapM (PGF.TypeCheck.eval env2) es2
  sequence_ [eqValue mzero suspend (scopeSize scope') v1 v2 | (v1,v2) <- zip vs1 vs2]
  es <- mapM (descend scope') args
  return (abs hypos1 (foldl EApp fe es))
  where
    suspend i c = do
      mv <- getMeta i
      case mv of
        MBound e -> c e
        MUnbound x scope tty cs -> setMeta i (MUnbound x scope tty (c:cs))

    abs []                e = e
    abs ((bt,x,ty):hypos) e = EAbs bt x (abs hypos e)

    exScope scope env []                = scope
    exScope scope env ((bt,x,ty):hypos) = 
       let env' | x /= wildCId = VGen (scopeSize scope) [] : env
                | otherwise    = env
       in exScope (addScopedVar x (TTyp env ty) scope) env' hypos

    mkEnv scope env []                = return (env,[])
    mkEnv scope env ((bt,x,ty):hypos) = do
      (env,arg) <- if x /= wildCId
                    then do i <- newMeta scope (TTyp env ty)
                            return (VMeta i (scopeEnv scope) [] : env,Right (EMeta i))
                    else return (env,Left (TTyp env ty))
      (env,args) <- mkEnv scope env hypos
      return (env,(bt,arg):args)

    descend scope (bt,arg) = do
      let dp' = fmap (flip (-) 1) dp
      e <- case arg of
             Right e  -> return e
             Left tty -> prove dp' scope tty
      e <- case bt of
             Implicit -> return (EImplArg e)
             Explicit -> return e
      return e


-- Helper function for random generation. After every
-- success we must restart the search to find sufficiently different solution.
restart :: RandomGen g => g -> (g -> [a]) -> [a]
restart g f =
  let (g1,g2) = split g
  in case f g1 of
       []     -> []
       (x:xs) -> x : restart g2 f


------------------------------------------------------------------------------
-- Selectors

instance Selector () where
  splitSelector s = (s,s)
  select cat scope dp = do
    gens <- typeGenerators scope cat
    TcM (\abstr k h -> iter k gens)
    where
      iter k []              ms s = id
      iter k ((_,e,tty):fns) ms s = k (e,tty) ms s . iter k fns ms s


instance RandomGen g => Selector (Identity g) where
  splitSelector (Identity g) = let (g1,g2) = split g
                               in (Identity g1, Identity g2)

  select cat scope dp = do
    gens <- typeGenerators scope cat
    TcM (\abstr k h -> iter k 1.0 gens)
    where
      iter k p []   ms (Identity g) = id
      iter k p gens ms (Identity g) = let (d,g')    = randomR (0.0,p) g
                                          (g1,g2)   = split g'
                                          (p',e_ty,gens') = hit d gens
                                      in k e_ty ms (Identity g1) . iter k (p-p') gens' ms (Identity g2)

      hit :: Double -> [(Double,Expr,TType)] -> (Double,(Expr,TType),[(Double,Expr,TType)])
      hit d (gen@(p,e,ty):gens)
        | d < p || null gens = (p,(e,ty),gens)
        | otherwise = let (p',e_ty',gens') = hit (d-p) gens
                      in (p',e_ty',gen:gens')


data Ontology a = Ontology [(Maybe Expr, Type)] a

instance RandomGen g => Selector (Ontology g) where
  splitSelector (Ontology args g) = let (g1,g2) = split g
                                    in (Ontology args g1, Ontology args g2)

  select cat scope dp = do
    Ontology args g <- get
    case pickArg [] cat args of
      []   -> do gens <- typeGenerators scope cat
                 TcM (\abstr k h -> iter k 1.0 gens)
      alts -> msum [  case mb_e of
                        Just e  -> do put (Ontology args g)
                                      return (e, TTyp [] ty)
                        Nothing -> mzero
                    | (mb_e,ty,args) <- alts]
    where
      iter k p []   ms (Ontology ce g) = id
      iter k p gens ms (Ontology ce g) =
        let (d,g')    = randomR (0.0,p) g
            (g1,g2)   = split g'
            (p',e_ty,gens') = hit d gens
        in k e_ty ms (Ontology ce g1) . iter k (p-p') gens' ms (Ontology ce g2)

      hit :: Double -> [(Double,Expr,TType)] -> (Double,(Expr,TType),[(Double,Expr,TType)])
      hit d (gen@(p,e,ty):gens) | d < p || null gens = (p,(e,ty),gens)
                                | otherwise          = let (p',e_ty',gens') = hit (d-p) gens
                                                       in (p',e_ty',gen:gens')

      pickArg args' cat' [] = []
      pickArg args' cat' (arg@(mb_e,ty@(DTyp _ cat _)):args)
        | cat' == cat = (mb_e, ty, foldl (flip (:)) args args') : 
                        pickArg (arg:args') cat' args
        | otherwise   = pickArg (arg:args') cat' args