-- 
-- (c) Susumu Katayama 2009
--

\begin{code}
{-# OPTIONS -fglasgow-exts -cpp #-}
module MagicHaskeller.ProgGenXF(ProgGenXF) where
import MagicHaskeller.Types
import MagicHaskeller.TyConLib
import Control.Monad
import MagicHaskeller.CoreLang
import Control.Monad.Search.Combinatorial
import MagicHaskeller.PriorSubsts
import Data.List(partition, sortBy, sort, nub, (\\))
import Data.Ix(inRange)

import MagicHaskeller.ClassifyDM

import System.Random(mkStdGen, StdGen)
import MagicHaskeller.Instantiate

import MagicHaskeller.ProgramGenerator

import MagicHaskeller.Expression

import MagicHaskeller.ClassifyTr
import MagicHaskeller.MyDynamic

import MagicHaskeller.T10(mergesortWithBy)

import MagicHaskeller.DebMT

import Debug.Trace
-- trace str = id

reorganize_ = reorganizer_
-- reorganize_ = id

classify = True
traceExpTy _ = id
-- traceExpTy fty = trace ("lookupexp "++ show fty)
traceTy _    = id
-- traceTy fty = trace ("lookup "++ show fty)

-- Memoization table, created from primitive components
newtype ProgGenXF = PGXF MemoDeb -- internal data representation.
-- ^ Program generator with experimental leveraged filtration, experimental variant of ProgGenSF. The following is explanation of ProgGenSF.
--   This program generator employs filtration by random testing, and rarely generate semantically equivalent expressions more than once, while different expressions will eventually appear (for most of the types, represented with Prelude types, whose arguments are instance of Arbitrary and which return instance of Ord).
--   The idea is to apply random numbers to the generated expressions, compute the quotient set of the resulting values at each depth of the search tree, and adopt the complete system of representatives for the depth and push the remaining expressions to one step deeper in the search tree.
--   (Thus, adoption of expressions that may be equivalent to another already-generated-expression will be postponed until their \"uniqueness\" is proved.)
--   As a result, (unlike "ProgGen",) expressions with size N may not appear at depth N but some deeper place.
--
--   "ProgGenSF" is more efficient along with a middle-sized primitive set (like @reallyall@ found in LibTH.hs),
--   but is slower than "ProgGen" for a small-sized one.
--
--   Also note that "ProgGenSF" depends on hard use of unsafe* stuff, so if there is a bug, it may segfault....

type ExpTip   = (Stream (Forest ([Dynamic], AnnExpr)), Stream (Forest ([Dynamic], AnnExpr)), Matrix AnnExpr)

type ExpTrie  = MapType ExpTip
type TypeTrie = MapType (Matrix (Matrix AnnExpr, Subst, Int))

type MemoTrie = (TypeTrie, ExpTrie)


lmt :: (Int, MemoDeb) -> Type -> Matrix AnnExpr
lmt (_,memoDeb@((_,mt),_,cmn)) fty =
  let (_,_,x) = traceExpTy fty $
                lookupMT mt fty -- こっちだとlookup
--              filtBF cmn fty $ matchFunctions (maxBound', memoDeb) fty --  こっちだとrecompute
  in x
--          filtBF ty = fmap fromAnnExpr . filterBF tcl rtrie ty . fmap (toAnnExprWind (execute opt) ty) . tabulate
filtBF cmn ty | classify  = filterTr cmn ty . fmap (toAnnExprWind (execute (opt cmn) (vl cmn)) ty)  -- .  mapDepth aUS
              | otherwise = \bf -> (undefined, undefined, toMx $ mapDepth aUS bf)

lmtty mt fty = traceTy fty $
               lookupMT mt fty

--memocond i = 3<i
-- memocond i = i<10
memocond i = True

-- memocond ty = size ty < 10 -- popArgs avail tしてからやるのはちょっと無駄．と思ったけど，実際は関係ないみたい．
-- memocond av ty = size ty + sum (map size av) < 10


instance ProgramGenerator ProgGenXF where
    mkTrieOpt cmn tcesopt tces = PGXF (mkTrieOptSF cmn tcesopt tces)
    matchingPrograms ty (dep, PGXF x)      = fromMx $ matchProgs (dep, x) ty
    unifyingPrograms ty (dep, px@(PGXF x)) = fromRc $ catBags $ fmap (\ (es,_,_) -> map (toAnnExpr $ reducer px) es) $ unifyingPossibilities ty (dep, x)
    extractCommon    (PGXF (_,_,cmn))      = cmn

unifyingPossibilities ty memodeb = unPS (unifyableExprs memodeb [] ty) emptySubst 0

-- quantifyされたやつがmemoizeされている訳だから，queryのentryではquantifyする必要はない
-- entry for query
-- てゆーか，quantifyしちゃうならMemoDeb.mguProgramsかなんかそのまま使えばいいって話もある．まあ，recursiveには上記のunifyableExprsを呼ばなきゃダメだけど．
-- なお，厳密な意味でmatchにするにはどうもquantifyは必要っぽい．
--matchProgs :: (Int, Memo) -> Type -> BF AnnExpr
matchProgs :: (Int,MemoDeb) -> Type -> Matrix AnnExpr
matchProgs memodeb ty = fmap (toAnnExprWindWind (reducer $ PGXF $ snd memodeb) ty) $ -- fmap meToAE $           --   koko
                                                                                     lmt memodeb $ normalize $ unquantify ty -- こっちだとlookup
{-
matchProgs memodeb ty = fmap toAnnExpr $ wind (fmap (mapCE Lambda)) (lookupFuns memodeb) [] (quantify ty)                 -- こっちだとrecompute というと語弊がある．recomputeしたきゃlmtのところを変えるべし．

-- matchProgsのみの下請け，matchFunsと交換可能
lookupFuns :: (Expression e, Ord e) => (Int, MemoDeb e) -> [Type] -> Type -> BF e
lookupFuns memodeb@(memodepth,((_,mt),_,tcl,rtrie)) avail reqret =
{-
#ifdef CLASSIFY
                                  fmap fromAnnExpr $ toRc $ filterDM tcl rtrie ty $ fromRc $ fmap (toAnnExprWind ty) $
#endif
-}
--                                  mapDepth uniqSort $
                                             matchFuns memodeb avail reqret
    where ty = popArgs avail reqret
-}

specializedPossibleTypes :: Type -> (Int, MemoDeb) -> Recomp Type
specializedPossibleTypes ty memodeb = runPS (fmap (\(av,t) -> popArgs av t) $ specializedTypes memodeb [] ty)
-- specializedPossibleTypes ty memodeb@(_,((mt,_),_,_,_)) = fmap (\(_,s,_) -> apply s ty) $ toRc $ lmtty mt ty


type MemoDeb = (MemoTrie, (([Prim],[Prim]),([Prim],[Prim])), Common)
-- TyConLibは[Typed [CoreExpr]]から自前で作るべき．場合によってはLinsCCLとかとTyConLibを共有できなくなるかも知れないけど，それはそれでOK．か？ やっぱ自前で作るのはLIBRARYのケースのみにしておくか．
-- あ，てゆーか，[Typed [CoreExpr]]を作るのにTyConLibが必要．
-- むしろ，[([CoreExpr],TypeRep)]からTyConLibと[Typed [CoreExpr]]を作る感じで．

 -- maxBound使うと多分効率悪いけど，まあ面倒だしいいか．
maxBound' = maxBound -- Setting this to some small value can sometimes be helpful when seeing the heap behavior.
mkTrieOptSF :: Common -> [Typed [CoreExpr]] -> [Typed [CoreExpr]] -> MemoDeb
mkTrieOptSF cmn txsopt txs
    = let
          memoDeb = (memoTrie, (qtlopt,qtl), cmn)
          -- memoTrie :: MemoTrie
          memoTrie = (typeTrie,expTrie)
          typeTrie = mkMTty (tcl cmn) (\ty -> freezePS ty (specTypes (maxBound', memoDeb) ty))
          expTrie = mkMTexp (tcl cmn) (\ty -> -- fmap (aeToME (tcl cmn) (rt cmn) ty) $
                                              filtBF cmn ty $ toMx $ matchFunctions (maxBound', memoDeb) ty)
      in memoDeb
    where qtlopt = splitPrims txsopt
          qtl    = splitPrims txs

mkMTty = mkMT
mkMTexp = mkMT

mondepth = zipDepthRc (\d xs -> trace ("depth="++show d++", and the length is "++show (length xs)) xs) -- depthと表示するなら+1するべきであった．(0から始まるので)


type BFT = Recomp
unBFM = unMx


{-
freezePS :: (Search m, Expression e) => Type -> PriorSubsts m (ExpTip e) -> Matrix (ExpTip e,Subst,Int)
freezePS ty ps
    = let mxty = maxVarID ty -- `max` maximum (map maxVarID avail)
      in Mx $ map (tokoro10ap ty) $ scanl1 (++) $ unMx $ toMx $ unPS ps emptySubst (mxty+1)
-}
freezePS :: Type -> PriorSubsts Recomp e -> Matrix (e,Subst,Int)
freezePS ty ps
    = let mxty = maxVarID ty -- `max` maximum (map maxVarID avail)
      in mapDepth (tokoro10ap ty) $ toMx $ unPS ps emptySubst (mxty+1)

-- MemoStingy.tokoro10 is different from T10.tokoro10, in that duplicates will be removed.
-- (Note that the type can be specialized to [(Type,k,i)] -> [(Type,k,i)])
tokoro10 :: (Eq k, Ord k) => [(a,k,i)] -> [(a,k,i)]
tokoro10 = mergesortWithBy const (\ (_,k,_) (_,l,_) -> k `compare` l)

-- tokoro10fstfst = mergesortWithBy const (\ ((k,_),_,_) ((l,_),_,_) -> k `compare` l)

tokoro10ap :: Type -> [(a,Subst,i)] -> [(a,Subst,i)]
tokoro10ap ty = mergesortWithBy const (\ (_,k,_) (_,l,_) -> normalize (apply k ty) `compare` normalize (apply l ty))

-- availにしろTypeにしろapplyされている．
-- だからこそ，runAnotherPS的にemptySubstに対して実行した方が効率的なはず？ でも，Substitutionってそんなにでかくならなかったのでは？FiniteMapでもassoc listでも変わらなかった気が．


specializedTypes :: (Search m) => (Int, MemoDeb) -> [Type] -> Type -> PriorSubsts m ([Type],Type)
specializedTypes memodeb avail t = do specializedCases memodeb avail t
                                      subst <- getSubst
                                      return (map (apply subst) avail, apply subst t)
-- specializedCases is the same as unifyableExprs, except that the latter returns PriorSubsts BF [CoreExpr], and that the latter considers memodepth.
specializedCases, specCases, specCases' :: (Search m) => (Int, MemoDeb) -> [Type] -> Type -> PriorSubsts m ()
specializedCases memodeb = applyDo (specCases memodeb)
specCases (_,memodeb) = wind_ (\avail reqret -> reorganize_ (\newavail -> uniExprs_ (maxBound',memodeb) newavail reqret) avail)
{- どっちがわかりやすいかは不明
specCases memodeb avail (t0:->t1) = specCases memodeb (t0 : avail) t1
specCases (_,memodeb) avail reqret = reorganize_ (\newavail -> uniExprs_ (maxBound',memodeb) newavail reqret) avail
-}





uniExprs_ :: (Search m) => (Int, MemoDeb) -> [Type] -> Type -> PriorSubsts m ()
uniExprs_ memodeb avail t
    = convertPS fromRc $ psListToPSRecomp lfp
    where lfp depth
              | memocond depth     = lookupUniExprs memodeb avail t depth >> return ()
              | otherwise          = makeUniExprs memodeb avail t depth >> return ()

lookupUniExprs :: (Int, MemoDeb) -> [Type] -> Type -> Int -> PriorSubsts [] (Matrix AnnExpr)
lookupUniExprs memodeb@(_,((mt,_),_,_)) avail t depth
    = lookupNormalized  (\tn -> unMx (lmtty mt tn) !! depth) avail t

makeUniExprs :: (Int, MemoDeb) -> [Type] -> Type -> Int -> PriorSubsts [] Type
makeUniExprs memodeb avail t depth
    = convertPS tokoro10fst $
                do psRecompToPSList (reorganize_ (\av -> specCases' memodeb av t) avail) depth
                   sub   <- getSubst
                   return $ quantify (apply sub $ popArgs avail t)

{-
makeUniExprs_ memodeb@(_,(_, _, tcl, rtrie)) avail t depth
    = t10PS (popArgs avail t) $ psRecompToPSList (specCases' memodeb avail t) depth
t10PS :: Type -> PriorSubsts [] a -> PriorSubsts [] ()
t10PS ty ps
    = do convertPS tokoro10fst $
                do ps
                   sub   <- getSubst
                   return (apply sub ty)
         return ()
-- ここでは同じ型になるものをまとめている訳だが，
-- - ここでまとめた方が速いのか，そんなことをせずに単にpsRecompToPSList (specCases' memodeb avail t) depthを使う方が速いのか
-- - typetrieにmemoizeするときもちゃんとまとめているのか，
-- 調べるべし．特に，typetrieでまとめていないからヒープが増えているのかも．
-}

-- entry point for memoization
specTypes :: (Search m) => (Int, MemoDeb) -> Type -> PriorSubsts m (Matrix AnnExpr)
specTypes memodeb@(_,((_,mt),_,_)) ty
                           = do let (avail,t) = splitArgs ty
                                reorganize_ (\av -> specCases' memodeb av t) avail
-- quantifyはmemo先で既にやられているので不要
                                typ <- applyPS ty
                                return (lmt memodeb $ normalize $ unquantify typ)


funApSub_ :: Search m => (Type -> PriorSubsts m ()) -> (Type -> PriorSubsts m ()) -> Type -> PriorSubsts m ()
funApSub_ lltbehalf behalf (t:>ts)  = do lltbehalf t
				         funApSub_ lltbehalf behalf ts
funApSub_ lltbehalf behalf (t:->ts) = do behalf t
				         funApSub_ lltbehalf behalf ts
funApSub_ lltbehalf behalf _t       = return ()

funApSub_spec behalf = funApSub_ behalf behalf

-- specCases' trie prims@(primgen,primmono) avail reqret = msum (map (retMono.fromPrim) primmono) `mplus` msum (map retMono fromAvail ++ map retGen primgen)
specCases' memodeb@(_,((ttrie,etrie), (prims@(primgen,primmono),_),cmn)) avail reqret
 = msum (map retPrimMono primmono ++ map retMono avail ++ map retGen primgen)
    where fas | constrL $ opt cmn = funApSub_ lltbehalf behalf
              | otherwise         = funApSub_spec       behalf
              where behalf    = specializedCases memodeb avail
                    lltbehalf = flip mguAssumptions_ avail
          -- retPrimMono :: (Int, Type, Int, Typed [CoreExpr]) -> PriorSubsts BFT ()
          retPrimMono (arity, retty, numtvs, _xs:::ty)
                                              = napply arity delayPS $
                                                do tvid <- reserveTVars numtvs
                                                   mguPS reqret (mapTV (tvid+) retty)
                                                   fas (mapTV (tvid+) ty)
          -- retMono :: Type -> PriorSubsts BFT ()
          retMono ty = napply (getArity ty) delayPS $
                       do mguPS reqret (getRet ty)
		          fas ty
          -- retGen :: (Int, Type, Int, Typed [CoreExpr]) -> PriorSubsts BFT ()
          retGen (arity, _r, numtvs, _s:::ty) = napply arity delayPS $
                                             do tvid <- reserveTVars numtvs -- この（最初の）IDそのもの（つまり返り値のtvID）はすぐに使われなくなる
                                               -- let typ = apply (unitSubst tvid reqret) (mapTV (tvid+) ty) -- mapTVとapplyはhylo-fusionできるはずだが，勝手にされる？
                                               --                                                              -- unitSubstをinlineにしないと駄目か
                                                mkSubsts tvid reqret
                                                fas (mapTV (tvid+) ty)

	                                        gentvar <- applyPS (TV tvid)

                                                guard (orderedAndUsedArgs gentvar)
                                                fas gentvar

type Generator m e = (Int,MemoDeb) -> [Type] -> Type -> PriorSubsts m [e]

unifyableExprs ::  Generator Recomp AnnExpr
unifyableExprs memodeb avails ty = unifyableExprs' memodeb avails ty
unifyableExprs' :: Generator Recomp AnnExpr
unifyableExprs' memodeb = applyDo (wind (fmap (map (mapCE Lambda))) (lookupNormalized (lookupTypeTrie memodeb)))


lookupTypeTrie :: (Int, MemoDeb) -> Type -> Recomp ([AnnExpr], Subst, Int)
lookupTypeTrie memodeb@(_, ((mt,_), _, _)) t
    = Rc $ 
      \depth -> let Mx xss = lmtty mt t
                in [ ({-map meToAE-} (yss!!depth), s, i) | (Mx yss, s, i) <- xss !! depth ]


lookupNormalized :: MonadPlus m => (Type -> m (e, Subst, Int)) ->  [Type] -> Type -> PriorSubsts m e
lookupNormalized fun avail t
    = do mx <- getMx
         let typ = popArgs avail t
             (tn, decoder) = encode typ mx
         (es, sub, m) <-  mkPS (fun tn)
         updatePS (retrieve decoder sub)
         updateMx (m+)
         return es

tokoro10fst :: (Eq k, Ord k) => [(k,s,i)] -> [(k,s,i)]
tokoro10fst = mergesortWithBy const (\ (k,_,_) (l,_,_) -> k `compare` l)

-- entry for memoization
matchFunctions :: (Int, MemoDeb) -> Type -> Recomp AnnExpr
matchFunctions memodeb ty = case splitArgs (quantify ty) of (avail,t) -> matchFuns memodeb avail t

matchFuns :: (Int,MemoDeb) -> [Type] -> Type -> Recomp AnnExpr
matchFuns memodeb avail reqret = catBags $ runPS (-- trace "matchFuns'" $ 
                                                  matchFuns' unifyableExprs memodeb avail reqret)

matchFuns' :: Generator Recomp AnnExpr -> Generator Recomp AnnExpr
-- matchFuns' = generateFuns matchPS filtExprs lookupListrie -- MemoDebの型の違いでこれはうまくいかなんだ．
matchFuns' rec memodeb@(_, md@(_, (_,(primgen,primmono)),cmn)) avail reqret
    = let behalf    = rec memodeb avail
          lltbehalf = lookupListrie lenavails rec memodeb avail -- heuristic filtration
          lenavails = length avail
--          fe :: Type -> Type -> [CoreExpr] -> [CoreExpr] -- ^ heuristic filtration
          fe        = filtExprs (guess $ opt cmn)
      in fromAssumptions (PGXF md) lenavails behalf matchPS reqret avail `mplus`
         msum (map (retPrimMono (PGXF md) lenavails lltbehalf behalf matchPS reqret) primmono
               ++ map ((    if tv0 $ opt cmn then retGenTV0 else
                        if tv1 $ opt cmn then retGenTV1
                                         else retGenOrd) (PGXF md) lenavails fe lltbehalf behalf reqret) primgen)

lookupListrie :: (Search m) => Int -> Generator m AnnExpr -> Generator m AnnExpr
lookupListrie lenavails rec memodeb@(_,md@(_,_,cmn)) avail t
                                    | constrL opts = matchAssumptions (PGXF md) lenavails t avail
                                    | guess opts = do
                                       args <- rec memodeb avail t
                                       let args' = filter (not.isClosed.toCE) args
                                       when (null args') mzero
                                       return args'
                                    | otherwise = rec memodeb avail t
    where opts = opt cmn
filtExprs :: Expression e => Bool -> Type -> Type -> [e] -> [e]
filtExprs g a b | g         = filterExprs a b
                | otherwise = id

\end{code}
