{-# Language TypeSynonymInstances,FlexibleInstances,MultiParamTypeClasses,FunctionalDependencies,RankNTypes,FlexibleContexts,KindSignatures,ScopedTypeVariables #-}

{-

Transformation to extract aggregations from complex expressions (and recollect dependency info.?) (5th step).

Assumption:
 - no function definition (other than vertex-computation and vertex-initialization)
 - unique names
 - let-bindings in the order consistent with their dependency
 - no nested aggregation

-}

module AggregatorExtraction where

import Spec
import ASTData

import Control.Monad.State
import Data.Maybe
import Data.List
import Numeric (showHex)
import Debug.Trace

-- uniq id, aggregations, in-aggr?, move aggregator flags
type DEnvAE = (DUnique, [DSmplDef DASTData], Bool, [Bool])

class AggExtractable a  where
    aggExtract :: a -> State DEnvAE a  -- returns a transformed expression
    aggExtract x = return x            -- default

getNewName :: String -> State (DEnvAE) DVarName
getNewName s = 
    do (i, as, af, mfs) <- get 
       let (n, i') = genNewName i s
       put (i', as, af, mfs)
       return n

-- entry point
aggExtraction :: forall (t :: * -> *) a . AggExtractable (t DASTData) => (t DASTData) -> DUnique -> (t DASTData, DUnique)
aggExtraction p uid = let (p', (uid', _, _, _)) = runState (aggExtract p) (uid, [], False, [True])
                      in (p', uid')

aggExtractionE :: (DExpr DASTData) -> DUnique -> (DExpr DASTData, [DSmplDef DASTData], DUnique)
aggExtractionE e uid = let (e', (uid', defs, _, _)) = runState (aggExtract e) (uid, [], False, [True])
                      in (e', defs, uid')


-- stores the given definition of a varible binding an aggregation to the env.
addAggDef :: DSmplDef DASTData -> State (DEnvAE) ()
addAggDef agg = 
    do (d, aggs, b, m) <- get
       put (d, agg:aggs, b, m)

popAggDefs :: State DEnvAE ([DSmplDef DASTData])
popAggDefs =
    do (d, aggs, f, m) <- get
       put (d, [], f, m)
       return aggs

unsetInAgg :: State DEnvAE ()
unsetInAgg =
    do (d, aggs, _, m) <- get
       put (d, aggs, False, m)

setInAgg :: State DEnvAE ()
setInAgg = 
    do (d, aggs, f, m) <- get
       if f then error "nested aggregation is not allowed." else put (d, aggs, True, m)

needsAggMove :: State DEnvAE Bool
needsAggMove =
    do x <- get
       let (d, aggs, f, m:ms) = x
       return m

setNeedsAggMove :: State DEnvAE ()
setNeedsAggMove =
    do (d, aggs, f, ms)<- get
       put(d, aggs, f, True:ms)

unsetNeedsAggMove :: State DEnvAE ()
unsetNeedsAggMove =
    do (d, aggs, f, ms) <- get
       put (d, aggs, f, False:ms)

restoreNeedsAggMove :: State DEnvAE ()
restoreNeedsAggMove =
    do x <- get
       let (d, aggs, f, _:ms) = x
       put(d, aggs, f, ms)

instance AggExtractable (DProgramSpec DASTData) where
    aggExtract (DProgramSpec rs p a) = 
        do p' <- aggExtract p
           return (DProgramSpec rs p' a)

instance AggExtractable (DConst DASTData) 

liftToGD def = DGDefSmpl def (getData def)

getDep :: forall (t :: * -> *) . (DAdditionalData (t DASTData) DASTData) => (t DASTData) -> [String]
getDep = (depOf . getData)


setDep :: forall (t :: * -> *) . (DAdditionalData (t DASTData) DASTData) => [String] -> (t DASTData) -> (t DASTData)
setDep deps x = setData ((getData x) { depOf = deps}) x

removeDep names deps = foldr delete deps names

instance AggExtractable (DProg DASTData) where
    aggExtract (DProg f defs e a) = 
        do defaggs' <- mapM (\def -> do def' <- aggExtract def; aggs <- popAggDefs; return (def', map liftToGD aggs)) defs
           let defs' = concatMap (\(def', aggs) -> aggs++[def']) defaggs'
           e' <- aggExtract e
           aggs <- popAggDefs
           let defs'' = (defs' ++ map liftToGD aggs) 
               deps = removeDep (concatMap getNames defs''++["g"]) $ normalizeDepsX $ map getDep defs'' ++ [depOf (getData e')]
               f' = setDep deps f
               a' = a { depOf = deps}
           return (DProg f' defs'' e' a')

instance AggExtractable (DGroundDef DASTData) where
    aggExtract (DGDefVI d a) = do d' <- aggExtract d
                                  return (DGDefVI d' (a {depOf = getDep d'}) )
    aggExtract (DGDefVC d a) = do d' <- aggExtract d
                                  return (DGDefVC d' (a {depOf = getDep d'}))
    aggExtract (DGDefGV d a) = do d' <- aggExtract d
                                  return (DGDefGV d' (a {depOf = getDep d'}))
    aggExtract (DGDefGF d a) = do d' <- aggExtract d
                                  return (DGDefGF d' (a {depOf = getDep d'}))
    aggExtract (DGDefSmpl d a) = do d' <- aggExtract d
                                    return (DGDefSmpl d' (a {depOf = getDep d'}))

