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

module ASTTrans
  (transIdElim, transCommElim, transInactivate, collectDPregel,
   dPregelStepArg, isDPregelFixStepArg, collectDConstr, DPregelStepArg,
   collectDependDSmplDef) where

import Spec
import ASTData

import Data.Maybe
import Data.List
import Numeric (showHex)

-- import GenSMT as GS
import Debug.Trace as DT

class ASTTrans a where
  elimId :: a -> a

instance ASTTrans (DProgramSpec a) where
  elimId (DProgramSpec rs p a) = DProgramSpec rs (elimId p) a

instance ASTTrans (DProg a) where
  elimId (DProg f defs e a) = DProg f (map elimId defs) e a

instance ASTTrans (DGroundDef a) where
  elimId (DGDefVC d a) = DGDefVC (elimId d) a
  elimId x = x

instance ASTTrans (DDefVertComp a) where
  elimId (DDefVertComp f ds e a) = DDefVertComp f (map elimId ds) e a

instance ASTTrans (DSmplDef a) where
  elimId (DDefVar f ds e a) = DDefVar f (map elimId ds) (elimId e) a
  elimId (DDefFun f vs ds e a) = DDefFun f vs (map elimId ds) (elimId e) a
  elimId x = x
  -- elimId (DDefTuple vs ds e a) = undefined

instance ASTTrans (DExpr a) where
  elimId (DFunAp f es a) = DFunAp f (map elimId es) a
  elimId (DAggr agg e g es a) = DAggr agg e g (idCheck e agg ++ es) a
  elimId x = x

idCheck :: DExpr a -> DAgg a -> [DExpr a]
idCheck e (DAggMin a) = [DFunAp (DBinOp "<" a) [e, DCExp (DCInt maxint a) a] a]
                        where maxint = 99999999
idCheck e (DAggMax a) =
  [DFunAp (DBinOp ">" a)
          [e, DFunAp (DFun "neg" a) [DCExp (DCInt maxint a) a] a] a]
  where maxint = 99999999
idCheck e (DAggSum a) = [DFunAp (DBinOp "!=" a) [e, DCExp (DCInt 0 a) a] a]
idCheck e (DAggProd a) = [DFunAp (DBinOp "!=" a) [e, DCExp (DCInt 1 a) a] a]
idCheck e (DAggAnd a) = [DFunAp (DBinOp "!=" a) [e, DCExp (DCBool True a) a] a]
idCheck e (DAggOr a) = [DFunAp (DBinOp "!=" a) [e, DCExp (DCBool False a) a] a]
idCheck e _ = []

transIdElim :: (Show a) => DProgramSpec a -> DProgramSpec a
transIdElim = elimId

class ASTTraverse a b | a -> b where
  collectDPregel :: a -> [b]
  collectDPregel _ = []

type DPregelStepArg a = (DFun a, DTermination a)

instance ASTTraverse (DProgramSpec a) (DPregelStepArg a) where
  collectDPregel (DProgramSpec _ prog _) = collectDPregel prog

instance ASTTraverse (DProg a) (DPregelStepArg a) where
  collectDPregel (DProg f gdefs gexpr _) =
    concatMap collectDPregel gdefs ++ collectDPregel gexpr

instance ASTTraverse (DGroundDef a) (DPregelStepArg a) where
  collectDPregel (DGDefGV gv _) = collectDPregel gv
  collectDPregel (DGDefGF gf _) = collectDPregel gf
  collectDPregel _ = []

instance ASTTraverse (DDefGraphVar a) (DPregelStepArg a) where
  collectDPregel (DDefGraphVar _ gexpr _) = collectDPregel gexpr

instance ASTTraverse (DDefGraphFun a) (DPregelStepArg a) where
  collectDPregel (DDefGraphFun _ _ gvs gexpr _) =
    concatMap collectDPregel gvs ++ collectDPregel gexpr

instance ASTTraverse (DGraphExpr a) (DPregelStepArg a) where
  collectDPregel (DPregel initfn stepfn term gexpr _) =
    (stepfn, term) : collectDPregel gexpr
  collectDPregel (DGMap _ gexpr _) = collectDPregel gexpr
  collectDPregel (DGZip gexpr1 gexpr2 _) =
    collectDPregel gexpr1 ++ collectDPregel gexpr2
  collectDPregel (DGIter _ _ _ gexpr _) = collectDPregel gexpr
  collectDPregel _ = []

dPregelStepArg :: String -> [DPregelStepArg a] -> Maybe (DTermination a)
dPregelStepArg name [] = Nothing
dPregelStepArg name ((DFun step _, term) : ds) =
  if name == step then Just term else dPregelStepArg name ds
dPregelStepArg name (_:ds) = dPregelStepArg name ds

isDPregelFixStepArg :: String -> [DPregelStepArg a] -> Bool
isDPregelFixStepArg name [] = False
isDPregelFixStepArg name ((DFun step _, DTermF _) : ds) =
  if name == step then True else isDPregelFixStepArg name ds
isDPregelFixStepArg name ((DFun step _, _) : ds) =
  if name == step then False else isDPregelFixStepArg name ds
isDPregelFixStepArg name (_:ds) = isDPregelFixStepArg name ds

