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

-- module GenSMT (genCommElimSpec, genInactivateSpec, PregelStepfn) where
module GenSMT
  (genCommElimSpec, genInactivateSpec, collectVCStepfns, collectVCFixStepfns,
   PregelStepfn, SMTOptimizable, stepfnName, applicableElimComm, applicableV2H)
where

import Spec
import ASTTrans
import ASTData
import TypeChecker
import Data.Char
import Data.List
import Debug.Trace
import System.IO.Unsafe

import DependencySimple (computeDependency, depAnnotate)

type Ty a = (DTypeTerm, a)

ppp v = unsafePerformIO $ do { putStrLn (show v); return v;}

sameConstructor :: DConstructor (Ty a) -> DConstructor (Ty b) -> Bool
sameConstructor (DConstructor con1 _) (DConstructor con2 _) = con1 == con2

vE, tRP, tBool :: String
vE = "|e|"    -- name of edge
tRP = "_RP"
tBool = "Bool"

prelude :: String
prelude = unlines $ 
  ["(declare-datatypes () ((_RP _pinf _ninf _nan (_rv (_val Real)))))",
   "(define-fun _isReal ((x _RP)) Bool",
   "  (not (or (= x _pinf) (or (= x _ninf) (= x _nan)))))",
   "(define-fun _eq ((x _RP) (y _RP)) Bool",
   "  (and (_isReal x) (and (_isReal y) (= (_val x) (_val y)))))",
   "(define-fun _neq ((x _RP) (y _RP)) Bool",
   "  (not (_eq x y)))",
   "(define-fun _gt ((x _RP) (y _RP)) Bool",
   "  (or (= x _pinf) (or (= _ninf y)",
   "  (and (_isReal x) (and (_isReal y) (> (_val x) (_val y)))))))",
   "(define-fun _lt ((x _RP) (y _RP)) Bool",
   "  (or (= x _ninf) (or (= y _pinf)",
   "  (and (_isReal x) (and (_isReal y) (< (_val x) (_val y)))))))",
   "(define-fun _add ((x _RP) (y _RP)) _RP",
   "  (ite (and (_isReal x) (_isReal y)) (_rv (+ (_val x) (_val y))) _nan))",
   "(define-fun _mult ((x _RP) (y _RP)) _RP",
   "  (ite (and (_isReal x) (_isReal y)) (_rv (* (_val x) (_val y))) _nan))",
   "(define-fun _div ((x _RP) (y _RP)) _RP",
   "  (ite (and (_isReal x) (_isReal y)) (_rv (/ (_val x) (_val y))) _nan))",
   "(define-fun _neg ((x _RP)) _RP",
   "  (ite (= x _pinf) _ninf (ite (= x _ninf) _pinf",
   "  (ite (= x _nan) _nan (_rv (- (_val x)))))))",
   "(define-fun _min ((x _RP) (y _RP)) _RP (ite (_lt x y) x y))",
   "(define-fun _max ((x _RP) (y _RP)) _RP (ite (_gt x y) x y))\n"]

genAuxDecls :: (Show a) =>
  DRecordSpec (Ty a) -> [DGroundDef (Ty a)] -> [String]
genAuxDecls rs gdefs = [genDataDecls rs] ++ concatMap genAuxFunc gdefs

genDataDecls :: DRecordSpec (Ty a) -> String
genDataDecls (DRecordSpec (DConstructor dname _) fields _) =
  genZ3DeclareDatatypes
    ("((" ++ dname ++ " (" ++ toHL dname ++ " " ++
     unwords (map genF fields) ++ ")))")
    True
  where genF (DField fname _, ty) = "(" ++ fname ++ " " ++ genTy ty ++ ")"

toHL :: String -> String
toHL (c:s) = toLower c : s

genTy :: DType a -> String
genTy (DTInt _ ) = tRP
genTy (DTBool _ ) = tBool
genTy (DTDouble _ ) = tRP
genTy _  = undefined -- others may not be supported.

genAuxFunc :: (Show a) => DGroundDef (Ty a) -> [String]
genAuxFunc (DGDefVI (DDefVertInit (DFun fn _) _ expr _) _) =  [] 
genAuxFunc (DGDefSmpl sd _) = [genSmplDef [] sd]
genAuxFunc (DGDefVC (DDefVertComp (DFun fn _) ldefs expr _) _) = []
genAuxFunc _ = []

