-- 
-- (c) Susumu Katayama
--
module MagicHaskeller.Expression(module MagicHaskeller.Expression, module MagicHaskeller.ExprStaged, CoreExpr) where
import MagicHaskeller.CoreLang
import MagicHaskeller.MyDynamic
import MagicHaskeller.Execute
-- import Reduce
import MagicHaskeller.Types
import MagicHaskeller.ExprStaged
import MagicHaskeller.Combinators

import MagicHaskeller.T10
import Control.Monad
import Data.Array((!), array)
import MagicHaskeller.ReadDynamic
import MagicHaskeller.TyConLib(defaultTCL, TyConLib)
-- import Debug.Trace

import MagicHaskeller.Instantiate(RTrie, uncurryDyn, uncurryTy, mkUncurry, mkCurry, curryDyn)
import MagicHaskeller.DebMT

-- AnnExpr remembers each Dynamic corresponding to the CoreExpr.
data AnnExpr = AE CoreExpr Dynamic deriving Show
instance Eq AnnExpr where
    a == b = toCE a == toCE b
instance Ord AnnExpr where
    compare a b = compare (toCE a) (toCE b)

-- MemoExpr further memoizes each dynamic function.
data MemoExpr = ME CoreExpr Dynamic -- memo table
                            Dynamic -- memoized function
aeToME :: TyConLib -> RTrie -> Type -> AnnExpr -> MemoExpr
aeToME tcl (_,_,_,_,mtdd) ty@(_:->_) (AE ce dyn)
    = case lookupMT mtdd argty of (m,a) -> let me@(ME _ memo _) = ME ce (dynApp m udyn) (curryDyn cur ty $ dynApp a memo)
                                           in me  -- make sure to use the memo table in the datatype by using letrec, or the table will be recomputed.
    where argty:->_ = uncurryTy tcl ty
          unc   = mkUncurry tcl
          udyn  = uncurryDyn unc ty dyn
          cur   = mkCurry tcl
aeToME _   _              _          (AE ce dyn) = ME ce undefined dyn -- non-functional case

meToAE :: MemoExpr -> AnnExpr
meToAE (ME ce _ f) = AE ce f

class (Ord e, Show e) => Expression e where
    fromCE         :: (CoreExpr->Dynamic) -> Int -> Int -> CoreExpr -> e -- この名前は誤解を招くかも．mkHeadとかそんな名前にする?
    toCE           :: e -> CoreExpr
    mapCE          :: (CoreExpr -> CoreExpr) -> e -> e  -- これも変．
    (<$>)          :: e -> e -> e
    appEnv         :: Int -> e -> e -> e
    toAnnExpr      :: (CoreExpr->Dynamic) -> e -> AnnExpr
    toAnnExprWind  :: (CoreExpr->Dynamic) -> Type -> e -> AnnExpr
    toAnnExprWindWind :: (CoreExpr->Dynamic) -> Type -> e -> AnnExpr
    fromAnnExpr    :: AnnExpr -> e
    reorganize     :: Monad m => ([Type] -> m [e]) -> [Type] -> m [e]
    reorganizeId   ::  ([Type] -> [e]) -> [Type] -> [e] -- reorganize for Id monad
instance Expression CoreExpr where
    fromCE _ _ _            = id
    toCE                  = id
    mapCE                 = id
    (<$>)                 = (:$)
    appEnv _              = (:$)
    toAnnExpr reduce e           = AE e (reduce e)
    toAnnExprWind reduce ty e    = AE e (reduce $ windType ty e)
    toAnnExprWindWind reduce ty e = let we = windType ty e in AE we (reduce we)
    fromAnnExpr (AE ce _) = ce
    reorganize = reorganizer
    reorganizeId = reorganizerId
instance Expression AnnExpr where
    fromCE _      lenavails arity ce@(X i) = AE ce (getDyn_LambdaBoundHead i lenavails arity)                  -- Note that 'dynss' and 'dynsss' uses
    fromCE reduce lenavails arity ce       = AE ce ((getDyn lenavails arity) `dynApp` reduce ce) -- 'unsafeExecute' instead of 'reduce'.
    toCE (AE ce _)                  = ce
    mapCE f (AE ce d)               = AE (f ce) d
    AE e1 h1   <$>   AE e2 h2       = AE (e1:$e2) (dynApp h1 h2)
    appEnv lenavails (AE e1 h1) (AE e2 h2) = AE (e1:$e2) (dynApp (dynApp (dynSn lenavails) h1) h2)
    toAnnExpr     _                 = id
    toAnnExprWind _ _               = id
    toAnnExprWindWind _ ty (AE ce d) = AE (windType ty ce) d
    fromAnnExpr                     = id
    reorganize = id
    reorganizeId = id
