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


{-

Depencency check, reordering of let-bindings, and alpha renaming with unique names. (1st step)

-}

module Analysis where

import Spec

import Debug.Trace

import Control.Monad.State
import Data.Maybe
import Data.List
import ASTData

type DNewName = DVarName
type DDepth = Int
type DDepVar = (DVarName, (DDepth, DNewName))
type DDeps = [DDepVar]
type DEnvDC a = ([(DDeps, DDepth)], DUnique, a, Bool)

-- extracted aggregations, in-aggregator flag, 
type NormalizeData1 a = ([DSmplDef (DDeps, a)], Bool, [Bool])
class DepCheckable a b c | a -> b c where
    depCheck :: a -> State (DEnvDC c) b  -- assigns [(name to which the given term depends, depth)] to the given term


-- to flattern definitions (within a vertex-compute/vertex-init function)
collectDef :: DSmplDef a -> [DSmplDef a]
collectDef (DDefFun f vs defs e a) = concatMap collectDef defs++[(DDefFun f vs [] e a)]
collectDef (DDefVar v defs e a) = concatMap collectDef defs++[(DDefVar v [] e a)]
collectDef (DDefTuple vs defs e a) = concatMap collectDef defs++[(DDefTuple vs [] e a)]

-- getter/setter on the environment

getEnvDepth' :: DEnvDC a -> DDepth
getEnvDepth' (bds, _, _, _) = snd (head bds)