collectDConstr :: DExpr a -> [DConstructor a]
collectDConstr (DIf e1 e2 e3 _) = collectDConstr e2 ++ collectDConstr e3
collectDConstr (DConsAp con es _) = [con]
collectDConstr _ = []

collectDependDSmplDef :: DSmplDef a -> [String]
collectDependDSmplDef (DDefFun f vars locals exp _) =
  collectDependLE (getNames f) vars locals exp

collectDependDSmplDef (DDefVar v locals exp _) =
  collectDependLE [] [v] locals exp

collectDependDSmplDef (DDefTuple vars locals exp _) =
  collectDependLE [] vars locals exp

collectDependLE :: [String] -> [DVar a] -> [DSmplDef a] -> DExpr a -> [String]
collectDependLE names vars locals exp =
  removeDup (freevs0 ++ freevs1)
  where boundvs = names ++ concatMap getNames vars ++ concatMap getNames locals
        freevs0 = removeDep boundvs (concatMap collectDependDSmplDef locals)
        freevs1 = removeDep boundvs (collectDependDExpr exp)

collectDependDExpr :: DExpr a -> [String]
collectDependDExpr (DIf e1 e2 e3 _) = concatMap collectDependDExpr [e1,e2,e3]
collectDependDExpr (DTuple es _) = concatMap collectDependDExpr es
collectDependDExpr (DFunAp fn es _) = concatMap collectDependDExpr es
collectDependDExpr (DConsAp con es _) = concatMap collectDependDExpr es
collectDependDExpr (DAggr ag e gen es _) = concatMap collectDependDExpr (e:es)
collectDependDExpr (DCheckTerm e _) = collectDependDExpr e
collectDependDExpr (DVExp v a) = getNames v
collectDependDExpr _ = []

removeDep :: (Eq a) => [a] -> [a] -> [a]
removeDep names deps = foldr delete deps names

removeDup :: (Eq a) => [a] -> [a]
removeDup [] = []
removeDup (x:xs) = x : delete x (removeDup xs)

transCommElim :: DGroundDef a -> DGroundDef a
transCommElim x = x

transInactivate :: DGroundDef a -> DGroundDef a
transInactivate = id

-- The following code is developed by Morihata

{-
transCommElim :: DProgramSpec a -> DProgramSpec a
transCommElim (DProgramSpec (r:rs) p a) = 
    let (con, flds, r') = transDataDecl r 
    in DProgramSpec (r' : rs) (transProg (con,flds) p) a

transDataDecl (DRecordSpec con flds a) = 
    let DField _ a' = fst (head flds)
    in (con, flds, 
        DRecordSpec con ((DField "__changed" a', DTBool a') : flds) a)

transProg con (DProg a defs b c) = DProg a (map (transDef con) defs) b c

transDef con (DGDefVI (DDefVertInit f defs  expr b) c) 
    = (DGDefVI (DDefVertInit f (map (transLet con False) defs) 
                                 (transExpr con False expr) b) c) 
transDef con (DGDefSmpl sd a) = DGDefSmpl (transLet con False sd) a
transDef con (DGDefVC (DDefVertComp f ldefs expr a) b) 
    = (DGDefVC (DDefVertComp f (map (transLet con True) ldefs) 
                                 (transExpr con True expr) a) b) 
transDef c a = a

transLet c flg (DDefFun f vars defs expr a) = 
    DDefFun f vars (map (transLet c flg) defs) (transExpr c flg expr) a
transLet c flg (DDefVar v defs expr a) = 
    DDefVar v (map (transLet c flg) defs) (transExpr c flg expr) a

transExpr c flg (DIf e1 e2 e3 a) = 
    DIf (transExpr c flg e1) (transExpr c flg e2) (transExpr c flg e3) a
transExpr c flg (DFunAp f exprs a) = 
    DFunAp f (map (transExpr c flg) exprs) a
transExpr c@(con,flds) flg (DConsAp con' exprs a) = 
    let exprs' = map (transExpr c flg) exprs
        equiv (DConstructor a _) (DConstructor b _) = a == b
    in if equiv con' con 
       then if flg 
            then let anyOf = foldr1 (\p r-> DFunAp (DBinOp "||" a) [p,r] a) 
                     preds = zipWith chkEq exprs' flds 
                     chkEq e (f,t) = DFunAp (DBinOp "==" a) 
                                     [e, DFieldAcc (DPrev (DVar "v" a) a) [f] a] a
                 in DConsAp con (anyOf preds : exprs') a
            else DConsAp con (DCExp (DCBool True a) a : exprs') a 
       else DConsAp con exprs' a
transExpr c flg  (DAggr agop e1 ge e2 a) = 
    let comChk = DFieldAcc (DPrev (DVar "u" a) a) [DField "__changed" a] a
    in DAggr agop e1 ge (comChk : e2) a
transExpr c flg a = a

transInactivate (DProgramSpec rs (DProg a defs expr c) d) = 
   (DProgramSpec rs (DProg a defs (transFregel expr) c) d)

transFregel (DPregel init step (DTermF a) expr b)  = 
    DPregel init step (DTermV2H a) expr b
-}