aggExtractLet
    :: forall (t :: * -> *) (v :: * -> *). 
       (DAdditionalData (v DASTData) DASTData,
        DAdditionalData (t DASTData) DASTData,
        AggExtractable (v DASTData)
       ) 
       => [(t DASTData)] -> [(DSmplDef DASTData)] -> (v DASTData) -> DASTData -> [String] -> Bool -> State DEnvAE ([(t DASTData)], [(DSmplDef DASTData)], (v DASTData), DASTData)
aggExtractLet vs defs e a names extract = 
        do defaggs' <- mapM (\def -> do def' <- aggExtract def; aggs <- if extract then popAggDefs else return []; return (def', aggs)) defs
           let defs' = concatMap (\(def', aggs) -> aggs++[def']) defaggs'
           e' <- aggExtract e
           aggs <- if extract then popAggDefs else return []
           let defs'' = (defs' ++ aggs) 
               deps = removeDep (concatMap getNames defs''++names) $ normalizeDepsX $ map getDep defs'' ++ [depOf (getData e')]
               vs' = map (setDep deps) vs
               a' = a { depOf = deps}
           return (vs', defs'', e', a')

instance AggExtractable (DDefVertComp DASTData) where
    aggExtract (DDefVertComp f defs e a) = 
        do x <- aggExtractLet [f] defs e a ["v", "prev", "curr"] True
           let ([f'], defs', e', a') = x
           return (DDefVertComp f' defs' e' a')

instance AggExtractable (DDefVertInit DASTData) where
    aggExtract (DDefVertInit f defs e a) = 
        do x <- aggExtractLet [f] defs e a ["v"] True
           let ([f'], defs', e', a') = x
           return (DDefVertInit f' defs' e' a')

instance AggExtractable (DDefGraphVar DASTData) where
    aggExtract (DDefGraphVar v e a) = 
        do e' <- aggExtract e
           let v' = setDep (getDep e') v
           return (DDefGraphVar v' e' (a {depOf = getDep v'}))

instance AggExtractable (DDefGraphFun DASTData) where
    aggExtract (DDefGraphFun f v defs e a) = 
        do defs' <- mapM aggExtract defs
           e' <- aggExtract e
           let f' = setDep (getDep e') f
           return (DDefGraphFun f' v defs' e' (a {depOf = getDep f'}))

instance AggExtractable (DSmplDef DASTData) where
    aggExtract (DDefFun f vs [] e a) = 
        do x <- aggExtractLet [f] [] e a (concatMap getNames vs) False
           let ([f'], _, e', a') = x
           return (DDefFun f' vs [] e' a')
    aggExtract (DDefVar v [] e a) = 
        do x <- aggExtractLet [v] [] e a [] False
           let ([v'], _, e', a') = x
           return (DDefVar v' [] e' a')
    aggExtract (DDefTuple vs [] e a) = 
        do x <- aggExtractLet vs [] e a (concatMap getNames vs) False
           let (vs', _, e', a') = x
           return (DDefTuple vs' [] e' a')

instance AggExtractable (DTermination DASTData) where
    aggExtract (DTermF a) = return (DTermF a)
    aggExtract (DTermI e a) = 
        do unsetNeedsAggMove
           e' <- aggExtract e
           restoreNeedsAggMove
           return (DTermI e' (a { depOf = getDep e'}))
    aggExtract (DTermU e a) = 
        do unsetNeedsAggMove
           e' <- aggExtract e
           restoreNeedsAggMove
           return (DTermU e' (a { depOf = removeDep ["g"] (getDep e')}))

getDeps4 t1 t2 t3 t4 = normalizeDepsX $ [getDep t1, getDep t2, getDep t3, getDep t4]
getDeps3 t1 t2 t3 = normalizeDepsX $ [getDep t1, getDep t2, getDep t3]
getDeps t1 t2 = normalizeDepsX $ [getDep t1, getDep t2]
getDepsX ts = normalizeDepsX $ map getDep ts

normalizeDeps = nub  -- remove duplicated dependency entries
normalizeDepsX = foldr (\a r -> normalizeDeps (a++r)) []

instance AggExtractable (DGraphExpr DASTData) where
    aggExtract (DPregel f0 ft x g a) = 
        do f0' <- aggExtract f0
           ft' <- aggExtract ft
           x' <- aggExtract x
           g' <- aggExtract g
           let ds = getDeps4 f0' ft' x' g'
           return (DPregel f0' ft' x' g' (a {depOf = ds}))
    aggExtract (DGMap f g a) = 
        do f' <- aggExtract f
           g' <- aggExtract g
           let ds = getDeps f' g'
           return (DGMap f' g' (a {depOf = ds}))
    aggExtract (DGZip g1 g2 a) = 
        do g1' <- aggExtract g1
           g2' <- aggExtract g2
           let ds = getDeps g1' g2'
           return (DGZip g1' g2' (a {depOf = ds}))
    aggExtract (DGIter f0 ft x g a) = 
        do f0' <- aggExtract f0
           ft' <- aggExtract ft
           x' <- aggExtract x
           g' <- aggExtract g
           let ds = getDeps4 f0' ft' x' g'
           return (DGIter f0' ft' x' g' (a {depOf = ds}))
    aggExtract (DGVar v a) = 
        do v' <- aggExtract v 
           return (DGVar v' (a {depOf = getDep v'}))

instance AggExtractable (DExpr DASTData) where
    aggExtract (DIf p t e a) = 
        do p' <- aggExtract p
           t' <- aggExtract t
           e' <- aggExtract e
           let ds = getDeps3 p' t' e'
           return (DIf p' t' e' (a {depOf = ds}))
    aggExtract (DTuple es a) =
        do es' <- mapM aggExtract es 
           let ds = getDepsX es'
           return (DTuple es' (a {depOf = ds}))
    aggExtract (DFunAp f es a) =
        do f' <- aggExtract f
           es' <- mapM aggExtract es
           let ds = normalizeDepsX (getDep f' : map getDep es')
           return (DFunAp f' es' (a {depOf = ds}))
    aggExtract (DConsAp c es a) =
        do c' <- aggExtract c
           es' <- mapM aggExtract es
           let ds = normalizeDepsX (getDep c' : map getDep es')
           return (DConsAp c' es' (a {depOf = ds}))
    aggExtract (DFieldAcc t fs a) =
        do t' <- aggExtract t
           fs' <- mapM aggExtract fs
           let ds = normalizeDepsX (getDep t' : map getDep fs')
           return (DFieldAcc t' fs' (a {depOf = ds}))
    aggExtract (DFieldAccE e fs a) =
        do e' <- aggExtract e
           fs' <- mapM aggExtract fs
           let ds = normalizeDepsX (getDep e' : map getDep fs')
           return (DFieldAccE e' fs' (a {depOf = ds}))
    aggExtract (DAggr a' e g es a) =
        do setInAgg
           a'' <- aggExtract a'
           g' <- aggExtract g
           e' <- aggExtract e
           es' <- mapM aggExtract es
           let ds = removeDep (case g of (DGenG _) -> ["u"]; (DGenTermG _) -> ["u"]; _ -> ["u","e"]) $ normalizeDepsX (getDep g' : getDep a'' : map (getDep) (e':es'))
           moveflag <- needsAggMove
           if moveflag then(do -- assigning a new name to this aggregation expression
                             newName <- getNewName "agg"
                             let d = [newName]
                                 v = DVar newName (a {depOf = d})  -- new variable biding the aggregation
                                 def = DDefVar v [] (DAggr a'' e' g' es' (a {depOf = ds})) (a {depOf = ds})
                             -- store the definition in the env
                             addAggDef def
                             unsetInAgg
                             return (DVExp v (a {depOf = d})))  -- returns the varible as the expression
           else do unsetInAgg
                   return (DAggr a'' e' g' es' (a {depOf = ds}))
    aggExtract (DVExp v a) =
        do v' <- aggExtract v
           return (DVExp v' (a {depOf = getDep v'}))
    aggExtract (DCExp c a) = 
        do c' <- aggExtract c 
           return (DCExp c' (a {depOf = getDep c'}))

instance AggExtractable (DAgg DASTData) where
    aggExtract (DAggMin a) = 
        do return (DAggMin (a { depOf = ["minimum"] }))
    aggExtract (DAggMax a) = 
        do return (DAggMax (a { depOf = ["maximum"] }))
    aggExtract (DAggSum a) = 
        do return (DAggSum (a { depOf = ["sum"] }))
    aggExtract (DAggProd a) = 
        do return (DAggProd (a { depOf = ["prod"] }))
    aggExtract (DAggAnd a) = 
        do return (DAggAnd (a { depOf = ["and"] }))
    aggExtract (DAggOr a) = 
        do return (DAggOr (a { depOf = ["or"] }))
    aggExtract (DAggChoice x a) =
        do x' <- aggExtract x
           let ds = getDep x'
           return (DAggChoice x' (a { depOf = "random":ds }))
    aggExtract (DTupledAgg as a) = 
        do as' <- mapM aggExtract as
           let ds = getDepsX as'
           return (DTupledAgg as' (a { depOf = ds }))

instance AggExtractable (DGen DASTData) where
    aggExtract (DGenI a) = return (DGenI (a {depOf = ["v"]}))
    aggExtract (DGenO a) = return (DGenO (a {depOf = ["v"]}))
    aggExtract (DGenG a) = return (DGenG (a {depOf = ["g"]}))
    aggExtract (DGenTermG a) = return (DGenTermG (a {depOf = ["g"]}))

instance AggExtractable (DEdge DASTData) where
    aggExtract (DEdge a) = return (DEdge (a {depOf = ["e"]}))

instance AggExtractable (DTableExpr DASTData) where
    aggExtract (DPrev v a) = return (DPrev v (a {depOf = [getName v, "prev"]}))
    aggExtract (DCurr v a) = return (DCurr v (a {depOf = [getName v, "curr"]}))
    aggExtract (DVal v a) = return (DVal v (a {depOf = [getName v, "val"]}))

instance AggExtractable (DField DASTData) where
--    aggExtract (DFfst a) = return (DFfst (a {depOf = ["_fst"]}))
--    aggExtract (DFsnd a) = return (DFsnd (a {depOf = ["_snd"]}))
    aggExtract (DField f a) = return (DField f (a {depOf = [f]}))

instance AggExtractable (DFun DASTData) where
    aggExtract (DBinOp f a) = return (DBinOp f (a {depOf = [f]}))
    aggExtract (DFun f a) = return (DFun f (a {depOf = [f]}))

instance AggExtractable (DVar DASTData) where
    aggExtract (DVar v a) = return (DVar v (a {depOf = [v]}))

instance AggExtractable (DConstructor DASTData) where
    aggExtract (DConstructor c a) = return (DConstructor c (a {depOf = [c]}))