getEnvDepth :: State (DEnvDC a) DDepth
getEnvDepth = 
    do env <- get 
       return (getEnvDepth' env)

getEnvBinds' :: DEnvDC a -> DDeps
getEnvBinds' (bds, _, _, _) = fst (head bds)

getEnvBinds :: State (DEnvDC a) DDeps
getEnvBinds = 
    do env <- get 
       return (getEnvBinds' env)

setEnvBindsDepth :: (DDeps, DDepth) -> State (DEnvDC a) ()
setEnvBindsDepth bd =
    do (bds, i, a, b) <- get
       put (bd:bds, i, a, b)

popBindsDepth :: State (DEnvDC a) ()
popBindsDepth =
    do (bds, i, a, b) <- get
       put (tail bds, i, a, b)

getNewName :: String -> State (DEnvDC a) DNewName
getNewName n = 
    do (bds, i, a, b) <- get 
       let (nn, i') = genNewName i n
       put (bds, i', a, b)
       return nn

doingAlpha :: State (DEnvDC a) Bool
doingAlpha =
    do (bds, i, a, b) <- get
       return b

replaceApostrophe n = map (\c -> if c == '\'' then 'X' else c) n
              
-- orgName = true => new name = original name  (for predefined identifiers)
getNewName' :: Bool -> (DVarName, DDepth) -> State (DEnvDC a) DDepVar
getNewName' orgName (n, d) = 
    do b <- doingAlpha
       if b then do let n' = replaceApostrophe n
                    nn <- getNewName n'
                    return (n, (d, if orgName then n else nn))
            else return (n, (d, n))

getDVDepth :: DDepVar -> DDepth
getDVDepth (_, (d, _)) = d

getDVVar :: DDepVar -> DVarName
getDVVar (v, (_, _)) = v

getDVNewName :: DDepVar -> DNewName
getDVNewName (_, (_, n)) = n

-- entry point for analysis chain (accepts the uniq ID)
dependencyC :: forall (t :: * -> *) a . 
               (DepCheckable (t a) (t (DDeps, a)) (NormalizeData1 a),
                DAdditionalData2 (t (DDeps, a)) (DDeps, a) ([String], a) (t ([String], a))) 
               => DUnique -> t a -> (t ([DVarName], a), DUnique)
dependencyC uid p = 
    let (ast, uid') = dependencyC' uid p
    in (mapData (\(ds::DDeps, a::a) -> (map (snd.snd) ds, a)) ast, uid')

dependencyC' :: forall (t :: * -> *) a . 
               (DepCheckable (t a) (t (DDeps, a)) (NormalizeData1 a),
                DAdditionalData2 (t (DDeps, a)) (DDeps, a) ([String], a) (t ([String], a))) 
               => DUnique -> t a -> (t (DDeps, a), DUnique)
dependencyC' uid p = 
    let (ast, (_,uid',_,_)) = runState (depCheck0 p) ([([], -2)], uid, ([], False, [True]), True)
    in (ast, uid')

-- no alpha-renaming
topoSortOnly :: forall (t :: * -> *) a . 
               (DepCheckable (t a) (t (DDeps, a)) (NormalizeData1 a),
                DAdditionalData2 (t (DDeps, a)) (DDeps, a) a (t a)) 
               => t a -> t a
topoSortOnly p = 
  let (ast, _) = runState (depCheck0 p) ([([], -2)], -100, ([], False, [True]), False)
  in mapData (snd::(DDeps,a)->a) ast

-- entry point for single use
dependency :: forall (t :: * -> *) a . 
               (DepCheckable (t a) (t (DDeps, a)) (NormalizeData1 a),
                DAdditionalData2 (t (DDeps, a)) (DDeps, a) ([String], a) (t ([String], a))
               ) => t a -> t (DDeps, a)
dependency p = fst $ dependencyC' 0 p

-- cleaned version
dependency' :: forall (t :: * -> *) a . 
               (DepCheckable (t a) (t (DDeps, a)) (NormalizeData1 a),
                DAdditionalData2 (t (DDeps, a)) (DDeps, a) ([String], a) (t ([String], a))
               ) => t a -> t ([String], a)
dependency' p = fst $ dependencyC 0 p


-- cleaned data

depCheck0 p = 
    do addNames' True initDCnames  -- names at level -1
       depCheck p
  
initDCnames = map fst initEnv

-- misc functions for dependencies
getDep :: forall (t :: * -> *) b . (DAdditionalData (t (DDeps, b)) (DDeps, b)) => (t (DDeps, b)) -> DDeps
getDep = fst . getData

getDepL :: forall (t :: * -> *) b . (DAdditionalData (t (DDeps, b)) (DDeps, b)) => Int -> (t (DDeps, b)) -> DDeps
getDepL depth = filter ((< depth) . getDVDepth ) . getDep    -- names in lower (outer) depth levels
getDepE depth = filter ((== depth) . getDVDepth ) . getDep   -- names at the depth level

getDeps a b = normalizeDepsX [getDep a, getDep b]
getDeps3 a b c = normalizeDepsX [getDep a, getDep b, getDep c]
getDeps4 a b c d = normalizeDepsX [getDep a, getDep b, getDep c, getDep d]
getDepsX es = normalizeDepsX $ map getDep es

normalizeDeps :: DDeps -> DDeps
normalizeDeps = nub  -- remove duplicated dependency entries

normalizeDepsX ::[DDeps] -> DDeps
normalizeDepsX = foldr (\a r -> normalizeDeps (a++r)) []

-- reordering the let bindings (or any named things) according to their dependencies among them
-- NOTE: run after the alpha-conversion
depTopoSort ::  forall (t :: * -> *) b . (DNamed (t (DDeps, b)), DAdditionalData (t (DDeps, b)) (DDeps, b)) =>
     DDepth -> [(t (DDeps, b)) ] -> [(t (DDeps, b)) ]
depTopoSort depth ts = depTopoSort' True depth ts

-- withNew == True => using new alpha-converted names 
depTopoSort' ::  forall (t :: * -> *) b . (DNamed (t (DDeps, b)), DAdditionalData (t (DDeps, b)) (DDeps, b)) =>
     Bool -> DDepth -> [(t (DDeps, b)) ] -> [(t (DDeps, b)) ]
depTopoSort' withNew depth ts = ds'
 where
  ts' = zip [0..] ts
  vs = map findDepIds ts'
  findDepIds (i, t) = (i, map findId $ getDepE depth t, t)
  findId :: DDepVar -> Int  -- this dv is defined in the returned position
  findId dv = let tt = findIndices (elem (if withNew then getDVNewName dv else getDVVar dv) . getNames) ts in if tt == [] then error ("Something wrong: " ++ show (map getNames ts) ++ " -- " ++ (show dv)) else head tt
  (ds, ok) = topologicalSort vs 
  ds' = if ok then map snd ds else error ("circular dependency among " ++ show (map getNames ts))

-- topological sort algorithm: https://ja.wikipedia.org/wiki/%E3%83%88%E3%83%9D%E3%83%AD%E3%82%B8%E3%82%AB%E3%83%AB%E3%82%BD%E3%83%BC%E3%83%88
-- (id, outNeighbors, data)
-- id starts from 0, contiguous =>  assumption: vs!!i = (i, _, _)
topologicalSort :: [(Int, [Int], a)] -> ([(Int, a)], Bool)
topologicalSort vs =
  let vs' = map (findIncoming vs) vs
      findIncoming us (v,ns,d) = (v, ns, d, length $ filter (\(_, os, _) -> elem v os) us)
      ins (_,_,_,i) = i
      decIns (x,ns,d,i) = (x, ns, d, i-1)
      noIncomings = filter ((==0) . ins)
      s = noIncomings vs'
      (l, g) = rec s [] vs'
      ok = all ((==0). ins)g
      rec [] l g = (l, g)
      rec (((x, ns, d, _)):s) l g 
          = let l' = (x, d):l
                (g', s') = foldr step (g, s) ns
                step j (g, s) = let (g1, y:g2) = splitAt j g
                                    g' = g1++((decIns y):g2)
                                in (g', if ins y == 1 then y:s else s)
            in rec s' l' g'
  in (l, ok)


-- adds names to the envitonment, returns the old env with the current depth
addNames :: [String] -> State (DEnvDC a) (Int, [DNewName])
addNames ns = addNames' False ns

addNames' :: Bool -> [String] -> State (DEnvDC a) (Int, [DNewName])
addNames' orgName ns = 
    do depth <- getEnvDepth
       binds <- getEnvBinds
       let ns' = if ns == nub ns then ns else error ("multiple definitions in " ++ show ns)
           depth' = depth + 1
           nds = zip ns' (repeat depth')
       binds' <- mapM (getNewName' orgName) nds
       setEnvBindsDepth (binds'++ binds, depth')
       return (depth', map getDVNewName binds')

-- removed added names by addNames
popNames :: State (DEnvDC a) ()
popNames = popBindsDepth

{- dependency check of expression e with simple-let-bindings defs -}
depCheckLet
    :: forall (u :: * -> *) b a . 
       (DepCheckable (u b) (u (DDeps, b)) (NormalizeData1 b),
        DAdditionalData (u (DDeps, b)) (DDeps, b)) 
       => [(DSmplDef b)] -> (u b) -> [DVarName] -> State (DEnvDC (NormalizeData1 b)) ([(DSmplDef (DDeps, b))], (u (DDeps, b)), DDeps, [DNewName])
depCheckLet defs e vs = depCheckLetG' False defs e vs

-- general version 
depCheckLetG' 
    :: forall (t :: * -> *) (u :: * -> *) b . 
       (DNamed (t b), 
        DNamed (t (DDeps, b)), 
        DepCheckable (t b) (t (DDeps, b)) (NormalizeData1 b),
        DepCheckable (u b) (u (DDeps, b)) (NormalizeData1 b),
        DAdditionalData (t (DDeps, b)) (DDeps, b),
        DAdditionalData (u (DDeps, b)) (DDeps, b)) 
       => Bool -> [(t b)] -> (u b) -> [DVarName] -> State (DEnvDC (NormalizeData1 b)) ([(t (DDeps, b))], (u (DDeps, b)), DDeps, [DNewName])
depCheckLetG' orgName defs e vs =
    do (_, nvs) <- addNames' orgName vs  -- add the argument variables first,
       let ons = concatMap getNames defs
       (depth, nns) <- addNames ons -- then add the let-bindings
       defs' <- mapM depCheck defs
       e' <- depCheck e
       let defs'' = replaceNames nns defs'  
           dss = map (getDepL depth) defs'' ++ [getDepL depth e'] 
           ds = normalizeDepsX dss
           defs''' = depTopoSort depth $ defs''  -- sort the bindings according to their dependency
       popNames
       popNames
       return (defs''', e', ds, nvs)

{- replaces old names with newnames at the depth -}
replaceNames nns dfs = dfs'
  where dfs' = evalState (mapM setNames dfs) nns

instance (Show a) => DepCheckable (DProgramSpec a) (DProgramSpec (DDeps, a)) (NormalizeData1 a) where
    depCheck (DProgramSpec rs p a) = 
        do (rs', p', ds, _) <- depCheckLetG' False rs p []
           return (DProgramSpec rs' p' (ds, a))

{- newly introduced names have no dependency -}
nullDep :: forall (k :: * -> *) a.
                 DAdditionalData2 (k a) a (DDeps, a) (k (DDeps, a)) =>
                 (k a) -> (k (DDeps, a))
nullDep f = mapData null f
  where null :: a -> (DDeps, a)
        null a = ([], a)

instance Show a => DepCheckable (DRecordSpec a) (DRecordSpec (DDeps, a)) (NormalizeData1 a) where
    depCheck (DRecordSpec c fts a) =
        do dst <- mapM depCheck (map snd fts)
           let dsf = map nullDep (map fst fts)
               c' = nullDep c
               fts' = zip dsf dst
               ds = normalizeDeps $ concatMap getDep dst
               ret = (DRecordSpec c' fts' (ds, a))
           return ret

instance DepCheckable (DType a) (DType (DDeps, a)) (NormalizeData1 a) where
    depCheck (DTInt a) = return (DTInt ([],a))
    depCheck (DTBool a) = return (DTBool ([],a))
    depCheck (DTString a) = return (DTString ([],a))
    depCheck (DTDouble a) = return (DTDouble ([],a))
    depCheck (DTTuple ts a) = 
        do ts' <- mapM depCheck ts
           let ds = normalizeDeps $ concatMap getDep ts'
           return (DTTuple ts' (ds, a))
    depCheck (DTRecord c ts a) = 
        do c' <- depCheck c
           ts' <- mapM depCheck ts
           let ds = normalizeDeps $ concatMap getDep ts' ++ getDep c'
           return (DTRecord c' ts' (ds,a))

instance DepCheckable (DConst a) (DConst (DDeps, a)) (NormalizeData1 a) where
    depCheck (DCInt c a) = return (DCInt c ([], a))
    depCheck (DCBool c a) = return (DCBool c ([], a))
    depCheck (DCString c a) = return (DCString c ([], a))
    depCheck (DCDouble c a) = return (DCDouble c ([], a))

instance DepCheckable (DProg a) (DProg (DDeps, a)) (NormalizeData1 a) where
    depCheck (DProg f defs e a) = 
        do (defs', e', ds, _) <- depCheckLetG' True defs e ["g"]
           let f' = nullDep f
           return (DProg f' defs' e' (ds, a))

instance DepCheckable (DGroundDef a) (DGroundDef (DDeps, a)) (NormalizeData1 a) where
    depCheck (DGDefVI d a) = do d' <- depCheck d
                                return (DGDefVI d' (getDep d', a))
    depCheck (DGDefVC d a) = do d' <- depCheck d
                                return (DGDefVC d' (getDep d', a))
    depCheck (DGDefGV d a) = do d' <- depCheck d
                                return (DGDefGV d' (getDep d', a))
    depCheck (DGDefGF d a) = do d' <- depCheck d
                                return (DGDefGF d' (getDep d', a))
    depCheck (DGDefSmpl d a) = do d' <- depCheck d
                                  return (DGDefSmpl d' (getDep d', a))

instance DepCheckable (DDefVertComp a) (DDefVertComp (DDeps, a)) (NormalizeData1 a) where
    depCheck (DDefVertComp f defs e a) = 
        do (defs', e', ds, _) <- depCheckLetG' True defs e ["v", "prev", "curr"]
           let f' = nullDep f
               defs'' = concatMap collectDef defs'
           return (DDefVertComp f' defs'' e' (ds, a))

instance DepCheckable (DDefVertInit a) (DDefVertInit (DDeps, a)) (NormalizeData1 a) where
    depCheck (DDefVertInit f defs e a) = 
        do (defs', e', ds, _) <- depCheckLetG' True defs e ["v"]
           let f' = nullDep f
               defs'' = concatMap collectDef defs'
           return (DDefVertInit f' defs'' e' (ds, a))

instance DepCheckable (DDefGraphVar a) (DDefGraphVar (DDeps, a)) (NormalizeData1 a) where
    depCheck (DDefGraphVar v e a) = 
        do e' <- depCheck e
           let v' = nullDep v
           return (DDefGraphVar v' e' (getDep e', a))

instance DepCheckable (DDefGraphFun a) (DDefGraphFun (DDeps, a)) (NormalizeData1 a) where
    depCheck (DDefGraphFun f v defs e a) = 
        do (defs', e', ds, vns) <- depCheckLetG' True defs e (getNames v)
           let f' = nullDep f
               [v'] = replaceNames vns [nullDep v]
           return (DDefGraphFun f' v' defs' e' (ds, a))

instance DepCheckable (DSmplDef a) (DSmplDef (DDeps, a)) (NormalizeData1 a) where
    depCheck (DDefFun f vs defs e a) = 
        do let vns = concatMap getNames vs
               f' = nullDep f
               vs' = map nullDep vs
           (defs', e', ds, vns') <- depCheckLet defs e vns
           let vs'' = replaceNames vns' vs'
           return (DDefFun f' vs'' defs' e' (ds, a))
    depCheck (DDefVar v defs e a) = 
        do (defs', e', ds, _) <- depCheckLet defs e []
           let v' = nullDep v
           return (DDefVar v' defs' e' (ds, a))
    depCheck (DDefTuple vs defs e a) = 
        do (defs', e', ds, _) <- depCheckLet defs e []
           let vs' = map nullDep vs
           return (DDefTuple vs' defs' e' (ds, a))

instance DepCheckable (DTermination a) (DTermination (DDeps, a)) (NormalizeData1 a) where
    depCheck (DTermF a) = return (DTermF ([], a)) 
    depCheck (DTermI e a) = 
        do e' <- depCheck e
           return (DTermI e' (getDep e', a))
    depCheck (DTermU e a) = 
        do addNames' True ["g"]
           e' <- depCheck e
           popNames
           return (DTermU e' (getDep e', a))

instance DepCheckable (DGraphExpr a) (DGraphExpr (DDeps, a)) (NormalizeData1 a) where
    depCheck (DPregel f0 ft x g a) = 
        do f0' <- depCheck f0
           ft' <- depCheck ft
           x' <- depCheck x
           g' <- depCheck g
           let ds = getDeps4 f0' ft' x' g'
           return (DPregel f0' ft' x' g' (ds, a))
    depCheck (DGMap f g a) = 
        do f' <- depCheck f
           g' <- depCheck g
           let ds = getDeps f' g'
           return (DGMap f' g' (ds, a))
    depCheck (DGZip g1 g2 a) = 
        do g1' <- depCheck g1
           g2' <- depCheck g2
           let ds = getDeps g1' g2'
           return (DGZip g1' g2' (ds, a))
    depCheck (DGIter f0 ft x g a) = 
        do f0' <- depCheck f0
           ft' <- depCheck ft
           x' <- depCheck x
           g' <- depCheck g
           let ds = getDeps4 f0' ft' x' g'
           return (DGIter f0' ft' x' g' (ds, a))
    depCheck (DGVar v a) = 
        do v' <- depCheck v 
           return (DGVar v' (getDep v', a))


instance DepCheckable (DExpr a) (DExpr (DDeps, a)) (NormalizeData1 a) where
    depCheck (DIf p t e a) = 
        do p' <- depCheck p
           t' <- depCheck t
           e' <- depCheck e
           let ds = getDeps3 p' t' e'
           return (DIf p' t' e' (ds, a))
    depCheck (DTuple es a) =
        do es' <- mapM depCheck es 
           let ds = getDepsX es'
           return (DTuple es' (ds, a))
    depCheck (DFunAp f es a) =
        do f' <- depCheck f
           es' <- mapM depCheck es
           let ds = normalizeDepsX (getDep f' : map getDep es')
           return (DFunAp f' es' (ds, a))
    depCheck (DConsAp c es a) =
        do c' <- depCheck c
           es' <- mapM depCheck es
           let ds = normalizeDepsX (getDep c' : map getDep es')
           return (DConsAp c' es' (ds, a))
    depCheck (DFieldAcc t fs a) =
        do t' <- depCheck t
           fs' <- mapM depCheck fs
           let ds = normalizeDepsX (getDep t' : map getDep fs')
           return (DFieldAcc t' fs' (ds, a))
    depCheck (DFieldAccE e fs a) =
        do e' <- depCheck e
           fs' <- mapM depCheck fs
           let ds = normalizeDepsX (getDep e' : map getDep fs')
           return (DFieldAccE e' fs' (ds, a))
    depCheck (DAggr a' e g es a) =
        do let a'' = nullDep a'
           g' <- depCheck g
           (depth, _) <- addNames' True (getNames g)
           e' <- depCheck e
           es' <- mapM depCheck es
           popNames
           let ds = normalizeDepsX (getDep g' : map (getDepL depth) (e':es'))
           return (DAggr a'' e' g' es' (ds, a))
    depCheck (DVExp v a) =
        do v' <- depCheck v
           return (DVExp v' (getDep v', a))
    depCheck (DCExp c a) = 
        do c' <- depCheck c 
           return (DCExp c' (getDep c', a))

instance DepCheckable (DGen a) (DGen (DDeps, a)) (NormalizeData1 a) where
    depCheck (DGenI a) = 
        do d <- checkDefined "v"           -- hard-coded...
           return (DGenI ([d], a))
    depCheck (DGenO a) = 
        do d <- checkDefined "v"
           return (DGenO ([d], a))
    depCheck (DGenG a) = 
        do return (DGenG ([], a))
    depCheck (DGenTermG a) = 
        do return (DGenTermG ([], a))

instance DepCheckable (DEdge a) (DEdge (DDeps, a)) (NormalizeData1 a) where
    depCheck (DEdge a) = 
        do d <- checkDefined "e"
           return (DEdge ([d], a))

instance DepCheckable (DTableExpr a) (DTableExpr (DDeps, a)) (NormalizeData1 a) where
    depCheck (DPrev v a) = 
        do v' <- depCheck v 
           d <- checkDefined "prev"
           return (DPrev v' (d:getDep v', a))
    depCheck (DCurr v a) = 
        do v' <- depCheck v 
           d <- checkDefined "curr"
           return (DCurr v' (d:getDep v', a))
    depCheck (DVal v a) = 
        do v' <- depCheck v 
           return (DVal v' (getDep v', a))

checkDefined n = 
  do binds <- getEnvBinds
     case (lookup n binds) of
       Just d -> return (n, d)
       Nothing -> error ("\nUndefined in dependency check " ++ n ++ " in " ++ show binds)

-- do not use these depChecks in field/fun/var/constructor definitions
instance DepCheckable (DField a) (DField (DDeps, a)) (NormalizeData1 a) where
    depCheck (DField f a) = 
        do d <- checkDefined f
           return (DField (getDVNewName d) ([d], a))
--    depCheck (DFfst a) = 
--        do return (DFfst ([], a))
--    depCheck (DFsnd a) = 
--        do return (DFsnd ([], a))

instance DepCheckable (DFun a) (DFun (DDeps, a)) (NormalizeData1 a) where
    depCheck (DFun f a) = 
        do d <- checkDefined f
           return (DFun (getDVNewName d) ([d], a))
    depCheck (DBinOp f a) = 
        do d <- checkDefined f
           return (DBinOp (getDVNewName d) ([d], a))

instance DepCheckable (DVar a) (DVar (DDeps, a)) (NormalizeData1 a) where
    depCheck (DVar v a) = 
        do d <- checkDefined v
           return (DVar (getDVNewName d) ([d], a))

instance DepCheckable (DConstructor a) (DConstructor (DDeps, a)) (NormalizeData1 a) where
    depCheck (DConstructor c a) = 
        do d <- checkDefined c
           return (DConstructor (getDVNewName d) ([d], a))