windType :: Type -> CoreExpr -> CoreExpr
windType (a:->b) e = Lambda (windType b e)
windType _       e = e

-- Sn = \f g x1 .. xn -> f x1 .. xn (g x1 .. xn)
dynSn lenavails = dynApp (getDyn lenavails 2) dynI

getDyn, mkDyn :: Int -> Int -> Dynamic
getDyn lenavails arity
--    | arity<=maxArity = case lenavails `divMod` maxLenavails of (d,m) -> napply d (dynApp (dynApp dynB (finiteDynar!(maxLenavails,arity)))) (finiteDynar!(m,arity)) -- なんか違うみたい．
    | lenavails<=maxLenavails && arity<=maxArity = -- trace (show (lenavails,arity)++show (maxLenavails,maxArity)) $
                                                   finiteDynar ! (lenavails,arity)
    | otherwise                                  = dynss !! lenavails !! arity

dynss :: [[Dynamic]]
dynss = [ [ mkDyn i j | j <- [0..] ] | i <- [0..] ]
mkDyn 0         _ = dynI
{-
mkDyn lenavails 0 = unsafeExecute (B :$ K) `dynApp` mkDyn (lenavails-1) 0 
mkDyn lenavails arity = unsafeExecute $ mkCE lenavails arity 
-}
-- mkDyn lenavails arity = napply lenavails (dynApp (dynB `dynApp` x arity)) dynI
mkDyn lenavails arity = dynApp (dynB `dynApp` x arity) (getDyn (lenavails-1) arity)

-- #ifdef TEMPLATE_HASKELL
-- x n | n<=maxArity = finiteDynar ! (1,n) -- なんか違うみたい．
-- #else
x 0 = dynK
x 1 = dynB
x 2 = dynS'
-- x 3 = unsafeToDyn (readType "(a->b->c->r)->(x->a)->(x->b)->(x->c)->x->r")    x3      ()
-- #endif
x n = napply n (dynApp dynB) dynS `dynApp` x (n-1)

finiteDynar  = array ((0,0),(maxLenavails,maxArity)) [ ((lenavails,arity), finiteDynss!!lenavails!!arity) | arity<-[0..maxArity], lenavails<-[0..maxLenavails] ]
finiteDynarr = array ((0,0,0),(maxDebindex,maxLenavails,maxArity)) [ ((debindex,lenavails,arity), finiteDynsss!!debindex!!(lenavails-debindex-1)!!arity) | arity<-[0..maxArity], debindex<-[0..maxDebindex], lenavails<-[debindex+1..maxLenavails] ]

finiteDynss = zipWith3 (zipWith3 (unsafeToDyn defaultTCL)) [ [ hdmnty  arity lenavails | arity <- [0..maxArity] ] | lenavails <- [0..maxLenavails] ]
                                              finiteHVss
                                              [ [ hdmnTHE arity lenavails | arity <- [0..maxArity] ] | lenavails <- [0..maxLenavails] ]
finiteDynsss = zipWith3 (zipWith3 (zipWith3 (unsafeToDyn defaultTCL)))
                        [ [ [ aimnty  debindex arity lenavails | arity <- [0..maxArity] ] | lenavails <- [debindex+1..maxLenavails] ] | debindex <- [0..maxDebindex] ]
                        finiteHVsss
                        [ [ [ aimnTHE debindex arity lenavails | arity <- [0..maxArity] ] | lenavails <- [debindex+1..maxLenavails] ] | debindex <- [0..maxDebindex] ]

