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

{-

Collecting dependency info. as the last phase of transformation on ASTs. (10th step)

-}

module DependencySimple where

import Spec
import ASTData

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


-- entry point

computeDependency p = let (p', _) = runState (depAnnotate p) ()
                      in p'

class DependencyAnnotatable a  where
    depAnnotate :: a -> State () a     -- reannotate [name to which the given term depends]
    depAnnotate x = return x           -- default


instance DependencyAnnotatable (DProgramSpec DASTData) where
    depAnnotate (DProgramSpec rs p a) = 
        do p' <- depAnnotate p
           return (DProgramSpec rs p' a)

instance DependencyAnnotatable (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 DependencyAnnotatable (DProg DASTData) where
    depAnnotate (DProg f defs e a) = 
        do defs' <- mapM depAnnotate defs
           e' <- depAnnotate e
           let 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 DependencyAnnotatable (DGroundDef DASTData) where
    depAnnotate (DGDefVI d a) = do d' <- depAnnotate d
                                   return (DGDefVI d' (a {depOf = getDep d'}) )
    depAnnotate (DGDefVC d a) = do d' <- depAnnotate d
                                   return (DGDefVC d' (a {depOf = getDep d'}))
    depAnnotate (DGDefGV d a) = do d' <- depAnnotate d
                                   return (DGDefGV d' (a {depOf = getDep d'}))
    depAnnotate (DGDefGF d a) = do d' <- depAnnotate d
                                   return (DGDefGF d' (a {depOf = getDep d'}))
    depAnnotate (DGDefSmpl d a) = do d' <- depAnnotate d
                                     return (DGDefSmpl d' (a {depOf = getDep d'}))

depAnnotateLet
    :: forall (t :: * -> *) (v :: * -> *). 
       (DAdditionalData (v DASTData) DASTData,
        DAdditionalData (t DASTData) DASTData,
        DependencyAnnotatable (v DASTData)
       ) 
       => [(t DASTData)] -> [(DSmplDef DASTData)] -> (v DASTData) -> DASTData -> [String] -> State () ([(t DASTData)], [(DSmplDef DASTData)], (v DASTData), DASTData)
depAnnotateLet vs defs e a names = 
        do defs' <- mapM depAnnotate defs
           e' <- depAnnotate e
           let 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 DependencyAnnotatable (DDefVertComp DASTData) where
    depAnnotate (DDefVertComp f defs e a) = 
        do x <- depAnnotateLet [f] defs e a ["v", "prev", "curr"]
           let ([f'], defs', e', a') = x
           return (DDefVertComp f' defs' e' a')

instance DependencyAnnotatable (DDefVertInit DASTData) where
    depAnnotate (DDefVertInit f defs e a) = 
        do x <- depAnnotateLet [f] defs e a ["v"]
           let ([f'], defs', e', a') = x
           return (DDefVertInit f' defs' e' a')

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

instance DependencyAnnotatable (DDefGraphFun DASTData) where
    depAnnotate (DDefGraphFun f v defs e a) = 
        do defs' <- mapM depAnnotate defs
           e' <- depAnnotate e
           let ds = removeDep (getNames v ++ concatMap getNames defs)$ normalizeDepsX $ [getDep e'] ++ map getDep defs'
               f' = setDep ds f
           return (DDefGraphFun f' v defs' e' (a {depOf = getDep f'}))

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

instance DependencyAnnotatable (DTermination DASTData) where
    depAnnotate (DTermF a) = return (DTermF a)
    depAnnotate (DTermI e a) = 
        do e' <- depAnnotate e
           return (DTermI e' (a { depOf = getDep e'}))
    depAnnotate (DTermU e a) = 
        do e' <- depAnnotate e
           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 DependencyAnnotatable (DGraphExpr DASTData) where
    depAnnotate (DPregel f0 ft x g a) = 
        do f0' <- depAnnotate f0
           ft' <- depAnnotate ft
           x' <- depAnnotate x
           g' <- depAnnotate g
           let ds = normalizeDepsX $ [getDeps4 f0' ft' x' g'] ++ [getNames f0', getNames ft'] --, getNames g' ]
           return (DPregel f0' ft' x' g' (a {depOf = ds}))
    depAnnotate (DGMap f g a) = 
        do f' <- depAnnotate f
           g' <- depAnnotate g
           let ds = normalizeDepsX $ [getDeps f' g'] -- ++ [getNames f', getNames g' ]
           return (DGMap f' g' (a {depOf = ds}))
    depAnnotate (DGZip g1 g2 a) = 
        do g1' <- depAnnotate g1
           g2' <- depAnnotate g2
           let ds = normalizeDepsX $ [getDeps g1' g2'] -- ++ [getNames g1', getNames g2' ]
           return (DGZip g1' g2' (a {depOf = ds}))
    depAnnotate (DGIter f0 ft x g a) = 
        do f0' <- depAnnotate f0
           ft' <- depAnnotate ft
           x' <- depAnnotate x
           g' <- depAnnotate g
           let ds = normalizeDepsX $ [getDeps4 f0' ft' x' g']++ [getNames f0', getNames ft'] -- , getNames g' ]
           return (DGIter f0' ft' x' g' (a {depOf = ds}))
    depAnnotate (DGVar v a) = 
        do v' <- depAnnotate v 
           return (DGVar v' (a {depOf = getDep v'}))

