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

{-

Instantiating VertComp/VertInit/GraphFun for each use of them (7th step)

Assumption: 
  - no function definition (other than the above)
  - let-bindings are located at the top level (of the above functions)

-}

module TypeInstantiation where

import Spec
import ASTData
import TypeChecker

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

type DTypeInstances = (DGroundDef DASTData, [(DTypeInfo, DGroundDef DASTData)])
type DTypeInstanceBinding = ([DVarName], DTypeInstances)
type DTypeInstanceBindings = [DTypeInstanceBinding]

-- uniq id, stack, stack head indices
type DEnvTI = (DUnique, DTypeInstanceBindings, [Int])

class TypeInstantiable a  where
    typeInstantiation :: a -> State DEnvTI a  -- functions are replaced with their instances
    typeInstantiation x = return x            -- default

-- entry point
runTypeInstantiation :: forall (t :: * -> *) a . TypeInstantiable (t DASTData) => (t DASTData) -> DUnique -> (t DASTData, DUnique)
runTypeInstantiation p uid = let (p', (uid', _, _)) = runState (typeInstantiation p) (uid, [], [])
                             in (p', uid')
-- TODO: built-in functions?

addNewBind :: DVarName -> DGroundDef DASTData -> State DEnvTI ()
addNewBind n def = do (i, bds, ks) <- get
                      let bds' = rec bds
                      put (i, bds', ks)
  where rec [] = []
        rec ((ns, (org, ins)):xs) = case elemIndex n ns of
                                      Just _ -> (ns, (org, (typeOf (getData def), def):ins)):xs
                                      Nothing -> (ns, (org, ins)):rec xs
        
getNewName :: String -> State (DEnvTI) DVarName
getNewName s = 
    do (i, bds, ks) <- get 
       let (n, i') = genNewName i (getBaseName s)
       put (i', bds, ks)
       return n

getBinds :: State DEnvTI DTypeInstanceBindings
getBinds = do (i, bds, ks) <- get
              return bds

addBindsByDefs :: [DGroundDef DASTData] -> State DEnvTI ()
addBindsByDefs defs = do addBinds (map (\def -> (getNames def, (def, []))) defs)

addBindsByDefsS :: [DSmplDef DASTData] -> State DEnvTI ()
addBindsByDefsS defs = addBindsByDefs (map (\def -> DGDefSmpl def (getData def)) defs)

addBindsByDefsV :: [DDefGraphVar DASTData] -> State DEnvTI ()
addBindsByDefsV defs = addBindsByDefs (map (\def -> DGDefGV def (getData def)) defs)

addBinds :: DTypeInstanceBindings -> State DEnvTI ()
addBinds bds = do (i, bds', ks) <- get
                  put (i, bds++bds', length bds:ks)


lookupBind :: DVarName -> State DEnvTI (Maybe (DTypeInstances, Int))
lookupBind n = do bds <- getBinds
                  return (rec bds)
    where rec [] = Nothing
          rec ((ns, ins):xs) = case elemIndex n ns of 
                                 Just i -> (Just (ins, if length ns == 1 then -1 else i))
                                 Nothing -> rec xs          

popBinds :: State DEnvTI DTypeInstanceBindings
popBinds = do xx <- get
              let (i, bds, k:ks) = xx
              put (i, drop k bds, ks)
              return (take k bds)

isSmplDef (DGDefSmpl def a) = True
isSmplDef _ = False

isVertFunDef (DGDefVI def a) = True
isVertFunDef (DGDefVC def a) = True
isVertFunDef _ = False

isGraphFunVarDef (DGDefGV def a) = True
isGraphFunVarDef (DGDefGF def a) = True
isGraphFunVarDef _ = False

instance TypeInstantiable (DProgramSpec DASTData) where
    typeInstantiation (DProgramSpec rs p a) = 
        do p' <- typeInstantiation p
           return (DProgramSpec rs p' a)

instance TypeInstantiable (DProg DASTData) where
    typeInstantiation (DProg f defs e a) = 
        do -- add the defs. to the env
           addBindsByDefs defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = concatMap (\(_, (_, ds)) -> map snd ds) defs'
           return (DProg f defs'' e' a)

instance TypeInstantiable (DGraphExpr DASTData) where
    typeInstantiation (DPregel f0 ft x g a) = 
        do f0' <- typeInstantiation f0
           ft' <- typeInstantiation ft
           x' <- typeInstantiation x
           g' <- typeInstantiation g
           return (DPregel f0' ft' x' g' a)
    typeInstantiation (DGMap f g a) = 
        do f' <- typeInstantiation f
           g' <- typeInstantiation g
           return (DGMap f' g' a)
    typeInstantiation (DGZip g1 g2 a) = 
        do g1' <- typeInstantiation g1
           g2' <- typeInstantiation g2
           return (DGZip g1' g2' a)
    typeInstantiation (DGIter f0 ft x g a) = 
        do f0' <- typeInstantiation f0
           ft' <- typeInstantiation ft
           x' <- typeInstantiation x
           g' <- typeInstantiation g
           return (DGIter f0' ft' x' g' a)
    typeInstantiation (DGVar v a) = 
        do v' <- typeInstantiation v 
           return (DGVar v' a)

------------------ giving new names when instantiating
class NameRefreshable a where
    refreshNames :: a -> State DEnvTI a

instance NameRefreshable (DGroundDef DASTData) where
    refreshNames (DGDefVC d a) = do d' <- refreshNames d
                                    return (DGDefVC d' a)
    refreshNames (DGDefVI d a) = do d' <- refreshNames d
                                    return (DGDefVI d' a)
    refreshNames (DGDefSmpl d a) = do d' <- refreshNames d
                                      return (DGDefSmpl d' a)
    refreshNames (DGDefGV d a) =
        do d' <- refreshNames d
           return (DGDefGV d' a)
    refreshNames (DGDefGF d a) = do d' <- refreshNames d
                                    return (DGDefGF d' a)

replaceNames :: [(DVarName, DVarName)] -> DExpr DASTData -> DExpr DASTData
replaceNames subsv e = rec e
  where
    rec (DIf p t e a) = 
        let p' = rec p
            t' = rec t
            e' = rec e
        in (DIf p' t' e' a)
    rec (DTuple es a) =
        let es' = map rec es 
        in (DTuple es' a)
    rec (DFunAp f es a) =
        let es' = map rec es
        in (DFunAp f es' a)
    rec (DConsAp c es a) =
        let es' = map rec es
        in (DConsAp c es' a)
    rec (x@(DFieldAcc t fs a)) = x
    rec (x@(DFieldAccE e fs a)) = x
    rec (DAggr a' e g es a) =
        let e'  = rec e
            es' = map rec es
        in (DAggr a' e' g es' a)
    rec (x@(DVExp (DVar v av) a)) = 
        case lookup v subsv of
          Just v' -> (DVExp (DVar v' av) a)
          Nothing -> x
    rec (x@(DCExp c a)) = x


instance NameRefreshable (DFun DASTData) where
    refreshNames (DFun f a) =
        do f' <- getNewName f
           return (DFun f' a)
    refreshNames (DBinOp f a) =
        do f' <- getNewName f
           return (DBinOp f' a)

instance NameRefreshable (DVar DASTData) where
    refreshNames (DVar v a) =
        do v' <- getNewName v
           return (DVar v' a)

instance NameRefreshable (DDefVertComp DASTData) where
    refreshNames (DDefVertComp f defs e a) = 
        do f' <- refreshNames f
           return (DDefVertComp f' defs e a)

instance NameRefreshable (DDefVertInit DASTData) where
    refreshNames (DDefVertInit f defs e a) = 
        do f' <- refreshNames f
           return (DDefVertInit f' defs e a)

instance NameRefreshable (DDefGraphVar DASTData) where
    refreshNames (DDefGraphVar v e a) = 
        do v' <- refreshNames v
           return (DDefGraphVar v' e a)

replaceNamesF subse (DFun f a) =
    case lookup f subse of
      Nothing -> (DFun f a)
      Just f' -> (DFun f' a)

replaceNamesV subse (DVar v a) =
    case lookup v subse of
      Nothing -> (DVar v a)
      Just v' -> (DVar v' a)

replaceNamesG subse (DPregel f0 ft x g a) = 
    let f0' = replaceNamesF subse f0
        ft' = replaceNamesF subse ft
        x' = x
        g' = replaceNamesG subse g
    in (DPregel f0' ft' x' g' a)
replaceNamesG subse (DGMap f g a) = 
    let f' = replaceNamesF subse f
        g' = replaceNamesG subse g
    in (DGMap f' g' a)
replaceNamesG subse (DGZip g1 g2 a) = 
    let g1' = replaceNamesG subse g1
        g2' = replaceNamesG subse g2
    in (DGZip g1' g2' a)
replaceNamesG subse (DGIter f0 ft x g a) = 
    let f0' = replaceNamesF subse f0
        ft' = replaceNamesF subse ft
        x' = x
        g' = replaceNamesG subse g
    in  (DGIter f0' ft' x' g' a)
replaceNamesG subse (DGVar v a) = 
    let v' = replaceNamesV subse v
    in (DGVar v' a)

replaceNamesGV subse (DDefGraphVar v e a) = 
    let e' = (replaceNamesG subse e)
    in (DDefGraphVar v e' a)

instance NameRefreshable (DDefGraphFun DASTData) where
    refreshNames (DDefGraphFun f v defs e a) = 
        do f' <- refreshNames f
           v' <- refreshNames v
           defs' <- mapM refreshNames defs 
           let subse = [(getName f, getName f'), (getName v, getName v')] ++ zipWith (\o n -> (getName o, getName n)) defs defs'
               e' = replaceNamesG subse e
               defs'' = map (replaceNamesGV subse) defs'
           return (DDefGraphFun f' v' defs'' e' a)

instance NameRefreshable (DSmplDef DASTData) where
    refreshNames (DDefVar v [] e a) = 
        do v' <- refreshNames v
           let subse = [(getName v, getName v')]
               e' = replaceNames subse e
           return (DDefVar v [] e' a)
    refreshNames (DDefTuple vs [] e a) = 
        do vs' <- mapM refreshNames vs
           let subse = zip (map getName vs) (map getName vs')
               e' = replaceNames subse e
           return (DDefTuple vs' [] e' a)
    refreshNames x = error $ "something wrong in refreshNames : x = " ++ show x

----------------------------------------------------------------

instance TypeInstantiable (DGroundDef DASTData) where
    typeInstantiation (DGDefVC d a) = do d' <- typeInstantiation d
                                         return (DGDefVC d' a)
    typeInstantiation (DGDefVI d a) = do d' <- typeInstantiation d
                                         return (DGDefVI d' a)
    typeInstantiation (DGDefSmpl d a) = do d' <- typeInstantiation d
                                           return (DGDefSmpl d' a)
    typeInstantiation (DGDefGV d a) = do d' <- typeInstantiation d
                                         return (DGDefGV d' a)
    typeInstantiation (DGDefGF d a) = do d' <- typeInstantiation d
                                         return (DGDefGF d' a)

makeDefs defs = concatMap (\(_, (_, ds)) -> map (\(DGDefSmpl def _) -> def) (map snd ds)) defs
makeDefsGV defs = concatMap (\(_, (_, ds)) -> map (\(DGDefGV def _) -> def) (map snd ds)) defs

instance TypeInstantiable (DDefVertComp DASTData) where
    typeInstantiation (DDefVertComp f defs e a) = 
        do addBindsByDefsS defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefs defs'
           return (DDefVertComp f defs'' e' a)

instance TypeInstantiable (DDefVertInit DASTData) where
    typeInstantiation (DDefVertInit f defs e a) = 
        do addBindsByDefsS defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefs defs'
           return (DDefVertInit f defs'' e' a)

instance TypeInstantiable (DDefGraphVar DASTData) where
    typeInstantiation (DDefGraphVar v e a) = 
        do e' <- typeInstantiation e
           return (DDefGraphVar v e' a)

instance TypeInstantiable (DDefGraphFun DASTData) where
    typeInstantiation (DDefGraphFun f v defs e a) = 
        do addBindsByDefsV defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefsGV defs'
           return (DDefGraphFun f v defs'' e' a)

instance TypeInstantiable (DSmplDef DASTData) where
    typeInstantiation (DDefFun f vs defs e a) = 
        do addBindsByDefsS defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefs defs'
           return (DDefFun f vs defs'' e' a)
    typeInstantiation (DDefVar v defs e a) = 
        do addBindsByDefsS defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefs defs'
           return (DDefVar v defs'' e' a)
    typeInstantiation (DDefTuple vs defs e a) = 
        do addBindsByDefsS defs
           e' <- typeInstantiation e
           defs' <- popBinds
           let defs'' = makeDefs defs'
           return (DDefTuple vs defs'' e' a)

instance TypeInstantiable (DTermination DASTData) where
    typeInstantiation (DTermF a) = return (DTermF a)
    typeInstantiation (DTermI e a) = 
        do e' <- typeInstantiation e
           return (DTermI e' a)
    typeInstantiation (DTermU e a) = 
        do e' <- typeInstantiation e
           return (DTermU e' a)


instance TypeInstantiable (DExpr DASTData) where
    typeInstantiation (DIf p t e a) = 
        do p' <- typeInstantiation p
           t' <- typeInstantiation t
           e' <- typeInstantiation e
           return (DIf p' t' e' a)
    typeInstantiation (DTuple es a) =
        do es' <- mapM typeInstantiation es 
           return (DTuple es' a)
    typeInstantiation (DFunAp f es a) =
        do f' <- typeInstantiation f
           es' <- mapM typeInstantiation es
           return (DFunAp f' es' a)
    typeInstantiation (DConsAp c es a) =
        do c' <- typeInstantiation c
           es' <- mapM typeInstantiation es
           return (DConsAp c' es' a)
    typeInstantiation (DFieldAcc t fs a) =
        do t' <- typeInstantiation t
           fs' <- mapM typeInstantiation fs
           return (DFieldAcc t' fs' a)
    typeInstantiation (DFieldAccE e fs a) =
        do e' <- typeInstantiation e
           fs' <- mapM typeInstantiation fs
           return (DFieldAccE e' fs' a)
    typeInstantiation (DAggr a' e g es a) =
        do a'' <- typeInstantiation a'
           g' <- typeInstantiation g
           e' <- typeInstantiation e
           es' <- mapM typeInstantiation es
           return (DAggr a'' e' g' es' a)
    typeInstantiation (DVExp v a) =
        do v' <- typeInstantiation v
           return (DVExp v' a)
    typeInstantiation (DCExp c a) = 
        do c' <- typeInstantiation c 
           return (DCExp c' a)

unifyUpdate def t = 
    do s <- unify' [(typeOf (getData def), t)]
       return (mapData (\a -> a { typeOf = apply s (typeOf a)}) def)

instance TypeInstantiable (DFun DASTData) where
    typeInstantiation (DBinOp f a) = instantiation DBinOp f a
    typeInstantiation (DFun f a) = instantiation DFun f a

instance TypeInstantiable (DVar DASTData) where
    typeInstantiation (DVar v a) = instantiation DVar v a

instantiation
  :: (DVarName -> DASTData -> b)
     -> DVarName
     -> DASTData
     -> State DEnvTI b
instantiation build v a =
        do r <- lookupBind v
           case r of 
             Nothing -> return (build v a)  -- temporary things?
             Just ((def, ins), pos) -> 
                 case def of 
                   (DGDefGF _ _) -> do def' <- refreshNames def        -- always duplicate it
                                       let defe'' = unifyUpdate def' (typeOf a)
                                       case defe'' of
                                         Right def'' ->
                                             do def''' <- typeInstantiation def'' -- instantiate it
                                                addNewBind v def'''              -- add it to the env
                                                return (build (getNames def''' !! (max pos 0)) a)  -- use it
                                         Left str -> error (str ++ "\nerror during unifying " ++ show def' ++ " and " ++ show (typeOf a) ++ "\nv="++show v ++ "\na="++show a)
                   otherwise -> 
                       case lookupAt pos (typeOf a) ins of
                         Just def' -> return (build (getNames def' !! (max pos 0)) a) -- already instantiated
                         Nothing -> if typeAt pos (typeOf (getData def)) == (typeOf a) 
                                     then do def'' <- typeInstantiation def
                                             addNewBind v def''  -- use the original one
                                             return (build (getNames def'' !! (max pos 0)) a)
                                     else do def' <- refreshNames def        -- gives new names
                                             let defe'' = unifyUpdate def' (typeOf a)
                                             case defe'' of
                                               Right def'' ->
                                                   do def''' <- typeInstantiation def'' -- instantiate it
                                                      addNewBind v def'''              -- add it to the env
                                                      return (build (getNames def''' !! (max pos 0)) a)  -- use it
                                               Left str -> error (str ++ "\nerror during unifying2 " ++ show def' ++ " and " ++ show (typeOf a) ++ "\nv="++show v ++ "\na="++show a ++ "\nins="++show ins ++ "\npos="++show pos )

typeAt (-1) t = t
typeAt pos (DTypeTerm "(,)" ts) = (ts!!pos)

lookupBy' :: Eq b => (a -> b) -> b -> [(a, x)] -> Maybe x
lookupBy' _ _ [] = Nothing
lookupBy' f key ((a,x):xs) = if f a == key then Just x else lookupBy' f key xs

lookupAt :: Int -> DTypeInfo -> [(DTypeInfo, x)] -> Maybe x
lookupAt pos t ins | pos == -1 = lookup t ins
                   | otherwise = lookupBy' (typeAt pos) t ins   -- tuple case




instance TypeInstantiable (DAgg DASTData)
instance TypeInstantiable (DGen DASTData)
instance TypeInstantiable (DEdge DASTData)
instance TypeInstantiable (DTableExpr DASTData)
instance TypeInstantiable (DField DASTData)
instance TypeInstantiable (DConstructor DASTData)
instance TypeInstantiable (DConst DASTData)