getDyn_LambdaBoundHead, mkDyn_LambdaBoundHead :: Int -> Int -> Int -> Dynamic
getDyn_LambdaBoundHead debindex lenavails arity
    | debindex<=maxDebindex && lenavails<=maxLenavails && arity<=maxArity = -- trace (show (debindex,lenavails,arity)++show (maxDebindex,maxLenavails,maxArity)) $
                                                                            finiteDynarr ! (debindex,lenavails,arity) -- こっちの方が効率的なんだけど，デバッグ中だけ一時的に．
                                                                            -- finiteDynsss !! debindex !! (lenavails-debindex-1) !! arity
    | otherwise                                  = dynsss !! debindex !! lenavails !! arity

dynsss :: [[[Dynamic]]]
dynsss = [ [ [ mkDyn_LambdaBoundHead i j k | k <- [0..] ] | j <- [0..] ] | i <- [0..] ]
mkDyn_LambdaBoundHead debindex lenavails arity = (getDyn lenavails (arity+1) `dynApp` dynI) `dynApp` select lenavails debindex
    where
      -- select lenavails debindex = unsafeExecute (napply lenavails Lambda $ X debindex)
      select lenavails debindex = napply (lenavails-1-debindex) (dynApp dynK) $ napply debindex (dynApp dynBK) dynI
dynBK = dynApp dynB dynK

-- dynF = dynApp dynC dynK

-- moved from ProgramGenerator.lhs

-- reorganize :: ([Type] -> PriorSubsts BF [CoreExpr]) -> [Type] -> PriorSubsts BF [CoreExpr]
-- として使われるのだが，特にexportされる訳でもないのでいちいちspecializeしない．
reorganizer :: Monad m => ([Type] -> m [CoreExpr]) -> [Type] -> m [CoreExpr]
reorganizer fun avail
    = case cvtAvails avail of
       (newavail, argss) ->
         do agentExprs <- fun newavail
            return [ result | e <- agentExprs, result <- replaceVars 0 e argss ]

reorganizerId :: ([Type] -> [CoreExpr]) -> [Type] -> [CoreExpr]
reorganizerId fun avail
    = case cvtAvails avail of
       (newavail, argss) ->
           [ result | e <- fun newavail, result <- replaceVars 0 e argss ]
replaceVars :: Int -> CoreExpr -> [[Int]] -> [CoreExpr]
replaceVars dep e@(X n)    argss = case argss !? (n - dep) of Nothing -> [e]
                                                              Just xs -> map (\ m -> X (m + dep)) xs
replaceVars dep (Lambda e) argss = map Lambda (replaceVars (dep+1) e argss)
replaceVars dep (f :$ e)   argss = liftM2 (:$) (replaceVars dep f argss) (replaceVars dep e argss)
replaceVars dep e          argss = [e]

cvtAvails = unzip . tkr10 . annotate

tkr10 :: [(Type,Int)] -> [(Type,[Int])]
tkr10 = mergesortWithBy (\ (k,is) (_,js) -> (k,is++js)) (\ (k,_) (l,_) -> k `compare` l) . map (\(k,i)->(k,[i]))


-- annotateはsplitAvailsの前処理としても使える．
annotate :: [Type] -> [(Type,Int)]
annotate ts = zipWith (,) ts [0..]
{-
annotate ts = an 0 ts
    where an n []     = []
          an n (t:ts) = (t,n) : an (n+1) ts
prop_annotate = \ts -> annotate ts == zipWith (,) ts [0..]
-}


-- Moved from T10
-- uniqSorter :: (Ord e) => [(e,Int)] -> [(e,Int)]
uniqSorter :: (Expression e) => [(e,Int)] -> [(e,Int)]
uniqSorter = annUniqSort -- swapUniqSort -- id -- uniqSort -- annUniqSort

uniqSort :: Ord a => [a] -> [a]
uniqSort = mergesortWithBy const compare
swapUniqSort :: (Ord a, Ord b) => [(a,b)] -> [(a,b)]
swapUniqSort = mergesortWithBy const (\(a,b) (c,d) -> compare (b,a) (d,c))
annUniqSort :: Expression e => [(e,Int)] -> [(e,Int)]
annUniqSort = map snd  .  mergesortWithBy const (\a b -> compare (fst a) (fst b))  .  map (\t@(ce,_i) -> (fromEnum $ toCE ce, t))
aUS :: Expression e => [e] -> [e]
aUS = map snd  .  mergesortWithBy const (\a b -> compare (fst a) (fst b))  .  map (\e -> (fromEnum $ toCE e, e))