instance DependencyAnnotatable (DExpr DASTData) where
    depAnnotate (DIf p t e a) = 
        do p' <- depAnnotate p
           t' <- depAnnotate t
           e' <- depAnnotate e
           let ds = getDeps3 p' t' e'
           return (DIf p' t' e' (a {depOf = ds}))
    depAnnotate (DTuple es a) =
        do es' <- mapM depAnnotate es 
           let ds = getDepsX es'
           return (DTuple es' (a {depOf = ds}))
    depAnnotate (DFunAp f es a) =
        do f' <- depAnnotate f
           es' <- mapM depAnnotate es
           let ds = normalizeDepsX (getDep f' : map getDep es')
           return (DFunAp f' es' (a {depOf = ds}))
    depAnnotate (DConsAp c es a) =
        do c' <- depAnnotate c
           es' <- mapM depAnnotate es
           let ds = normalizeDepsX (getDep c' : map getDep es')
           return (DConsAp c' es' (a {depOf = ds}))
    depAnnotate (DFieldAcc t fs a) =
        do t' <- depAnnotate t
           fs' <- mapM depAnnotate fs
           let ds = normalizeDepsX (getDep t' : map getDep fs')
           return (DFieldAcc t' fs' (a {depOf = ds}))
    depAnnotate (DFieldAccE e fs a) =
        do e' <- depAnnotate e
           fs' <- mapM depAnnotate fs
           let ds = normalizeDepsX (getDep e' : map getDep fs')
           return (DFieldAccE e' fs' (a {depOf = ds}))
    depAnnotate (DAggr a' e g es a) =
        do a'' <- depAnnotate a'
           g' <- depAnnotate g
           e' <- depAnnotate e
           es' <- mapM depAnnotate es
           let ds = removeDep (case g of (DGenG _) -> ["u"]; (DGenTermG _) -> ["u"]; _ -> ["u","e"]) $ normalizeDepsX (getDep g' : getDep a'' : map (getDep) (e':es'))
           return (DAggr a'' e' g' es' (a {depOf = ds}))
    depAnnotate (DVExp v a) =
        do v' <- depAnnotate v
           return (DVExp v' (a {depOf = getDep v'}))
    depAnnotate (DCExp c a) = 
        do c' <- depAnnotate c 
           return (DCExp c' (a {depOf = getDep c'}))

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

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

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

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

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

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

instance DependencyAnnotatable (DVar DASTData) where
    depAnnotate (DVar v a) = return (DVar v (a {depOf = [v]}))

instance DependencyAnnotatable (DConstructor DASTData) where
    depAnnotate (DConstructor c a) = return (DConstructor c (a {depOf = [c]}))