{-# OPTIONS -fglasgow-exts #-}

module LISS.Analyzer where

import Data.Maybe
import qualified Data.Map as Map
import qualified Data.List as List

import Control.Monad
import Control.Monad.Error
import Control.Monad.State
import Control.Monad.Identity

import LISS.Syntax
import qualified Why.SyntaxFlattened as WhyF

freshvars = ["a"++n | n <- map show [1..10000]]

type Size = Int
data Env = Env { vars :: [String], arrays :: Map.Map String Size, fdecls :: [WhyF.WhyF], label :: String }

remove :: Monad m => String -> StateT Env m ()
remove v = modify (\(Env fresh sizes fd l) -> Env (List.delete v fresh) sizes fd l)

setLabel :: Monad m => String -> StateT Env m ()
setLabel v = modify (\(Env fresh sizes fd l) -> Env fresh sizes fd v)

addArray :: Monad m => String -> Size -> StateT Env m ()
addArray v sz = modify (\(Env fresh sizes fd l) -> Env fresh (Map.insert v sz sizes) fd l)

addFDecl :: Monad m => WhyF.WhyF -> StateT Env m ()
addFDecl f = modify (\(Env fresh sizes fd l) -> Env fresh sizes (f:fd) l)

fresh :: Monad m => StateT Env m String
fresh =
     do name <- gets (head . vars)
     	modify (\(Env fresh sizes fd l) -> Env (tail fresh) sizes fd l)
	return name

class Logic a where
  analyze :: (Error e, MonadError e m) => a -> StateT Env m WhyF.WhyF

-- instances

analyzeLiss :: Liss -> IO WhyF.Why
analyzeLiss (Liss name body) =
     do (whyf,env) <- runStateT (analyzeBody body (WhyF.PIdent "nop")) (Env freshvars Map.empty [] "")
     	let arrs = Map.keys $ arrays env
     	let whyarrays = if null arrs then [] else [WhyF.Parameter False arrs (WhyF.Simple (WhyF.SVTArray WhyF.TInt))]
     	let res = [WhyF.Let name Nothing whyf]
	return (whyarrays++res)

analyzeBody :: (Error e, MonadError e m) => Body -> WhyF.WhyF -> StateT Env m WhyF.WhyF
analyzeBody (Body decs st) whyr =
     do whyst <- analyze st
     	foldM (flip $ curry analyze) (WhyF.Block (whyst `WhyF.Seq` whyr)) $ reverse decs

-- paired with WhyF for implementation purposes
instance Logic (Declaration,WhyF.WhyF) where
	analyze ((VarDec (Arg name Integer)),prog) = return $ WhyF.LetInRef name (WhyF.PConst (WhyF.IntConst 0)) prog
	analyze ((VarDec (Arg name Boolean)),prog) = return $ WhyF.LetInRef name (WhyF.PConst (WhyF.BoolConst False)) prog
	analyze ((VarDec (Arg name Sequence)),prog) = return $ WhyF.LetInRef name (WhyF.PIdent "seqnil") prog
	analyze ((VarDec (Arg name Set)),prog) = return $ WhyF.LetInRef name (WhyF.PIdent "setnil") prog
	analyze ((VarDec (Arg name (Array size))),prog) = addArray name size >> return prog
	analyze ((SubprogDec specs name args Integer body (ReturnInt ie1)),prog) = analyzeFun specs name args body ie1 prog
	analyze ((SubprogDec specs name args Boolean body (ReturnBool be1)),prog) = analyzeFun specs name args body be1 prog
	analyze ((SubprogDec specs name args Sequence body (ReturnSeq se1)),prog) = analyzeFun specs name args body se1 prog
	analyze ((SubprogDec specs name args Set body (ReturnSet se1)),prog) = analyzeFun specs name args body se1 prog
	analyze ((SubprogDec specs name args (Array sz) body (ReturnArray ae1)),prog) = analyzeFun specs name args body ae1 prog
	analyze _ = throwError (strMsg "Inconsistent function declaration")

analyzeFun specs name args body re prog =
     do whyargs <- mapM analyzeArg args
	re' <- analyze re
	whybody <- analyzeBody body re'
	pres <- mapM analyze $ filter isPre specs
	posts <- mapM analyze $ filter isPost specs
	let whybodypre = if null pres then whybody else WhyF.PPre (foldr1 WhyF.Conj pres, Nothing) whybody
	let whybodypos = if null posts then whybodypre else WhyF.PPost1 whybodypre ((foldr1 WhyF.Conj posts,Nothing),[])
	return $ WhyF.LetBind name whyargs whybodypos prog


analyzeArg :: (Error e, MonadError e m) => Arg -> StateT Env m (WhyF.Identifier,WhyF.WhyF)
analyzeArg (Arg name Integer) = return $ (name,WhyF.Simple (WhyF.SVTRef WhyF.TInt))
analyzeArg (Arg name Boolean) = return $ (name,WhyF.Simple (WhyF.SVTRef WhyF.TBool))
analyzeArg (Arg name Sequence) = return $ (name,WhyF.Simple (WhyF.SVTRef (WhyF.User "sequence")))
analyzeArg (Arg name Set) = return $ (name,WhyF.Simple (WhyF.SVTRef (WhyF.User "set")))
analyzeArg (Arg name (Array sz)) = addArray name sz >> return (name,WhyF.Simple (WhyF.SVTArray WhyF.TInt))

instance Logic Statement where
	analyze (IntAssign ivar ie1) = analyze ie1 >>= return . WhyF.Assign ivar
	analyze (BoolAssign bvar be1) = analyze be1 >>= return . WhyF.Assign bvar
	analyze (SeqAssign svar se1) = analyze se1 >>= return . WhyF.Assign svar
	analyze (SetAssign svar se1) = analyze se1 >>= return . WhyF.Assign svar
	analyze (ArrayPosAssign avar ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.ArrayAssign avar iv1 iv2
	analyze (SequenceS c1 c2) =
	     do cv1 <- analyze c1
	     	cv2 <- analyze c2
		return $ WhyF.Seq cv1 cv2
	analyze (IfThenElse be1 c1 c2) =
	     do bv1 <- analyze be1
	     	cv1 <- analyze c1
		cv2 <- analyze c2
		return $ WhyF.ProgIfThenElse bv1 (WhyF.Block cv1) (Just (WhyF.Block cv2))
	analyze (While la be1 c1) =
	     do	bv1 <- analyze be1
	     	cv1 <- analyze c1
		oldl <- gets label
		l <- fresh
		setLabel l
                (a,b) <- analyzeLA la
		setLabel oldl
		return $ WhyF.Label l $ WhyF.While bv1 ((a, Nothing), b) cv1 -- correct?
	analyze (For ivar ll ul step be1 c1) =
	     analyze (
		(IntAssign ivar ll) `SequenceS`
		(While (PBoolConst True, Nothing) -- body could change ivar, yelding this variant probably incorrect
                        (LessEq (IntVarExp ivar) ul)
			((IfThenElse be1 c1 Skip) `SequenceS`
			(IntAssign ivar (Sum (IntVarExp ivar) (IntConst step))))))
	analyze (ForEach ivar se1 be1 c1) =
	     do nv <- fresh
		sv1 <- analyze se1
		cv1 <- analyze (
			(While (PBoolConst True, Nothing) -- probably wrong
                                (Negation (EmptySeq (SeqVarExp nv)))
				((IntAssign ivar (Head (SeqVarExp nv))) `SequenceS`
				((IfThenElse be1 c1 Skip) `SequenceS`
				(SeqAssign nv (Tail (SeqVarExp nv)))))))
		return $ WhyF.LetInRef nv sv1 cv1
	analyze Skip = return $ WhyF.PIdent "nop" -- correct
	analyze (Read ivar) = return $ WhyF.TFunction "read" [WhyF.TVar (ivar,Nothing)]
	analyze (Write ie1) =
	     do iv1 <- analyze ie1
	     	return $ WhyF.TFunction "write" [iv1]
	analyze (WriteString s) =
	     do return $ WhyF.PIdent "nop" -- correct?

analyzeLA :: (Error e, MonadError e m) => LoopAnnot -> StateT Env m (WhyF.WhyF, Maybe (WhyF.WhyF, Maybe Identifier))
analyzeLA (a, Just b) = do 
	a' <- analyze a
	b' <- analyze b
	return (a', Just (b',Nothing))
analyzeLA (a, Nothing) = do
	a' <- analyze a
	return (a', Nothing)

instance Logic IntExp where
	analyze (IntConst i)  = return $ WhyF.IntConst i
	analyze (IntVarExp v) = return $ WhyF.RefAccess v
	analyze (ArrayAccess v e1) =
	     do v1 <- analyze e1
	    	return $ WhyF.ArrayAccess v v1
	analyze (Negative e1) = analyze e1 >>= return . WhyF.TNegation
	analyze (Sum e1 e2)   =
	     do v1 <- analyze e1
	     	v2 <- analyze e2
		return $ WhyF.Sum v1 v2
	analyze (Sub e1 e2)   =
	     do v1 <- analyze e1
	     	v2 <- analyze e2
		return $ WhyF.Sub v1 v2
	analyze (Mult e1 e2)  =
	     do v1 <- analyze e1
	     	v2 <- analyze e2
		return $ WhyF.Mult v1 v2
	analyze (Div e1 e2)   =
	     do v1 <- analyze e1
	     	v2 <- analyze e2
		return $ WhyF.Div v1 v2
	analyze (Rem e1 e2)   =
	     do v1 <- analyze e1
	     	v2 <- analyze e2
		return $ WhyF.Rem v1 v2
	analyze (Head seq) =
	     do seqv <- analyze seq
	     	return $ WhyF.TFunction "why_head" [seqv]
	analyze (IntFunction name vars) = return $ WhyF.TFunction name (map (\v -> WhyF.TVar (v,Nothing)) vars)

instance Logic BoolExp where
	analyze (BoolConst True)  = return $ WhyF.PredTrue
	analyze (BoolConst False) = return $ WhyF.PredFalse
	analyze (BoolVarExp v)  = return $ WhyF.RefAccess v
	analyze (Equal ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpEq iv2
	analyze (Diff ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpNEq iv2
	analyze (Lesser ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpL iv2
	analyze (LessEq ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpLE iv2
	analyze (Greater ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpG iv2
	analyze (GreatEq ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Infix iv1 WhyF.OpGE iv2
	analyze (Negation be1) = analyze be1 >>= return . WhyF.PNegation
	analyze (Conj  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Infix bv1 WhyF.OpAnd bv2
	analyze (Disj  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Infix bv1 WhyF.OpOr bv2
	analyze (Imp  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Imp bv1 bv2
	analyze (Equiv  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Equiv bv1 bv2
	analyze (EmptySeq se1) = analyze se1 >>= return . WhyF.PFunction "why_is_seqnil" . (:[])
	analyze (ElemSet ie1 set) =
	     do iv1 <- analyze ie1
	     	setv <- analyze set
		return $ WhyF.PFunction "inset" [iv1, setv]
	analyze (BoolFunction name vars) = return $ WhyF.PFunction name (map WhyF.PVar vars)

instance Logic SeqExp where
	analyze (SeqVarExp v) = return $ WhyF.RefAccess v
--	analyze (SeqConst l) =
	analyze SeqNull = return $ WhyF.PVar "seqnil"
	analyze (Tail seq) =
	     do seqv <- analyze seq
	     	return $ WhyF.TFunction "why_tail" [seqv]
	analyze (Cons ie1 seq) =
	     do iv1 <- analyze ie1
	     	seqv <- analyze seq
		return $ WhyF.TFunction "cons" [iv1, seqv]
	analyze (SeqFunction name vars) = return $ WhyF.TFunction name (map (\v -> WhyF.TVar (v,Nothing)) vars)

instance Logic SetExp where
	analyze (SetVarExp v) = return $ WhyF.RefAccess v
--	analyze (SetConst x p) =
--	analyze SetNull =
	analyze (Union se1 se2) =
	     do sv1 <- analyze se1
	     	sv2 <- analyze se2
		return $ WhyF.TFunction "union" [sv1, sv2]
	analyze (Intersect se1 se2) =
	     do sv1 <- analyze se1
	     	sv2 <- analyze se2
		return $ WhyF.TFunction "intersect" [sv1, sv2]
	analyze (SetFunction name vars) = return $ WhyF.TFunction name (map (\v -> WhyF.TVar (v,Nothing)) vars)

instance Logic ArrayExp where
--	analyze (ArrayVarExp v) = return $ WhyF.RefAccess v
	analyze (ArrayFunction name vars) = return $ WhyF.TFunction name (map (\v -> WhyF.TVar (v,Nothing)) vars)

instance Logic Spec where
	analyze (Pre p) = analyze p
	analyze (Post p) = analyze p

instance Logic Term where 
	analyze (TIntConst i) = return $ WhyF.Const $ WhyF.IntConst i
	analyze (TIntVar iv) = return $ WhyF.TVar (iv,Nothing)
	analyze (TArrayAccess v t) =
	     do t' <- analyze t
	     	return $ WhyF.TArrayAccess (v,Nothing) t'
	analyze (TNegative t) =
	     do t' <- analyze t
		return $ WhyF.TNegation t'
	analyze (TSum t1 t2) =
	     do t1' <- analyze t1
		t2' <- analyze t2
		return $ WhyF.Sum t1' t2'
	analyze (TSub t1 t2) =
	     do t1' <- analyze t1
		t2' <- analyze t2
		return $ WhyF.Sub t1' t2'
	analyze (TMult t1 t2) =
	     do t1' <- analyze t1
		t2' <- analyze t2
		return $ WhyF.Mult t1' t2'
	analyze (TDiv t1 t2) =
	     do t1' <- analyze t1
		t2' <- analyze t2
		return $ WhyF.Div t1' t2'
	analyze (TRem t1 t2) =
	     do t1' <- analyze t1
		t2' <- analyze t2
		return $ WhyF.Rem t1' t2'
	analyze (THead seq) =
	     do seqv <- analyze seq
	     	return $ WhyF.TFunction "why_head" [seqv]
	analyze (TIntFun name vars) = return $ WhyF.TFunction name (map (\v -> WhyF.TVar (v,Nothing)) vars)
	analyze (TOld v) = gets label >>= return . (\l -> WhyF.TVar (v,Just l))
--	analyze (TOldPos v t) =
--	     do t' <- analyze t
--	     	return $ WhyF.TOld $ WhyF.ArrayAccess v t'
	analyze Result = return WhyF.TResult

instance Logic Predicate where
	analyze (PBoolConst True)  = return $ WhyF.PredTrue
	analyze (PBoolConst False) = return $ WhyF.PredFalse
	analyze (PBoolVar v)  = return $ WhyF.RefAccess v
	analyze (PEqual ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Equal iv1 iv2
	analyze (PDiff ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Diff iv1 iv2
	analyze (PLesser ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Lesser iv1 iv2
	analyze (PLessEq ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.LessEq iv1 iv2
	analyze (PGreater ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.Greater iv1 iv2
	analyze (PGreatEq ie1 ie2) =
	     do iv1 <- analyze ie1
	     	iv2 <- analyze ie2
		return $ WhyF.GreatEq iv1 iv2
	analyze (PNegation be1) = analyze be1 >>= return . WhyF.PNegation
	analyze (PConj  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Conj bv1 bv2
	analyze (PDisj  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Disj bv1 bv2
	analyze (PImp  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Imp bv1 bv2
	analyze (PEquiv  be1 be2) =
	     do bv1 <- analyze be1
	     	bv2 <- analyze be2
		return $ WhyF.Equiv bv1 bv2
	analyze (PEmptySeq se1) = analyze se1 >>= return . WhyF.Diff (WhyF.PVar "seqnil")
	analyze (PElemSet ie1 set) =
	     do iv1 <- analyze ie1
	     	setv <- analyze set
		return $ WhyF.PFunction "inset" [iv1, setv]
	analyze (PBoolFunction name vars) = return $ WhyF.PFunction name (map WhyF.PVar vars)
	analyze (PIfThenElse t1 p1 p2) =
	     do tv1 <- analyze t1
	     	pv1 <- analyze p1
	     	pv2 <- analyze p2
		return $ WhyF.PIfThenElse tv1 pv1 pv2
	analyze (Forall args p1) =
	     do let (ints,bools,seqs,sets) = separateTypes args
	     	pv1 <- analyze p1
		let pv2 = if null sets then pv1 else WhyF.Forall sets (WhyF.User "set") [] pv1
		let pv3 = if null seqs then pv2 else WhyF.Forall seqs (WhyF.User "sequence") [] pv2
		let pv4 = if null bools then pv3 else WhyF.Forall bools WhyF.TBool [] pv3
		let pv5 = if null ints then pv4 else WhyF.Forall ints WhyF.TInt [] pv4
		return $ pv5
	analyze (Exists args p1) =
	     do let (ints,bools,seqs,sets) = separateTypes args
	     	pv1 <- analyze p1
		let pv2 = if null sets then pv1 else WhyF.Exists sets (WhyF.User "set") pv1
		let pv3 = if null seqs then pv2 else WhyF.Exists seqs (WhyF.User "sequence") pv2
		let pv4 = if null bools then pv3 else WhyF.Exists bools WhyF.TBool pv3
		let pv5 = if null ints then pv4 else WhyF.Exists ints WhyF.TInt pv4
		return $ pv5
	analyze (ValidIndex v ie1) =
	     do iv1 <- analyze ie1
	     	sz' <- gets (Map.lookup v . arrays)
		case sz' of
			Nothing -> throwError (strMsg ("Unknown array: "++v))
			Just sz -> return $ WhyF.GreatEq iv1 (WhyF.Const (WhyF.IntConst 0)) `WhyF.Conj` WhyF.Lesser iv1 (WhyF.Const (WhyF.IntConst sz))

separateTypes args =
    let ints = map (\(Arg v t) -> v) $ filter (\(Arg v t) -> t == Integer) args
    	bools = map (\(Arg v t) -> v) $ filter (\(Arg v t) -> t == Boolean) args
    	seqs = map (\(Arg v t) -> v) $ filter (\(Arg v t) -> t == Sequence) args
    	sets = map (\(Arg v t) -> v) $ filter (\(Arg v t) -> t == Set) args
     in (ints,bools,seqs,sets)

isPre (Pre p) = True
isPre _ = False

isPost (Post p) = True
isPost _ = False