getVC :: DProgramSpec (Ty a) -> [(String, [DSmplDef (Ty a)], DExpr (Ty a))]
getVC pr@(DProgramSpec rs (DProg _ defs _ _) _) = rs
  where genVCdefs :: DGroundDef a -> [(String, [DSmplDef a], DExpr a)]
        genVCdefs (DGDefVC (DDefVertComp (DFun fn _) ldefs expr _) _) =
          if isDPregelFixStepArg fn dps then [(fn, ldefs, expr)] else []
        genVCdefs _ = []
        dps = collectDPregel pr
        rs = concatMap genVCdefs defs

type PregelStepfn a =
  (String, [DSmplDef a], DExpr a, DTermination a, DRecordSpec a)

type SMTOptimizable = (String, Bool, Bool)
-- (name of a step function,
--  applicable eliminate-commucation optimzation,
--  applicable vote-to-halt optimzation)

stepfnName :: SMTOptimizable -> String
stepfnName (name, _, _) = name

applicableElimComm :: SMTOptimizable -> Bool
applicableElimComm (_, b, _) = b

applicableV2H :: SMTOptimizable -> Bool
applicableV2H (_, _, b) = b

-- If pr is a DProgramSpec after typechecking, a is (ASTData.DTypeInfo, a').
-- collectVCStepfns :: (Show a) => DProgramSpec a -> [PregelStepfn a]
collectVCStepfns :: (Show a) => DProgramSpec (Ty a) -> [PregelStepfn (Ty a)]
collectVCStepfns pr@(DProgramSpec rs (DProg _ defs _ _) _) =
  -- ppp "collectDPregel = " `seq` ppp dps `seq`
  concatMap genVCStep defs
  where dps = collectDPregel pr
        -- genVCStep :: DGroundDef a -> [PregelStepfn a]
        genVCStep (DGDefVC (DDefVertComp (DFun fn _) ldefs expr _) _) =
          case dPregelStepArg fn dps of
            Just term -> [(fn, ldefs, expr, term, getRecSpec rs ldefs expr)]
            Nothing   -> []
        genVCStep _ = []
        -- getRecSpec :: [DRecordSpec t] -> [DSmplDef t] -> DExpr t -> DRecordSpec t
        getRecSpec rs ldefs expr =
          if null cs then undefined
          else if null (tail cs) then findRS (head cs) rs
          else ppp "DConstructor is not uniq" `seq` findRS (head cs) rs
          where cs = uniqList (getDConstrs ldefs expr)
                -- uniqList :: Eq t => [t] -> [t]
                uniqList [] = []
                uniqList (x:xs) =
                  (if or (map (sameConstructor x) xs) then [] else [x])
                  ++ uniqList xs
                -- findRS :: DConstructor a -> [DRecordSpec a] -> DRecordSpec a
                findRS con [] = undefined
                findRS con (r@(DRecordSpec c mems _):rs)
                  | sameConstructor con c = r
                  | otherwise             = findRS con rs
        -- getDConstrs :: [DSmplDef t] -> DExpr t -> [DConstructor t]
        getDConstrs ldefs expr = cs
          where cs = concatMap col ldefs ++ collectDConstr expr
                col (DDefVar _ _ e _) = collectDConstr e
                col _ = []

collectVCFixStepfns :: (Show a) => DProgramSpec (Ty a) -> [PregelStepfn (Ty a)]
collectVCFixStepfns pr = filter isTermFix (collectVCStepfns pr)
  where isTermFix (_, _, _, DTermF _, _) = True
        isTermFix _ = False

genVC :: (Show a) =>
  String -> (String, [DSmplDef (Ty a)], DExpr (Ty a)) ->
  (String, [(DAgg (Ty a), DExpr (Ty a), DGen (Ty a), [DExpr (Ty a)])],
   [String], String, [String])
genVC dn (fn, ldefs, expr) = 
    let mds = messageAggrs dn ldefs expr
        mdts = map getType mds
        vars = map (\n -> genZ3var ("m" ++ show n)) [1 .. length mds]
        vc = genZ3DefineFun
               fn
               (genParenedStr
                 (genZ3VarTy vx dn : zipWith genZ3VarTy vars mdts))
               dn
               (snd (genLetsExpr vars ldefs expr))
               True
        vx = genZ3var "vx"
        getType (DAggMin _, _, _, _) = tRP
        getType (DAggMax _, _, _, _) = tRP
        getType (DAggSum _, _, _, _) = tRP
        getType (DAggProd _, _, _, _) = tRP
        getType (DAggAnd _, _, _, _) = tBool
        getType (DAggOr _, _, _, _) = tBool
        getType (_, _, _, _) = "None"
        {-
        vars = map (\n -> "|m" ++ show n ++"|") [1 .. length mds]
        vc = "(define-fun |" ++ fn ++ "| ((|vx| " ++ dn ++ ") " ++
             (unwords $ map (\vn -> "(" ++ vn ++ " _RP)") vars) ++ ") " ++ 
             dn ++ "\n  " ++ snd (genLetsExpr vars ldefs expr) ++ ")\n"
        -}
    in (fn, mds, mdts, vc, vars) 

-- genSmplDef has to use the typing information to produce the
-- correct type.  This is TEKITOU.
genSmplDef :: (Show a) => [String] -> DSmplDef (Ty a) -> String
genSmplDef env (DDefFun (DFun fn _) vars defs expr _) =
  genZ3DefineFun
    fn
    (genParenedStr (map genArg vars))
    tRP
    (snd (genLetsExpr env defs expr))
    True
  where genArg (DVar vn _) = genParenedStr [genZ3var vn, tRP]
--  where genArg (DVar vn _) = genParenedStr [vn, tRP]

genSmplDef env (DDefVar (DVar vn _) defs expr _) =
  genZ3DeclareConst vn' tRP True ++
  genZ3AssertEq vn' (snd (genLetsExpr env defs expr)) True
  where vn' = genZ3var vn

{-
genSmplDef env (DDefVar (DVar vn _) defs expr _) =
  genZ3DeclareConst vn tRP True ++
  genZ3AssertEq vn (snd (genLetsExpr env defs expr)) True
-}

type Env = [String]

mapE :: (Show a) =>
  (Env -> m a -> (Env, String)) -> Env -> [m a] -> (Env, [String])
mapE f env [] = (env,[])
mapE f env (a:x) = let (env',a') = f env a 
                       (env'', x') = mapE f env' x
                   in (env'', a':x')

genLetsExpr :: (Show a) =>
  Env -> [DSmplDef (Ty a)] -> DExpr (Ty a) -> (Env, String)
genLetsExpr env [] expr = genExpr env expr
genLetsExpr env defs expr = 
  let ds = zip defs (map collectDependDSmplDef defs)
      defs' = map fst (sortBy compareDefs ds)
  in -- ppp defs `seq` ppp ds `seq` ppp defs' `seq`
     genLetsExpr' env defs' expr

genLetsExpr' :: (Show a) =>
  Env -> [DSmplDef (Ty a)] -> DExpr (Ty a) -> (Env, String)
genLetsExpr' env [] expr = genExpr env expr
genLetsExpr' env (d:ds) expr = 
  let (env', dr) = genLet env d
      (env'', er) = genLetsExpr' env' ds expr
  in (env'', genSexpr "let" [genParenedStr [dr], er] False)

{-
genLetsExpr env defs expr = 
  let (env', dr) = mapE genLet env defs
      (env'', er) = genExpr env' expr
      ds = zip defs (map collectDependDSmplDef defs)
      defs' = map fst (sortBy compareDefs ds)
  in ppp defs `seq` ppp ds `seq` ppp defs' `seq`
     (env'', genSexpr "let" [genParenedStr dr, er] False)
-}

compareDefs :: (DSmplDef (Ty a), [String]) -> (DSmplDef (Ty a), [String]) ->
               Ordering
compareDefs (sdef1, dep1) (sdef2, dep2) =
  let (n1:_) = getNames sdef1
      (n2:_) = getNames sdef2
  in if n2 `elem` dep1 then GT
     else if n1 `elem` dep2 then LT
     else EQ

genLet :: (Show a) => Env -> DSmplDef (Ty a) -> (Env, String)
genLet env (DDefVar (DVar vn _) defs expr _) =
  let (env', dr) = genLetsExpr env defs expr
  in  (env', genSexpr (genZ3var vn) [dr] False)
-- in  (env', genSexpr vn [dr] False)

genExpr :: (Show a) => Env -> DExpr (Ty a) -> (Env, String)
genExpr env (DIf e1 e2 e3 _) =
  let (env', es) = mapE genExpr env [e1,e2,e3]
  in (env', genSexpr "ite" es False)

genExpr env (DFunAp f@(DBinOp "!=" (tne,_)) exprs _) =
  let at1 = argTypeName1 tne
      (env', es) = mapE genExpr env exprs
  in (env', genSexpr (genEqOp "!=" at1) es False)

genExpr env (DFunAp f exprs _) 
    = let (env', es) = mapE genExpr env exprs
      in (env', genSexpr (genFunOp f) es False)

genExpr env (DConsAp (DConstructor cn _) exprs _) 
    = let (env', es) = mapE genExpr env exprs
      in (env', genSexpr (toHL cn) es False)
genExpr env (DFieldAcc _ fns _) 
    = let es = foldr (\(DField fn _) r -> genSexpr fn [r] False) "|vx|" fns
      in (env, es) 
      -- it is assumed that Curr is not used

genExpr env (DFieldAccE (DEdge _) [] _) = (env, vE)
genExpr (m : env) (DAggr agop e1 ge e2 _) = (env, m)
-- genExpr env (DVExp (DVar vn _) _) = (env, vn)
genExpr env (DVExp (DVar vn _) _) = (env, genZ3var vn)
genExpr env (DCExp c _) = (env, genConst c)
genExpr _ d = traceShow d undefined

genConst :: (Show a) => DConst (Ty a) -> String
genConst (DCInt i _) = "(_rv " ++ show i ++ ")"
genConst (DCBool True _) = "true"
genConst (DCBool False _) = "false"
genConst (DCDouble d _) = show d

genFunOp :: (Show a) => DFun (Ty a) -> String
genFunOp (DFun fn t) = genOF fn (argTypeName1 . fst $ t)
genFunOp (DBinOp opn t) = genOF opn (argTypeName1 . fst $ t)

genOF :: String -> String -> String
genOF "==" "Bool" = "="
genOF "==" _ = "_eq"
genOF "+"  _ = "_add"
genOF "*"  _ = "_mult"
genOF "/"  _ = "_div"
genOF "neg"_ = "_neg"
genOF "min"_ = "_min"
genOF "max"_ = "_max"
genOF "||" _ = "or"
genOF "&&" _ = "and"
genOF ">"  _ = "_gt"
genOF "<"  _ = "_lt"
genOF a    _ = a

genEqOp :: String -> String -> String
genEqOp "!=" "Int" = "_neq"
genEqOp "!=" "Float" = "_neq"
genEqOp "!=" "Bool" = "xor"

-- extracts aggregators concering message
messageAggrs :: (Show a) =>
  String -> [DSmplDef (Ty a)] -> DExpr (Ty a) ->
  [(DAgg (Ty a), DExpr (Ty a), DGen (Ty a), [DExpr (Ty a)])]
messageAggrs dn ldefs expr = msag expr ++ concatMap msagLet ldefs
  where msag (DIf e1 e2 e3 _)  = concatMap msag [e1, e2, e3]
        msag (DFunAp _ exprs _) = concatMap msag exprs
        msag (DConsAp (DConstructor cn _) exprs _)  = concatMap msag exprs
        msag (DAggr agop e1 ge e2 _) = [(agop, e1, ge, e2)]
        msag a = []
        msagLet (DDefVar (DVar vn _) defs expr _)
          = msag expr ++ concatMap msagLet defs 

genAggr :: DAgg (Ty a) -> String
genAggr (DAggMin _) = "_min"
genAggr (DAggMax _) = "_max"
genAggr (DAggSum _) = "_add"
genAggr (DAggProd _) = "_mult"
genAggr (DAggAnd _) = "and"
genAggr (DAggOr _) = "or"

genCommElimSpec :: (Show a) =>
  [DGroundDef (Ty a)] -> PregelStepfn (Ty a) -> String
genCommElimSpec gdefs ps@(fname, ldefs, expr, term, recspec) =
  prelude ++ unlines (genAuxDecls recspec gdefs) ++ mdfs ++ 
  vc ++ vdefs ++ cnd ++ genZ3CheckSat True
  where DRecordSpec (DConstructor dn _) dflds _ = recspec
        mdfs = unlines $ zipWith3 (genAop dn) mds mdts [1..]
        (fn, mds, mdts, vc, vars) = genVC dn (fname, ldefs, expr)
        vars2 = map (\(a:vn) -> a : '2' : vn) vars
        dcfld v t = genDC v t ++ genFldcon v dflds ++ "\n"
        vdefs = dcfld v1 dn ++ dcfld v2 dn ++
                genDC e tRP ++ genDC w dn ++ "\n" ++
                unwords (zipWith genDC vars mdts) ++ "\n" ++
                unwords (zipWith genDC vars2 mdts) ++ "\n"
        aopApp (v,n) =
          genSexpr (genZ3var ("aop" ++ show n)) [v, w, e] False
        vcApp v f vs = genSexpr (genZ3var f) (v:vs) False
        cnd = genZ3AssertEq v2 (vcApp v1 fn (map aopApp (zip vars [1..]))) True
              ++
              genZ3AssertNotEq
                (vcApp v2 fn vars2)
                (vcApp v2 fn (map aopApp (zip vars2 [1..])))
                True
        v1 = genZ3var "v1"
        v2 = genZ3var "v2"
        e = genZ3var "e"
        w = genZ3var "w"

{-
genInactivateSpec pr@(DProgramSpec rs p _) = 
    let (DRecordSpec (DConstructor dn _) dflds _) : _ = rs
        [(fn, mds, vc, vars)] = map (genVC dn) $ getVC pr
        vcApp v f vs = "(|" ++ f ++ "| " ++ v ++ " " ++ unwords vs ++")"
        fldcon v = concat ["(assert (_isReal (" ++ fn ++ " " ++ v ++ ")))" | 
                           (DField fn _, DTInt _) <- dflds] 
        dc t v = let ass = if t == tRP 
                           then "(assert (_isReal " ++ v ++ ")) "
                           else ""
                 in "(declare-const "  ++ v ++ " " ++ t ++ ") " ++ ass
        units = map (\(aop,_,_,_) -> toUnit aop) mds
        allOf = foldr (\a r -> "(and " ++ a ++ " " ++ r ++")") "true"
        genFilter (_,_,_,es) = map (snd . genExpr []) es 
        cnd1 = "(assert " ++ allOf (concatMap genFilter mds) ++ ")\n"
        cnd2 = "(assert (not (= |vx| " ++ vcApp "|vx|" fn units ++ ")))\n"
    in prelude ++ unlines (genAuxDecls pr) ++ vc ++ 
       dc dn "|vx|" ++ fldcon "|vx|" ++ "\n" ++ 
       cnd1 ++ cnd2  ++ "(check-sat)\n"
-}

genInactivateSpec :: (Show a) =>
  [DGroundDef (Ty a)] -> PregelStepfn (Ty a) -> String
genInactivateSpec gdefs (fname, ldefs, expr, term, recspec) =
  prelude ++ unlines (genAuxDecls recspec gdefs) ++ vc ++ dcfld vx dn ++
  "\n" ++ vdefs ++ cnd1 ++ cnd2 ++ genZ3CheckSat True
  where DRecordSpec (DConstructor dn _) dflds _ = recspec
        (fn, mds, mdts, vc, vars) = genVC dn (fname, ldefs, expr)
        dcfld v t = genDC v t ++ genFldcon v dflds ++ "\n"
        vdefs = genDC e tRP ++ "\n"
        vcApp v f vs = genSexpr (genZ3var f) (v:vs) False
        units = map (\(aop,_,_,_) -> toUnit aop) mds
        allOf = foldr (\a r -> genSexpr "and" [a,r] False) "true"
        genFilter (_,_,_,es) = map (snd . genExpr []) es 
        cnd1 = genSexpr "assert" [allOf (concatMap genFilter mds)] True
        cnd2 = genZ3AssertNotEq vx (vcApp vx fn units) True
        e = genZ3var "e"
        vx = genZ3var "vx"

toUnit :: DAgg (Ty a) -> String
toUnit (DAggMin _) = "_pinf"
toUnit (DAggMax _) = "_ninf"
toUnit (DAggSum _) = "(_rv 0)"
toUnit (DAggProd _) = "(_rv 1)"
toUnit (DAggAnd _) = "true"
toUnit (DAggOr _) = "false"

genParenedStr :: [String] -> String
genParenedStr ss = "(" ++ unwords ss ++ ")"

genSexpr :: String -> [String] -> Bool -> String
genSexpr fn args nl = genParenedStr (fn : args) ++ if nl then "\n" else ""

genAop :: (Show a) =>
  String -> (DAgg (Ty a), DExpr (Ty a), DGen (Ty a), [DExpr (Ty a)]) ->
  String -> Int -> String
genAop dn p@(aop, e1, _ , _) t n =
  genZ3DefineFun
    ("aop" ++ show n)
    (genSexpr (genZ3VarTy mx t) [genZ3VarTy vx dn, genZ3VarTy vE tRP] False)
    t     -- Is this tekitou?
    (genSexpr (genAggr aop) [snd (genExpr [] e1), mx] False)
    True
  where mx = genZ3var "mx"
        vx = genZ3var "vx"

genFldcon :: (Show a) => String -> [(DField (Ty a), DType (Ty a))] -> String
genFldcon v dflds =
   -- concat [genZ3AssertIsReal ("(" ++ fn ++ " " ++ v ++ ")") False
   concat [genZ3AssertIsReal (genSexpr fn [v] False) False
             | (DField fn _, DTInt _) <- dflds] 

genDC :: String -> String -> String
genDC v t = let ass = if t == tRP
                      then genZ3AssertIsReal v False ++ " "
                      else ""
            in genZ3DeclareConst v t False ++ " " ++ ass

genZ3DefineFun :: String -> String -> String -> String -> Bool -> String
genZ3DefineFun fname formals rty body nl =
  genSexpr "define-fun" [genZ3var fname, formals, rty, body] nl

genZ3DeclareDatatypes :: String -> Bool -> String
genZ3DeclareDatatypes dt nl = genSexpr "declare-datatypes" ["()", dt] nl
         
genZ3DeclareConst :: String -> String -> Bool -> String
genZ3DeclareConst cname ty nl = genSexpr "declare-const" [cname, ty] nl

genZ3AssertEq :: String -> String -> Bool -> String
genZ3AssertEq op1 op2 nl =
  genSexpr "assert" [genSexpr "=" [op1, op2] False] nl

genZ3AssertNotEq :: String -> String -> Bool -> String
genZ3AssertNotEq op1 op2 nl =
  genSexpr "assert" [genSexpr "not" [genSexpr "=" [op1, op2] False] False] nl

genZ3AssertIsReal :: String -> Bool -> String
genZ3AssertIsReal op nl =
  genSexpr "assert" [genSexpr "_isReal" [op] False] nl

genZ3VarTy :: String -> String -> String
genZ3VarTy v t = genParenedStr [v,t]

genZ3var :: String -> String
genZ3var v = "|" ++ v ++ "|"

genZ3CheckSat :: Bool -> String
genZ3CheckSat nl = "(check-sat)" ++ if nl then "\n" else ""

nthArgType :: Int -> DTypeTerm -> DTypeTerm
nthArgType n (DTypeTerm "->" [t1, t2]) | n == 1 = t1
                                       | otherwise = nthArgType (n - 1) t2
nthArgType _ _ = undefined

argType1, argType2 :: DTypeTerm -> DTypeTerm
argType1 = nthArgType 1
argType2 = nthArgType 2

resultType :: DTypeTerm -> DTypeTerm
resultType (DTypeTerm "->" [_, t2]) = resultType t2
resultType t = t

dTypeName :: DTypeTerm -> String
dTypeName (DTypeTerm name _) = name
dTypeName (DTypeVar id) = id

argTypeName1, argTypeName2, resultTypeName :: DTypeTerm -> String
argTypeName1 = dTypeName . argType1
argTypeName2 = dTypeName . argType2
resultTypeName = dTypeName . resultType