-- 
-- (c) Susumu Katayama 2009
--
\begin{code}
{-# OPTIONS -cpp -fglasgow-exts #-}
module MagicHaskeller.ProgGen(ProgGen(PG)) where

import MagicHaskeller.Types
import MagicHaskeller.TyConLib
import Control.Monad
import Data.Monoid
import MagicHaskeller.CoreLang
import Control.Monad.Search.Combinatorial
import MagicHaskeller.PriorSubsts
import Data.List(partition, sortBy)
import Data.Ix(inRange)

import MagicHaskeller.ProgramGenerator

import MagicHaskeller.Classify
import System.Random(mkStdGen)
import MagicHaskeller.Instantiate

import MagicHaskeller.Expression

import MagicHaskeller.T10
import qualified Data.Map as Map

import MagicHaskeller.DebMT

import Debug.Trace

import Data.Monoid


-- x #define DESTRUCTIVE
traceTy _    = id
-- traceTy fty = trace ("lookup "++ show fty)


type BF = Recomp
-- type BF = DBound

type BFM = Matrix
-- type BFM = DBMemo
fromMemo :: Search m => Matrix a -> m a
fromMemo = fromMx
toMemo :: Search m => m a -> Matrix a
toMemo = toMx


-- Memoization table, created from primitive components

-- | The vanilla program generator corresponding to Version 0.7.*
newtype ProgGen = PG (MemoDeb CoreExpr) -- ^ internal data representation


mapStrTyCon :: MemoDeb CoreExpr -> Map.Map String TyCon
mapStrTyCon = fst . extractTCL . PG

#ifdef DESTRUCTIVE
type MemoTrie a = (MapType (BFM (Possibility a)), MapType (BFM (Possibility a)))
#else
type MemoTrie a = MapType (BFM (Possibility a))
#endif

#ifdef DESTRUCTIVE
lmt (mt,_) fty =
#else
lmt mt fty =
#endif
       traceTy fty $
       lookupMT mt fty

lookupFunsPoly :: (Search m, Expression e) => Generator m e -> Generator m e
lookupFunsPoly behalf memodeb@(memodepth,(mt,_,_)) avail reqret
    = PS (\subst mx ->
              let (tn, decoder) = encode (popArgs avail reqret) mx
              in ifDepth (<= memodepth)
                         (fmap (\ (exprs, sub, m) -> (exprs, retrieve decoder sub `plusSubst` subst, mx+m)) $ fromMemo $ lmt mt tn)
                         (unPS (behalf memodeb avail reqret) subst mx) )

instance ProgramGenerator ProgGen where
    mkTrie cmn tces = PG (mkTrieMD cmn tces)
    unifyingPrograms ty (dep, px@(PG x)) = fmap (toAnnExpr $ reducer px) $ catBags $ fmap (\ (es,_,_) -> es) $ unifyingPossibilities ty (dep, x)
    extractCommon     (PG (_,_,cmn)) = cmn

unifyingPossibilities :: Search m => Type -> (Int, MemoDeb CoreExpr) -> m ([CoreExpr],Subst,Int)
unifyingPossibilities ty memodeb = unPS (mguProgs memodeb [] ty) emptySubst 0

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

 -- maxBound使うと多分効率悪いけど，まあ面倒だしいいか．
mkTrieMD :: Common -> [Typed [CoreExpr]] -> MemoDeb CoreExpr
mkTrieMD cmn txs
    = let
          memoDeb = (memoTrie, qtl, cmn)
          -- memoTrie :: MemoTrie a
-- monomorphicなののみをmemoしたい時は，MemoMT.fps/freezePSを使うべし．
#ifdef DESTRUCTIVE
          memoTrie = (allTrie,listrie)
          listrie = mkMT (tcl cmn) (\ty -> freezePS ty (let (avail,t) = splitArgs ty in mguFuns (maxBound, (opt, (listrie, listrie), (fst qtl,[]), tcl cmn, rt cmn)) avail t :: PriorSubsts BF [CoreExpr]))
          allTrie = mkMT (tcl cmn) (\ty -> freezePS ty (let (avail,t) = splitArgs ty in mguFuns (maxBound, memoDeb)          avail t :: PriorSubsts BF [CoreExpr]))
#else
          memoTrie = mkMT (tcl cmn) (\ty -> freezePS ty (let (avail,t) = splitArgs ty in mguFuns (maxBound, memoDeb) avail t :: PriorSubsts BF [CoreExpr]))
#endif
-- We need to specialize the type (to BF) in order to avoid ambiguity.
      in memoDeb
    where qtl = splitPrims txs


-- moved from DebMT.lhs to avoid cyclic modules.
freezePS :: Search m => Type -> PriorSubsts m (Bag e) -> BFM (Possibility e)
freezePS ty ps
    = let mxty = maxVarID ty -- `max` maximum (map maxVarID avail)
      in toMemo $ mergesortDepthWithBy (\(xs,k,i) (ys,_,_) -> (xs `mappend` ys, k, i)) (\(_,k,_) (_,l,_) -> k `compare` l) $ fps mxty ps
fps :: Search m => Int -> PriorSubsts m es -> m (es,[(Int, Type)],Int)
fps mxty (PS f) = do (exprs, sub, m) <- f emptySubst (mxty+1)
                     return (exprs, filterSubst sub mxty, m)
    where filterSubst :: Subst -> Int -> [(Int, Type)]
	  filterSubst sub  mx = [ t | t@(i,_) <- sub, inRange (0,mx) i ] -- note that the assoc list is NOT sorted.

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

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



mguPrograms, mguProgs, mguFuns :: Search m => Generator m CoreExpr

mguPrograms memodeb = applyDo (mguProgs memodeb)

mguProgs memodeb = wind (>>= (return . fmap Lambda)) (\avail reqret -> reorganize (\newavail -> lookupFunsPoly mguFuns memodeb newavail reqret) avail)
{- どっちがわかりやすいかは不明
mguProgs memodeb avail (t0:->t1) = do result <- mguProgs memodeb (t0 : avail) t1
				      return (fmap Lambda result)
mguProgs memodeb avail reqret = reorganize (\newavail -> lookupFunsPoly mguFuns memodeb newavail reqret) avail
-}

mguFuns memodeb = generateFuns  mguPrograms memodeb

-- MemoDebの型が違うと使えない．
generateFuns :: (Search m) =>
                Generator m CoreExpr                               -- ^ recursive call
                -> Generator m CoreExpr
generateFuns rec memodeb@(_, md@(_, (primgen,primmono),cmn)) avail reqret
    = let behalf    = rec memodeb avail
          lltbehalf = lookupListrie (opt cmn) rec memodeb avail -- heuristic filtration
          lenavails = length avail
          fe :: Type -> Type -> [CoreExpr] -> [CoreExpr] -- ^ heuristic filtration
          fe        = filtExprs (guess $ opt cmn)
          rg        =    if tv0 $ opt cmn then retGenTV0 else
                      if tv1 $ opt cmn then retGenTV1 else retGen
      in fromAssumptions (PG md) lenavails behalf mguPS reqret avail `mplus` msum (map (rg (PG md) lenavails fe lltbehalf behalf reqret) primgen ++ map (retPrimMono (PG md) lenavails lltbehalf behalf mguPS reqret) primmono )

#ifdef DESTRUCTIVE
lookupListrie rec (memodepth, (trie, (primgen,_),tcl,rtrie)) avail t = rec (memodepth,((snd trie, snd trie), (primgen,[]), tcl, rtrie)) avail t
filtExprs = filterExprs
#else
lookupListrie opt rec memodeb avail t
                      | constrL opt = mguAssumptions t avail
                      | guess opt = do args <- rec memodeb avail t
                                       let args' = filter (not.isClosed.toCE) args
                                       when (null args') mzero
                                       return args'
                      | otherwise = rec memodeb avail t
filtExprs g a b | g         = filterExprs a b
                | otherwise = id
#endif

\end{code}
