import Control.Monad.State.Strict
import Control.Monad.Identity
import MiniWhile
import Data.Map 



------Pseudo random number generator------

type Random a = State Integer a
 
 

-- data ST s a = [s] -> [(a, [s])]

--------state is the iteration
fresh :: Random Integer
fresh = StateT $ \i -> Identity ((g i), i + 1)
        where g i = if i >=1 then (6364136223846793005 * (evalState fresh (i - 1)) + 1442695040888963407) `mod` (2^64) 
                     else (1442695040888963407 `mod` (2^64))

--fresh :: Random Integer
--fresh = do 
--           i <- get
--           if i > 1 then
--           put (i-1)


runPRNG :: Random a -> Integer -> a
runPRNG fresh i = evalState fresh i

--------- state stores x_n for the next round
fresh' :: Random Integer
fresh' = do 
           x <- get
           let nextX = (6364136223846793005 * x + 1442695040888963407) * (2 ^ 64)
           put nextX 
    --       put :: s -> m ()
           return nextX

runPRNG' :: Random a -> Integer -> a
runPRNG' fresh' i = evalState fresh' i
--------While Evaluation

type Value = Integer

type Memory = Map Id Value

-- Map k :: functor

type InterpM = State Memory

--State s a :: StateT s Identity a = StateT $ \s -> Identity (a , s)
--data Stmt = Asgn Aexp Exp | While Exp [Stmt]

evalStmts :: [Stmt] -> InterpM ()
evalStmts [] = return ()
evalStmts (s:sts) = evalStatement s >> evalStmts sts

evalStatement :: Stmt -> InterpM ()
evalStatement (Asgn (Var id) ex) =  do 
                                       v <- evalExp ex
                                       memory <- get
                                       put (insert id v memory)
evalStatement (Asgn _ ex) = return ()

evalStatement w@(While exp sts) = do 
                                   v <- evalExp exp
                                   evalStmts sts
                                   evalStatement (While exp sts)

evalExp :: Exp -> InterpM Integer
evalExp (IfTE e1 e2 e3) = do 
                              v1 <- evalExp e1 
                              v2 <- evalExp e2
                              v3 <- evalExp e3
                              if v1 == 1 then return v2 else return v3
evalExp (ComExp aex1 cop aex2) = do 
                                    v1 <- evalAexp aex1
                                    v2 <- evalAexp aex2
                                    return (comop cop v1 v2)
evalExp (Not ex) = do 
                        v <- evalExp ex 
                        if v == 1 then return 0 else return 1
evalExp (AExp aex) = evalAexp aex


evalAexp :: Aexp -> InterpM Integer
evalAexp (Num i)  = return i
evalAexp (Var id) = findWithDefault 0 id <$> get
      --              findWithDefault 0 id :: (Map Id Int -> Int) -> InterpM (Map Id Int) -> InterpM Integer
      --              fmap :: a -> b -> m a -> m b
evalAexp (OpA aex1 op aex2) = do 
                                 v1 <- evalAexp aex1
                                 v2 <- evalAexp aex2
                                 return (aop op v1 v2)

aop :: Op -> Integer -> Integer -> Integer
aop Plus x y = x + y
aop Minus x y = x - y
aop Multi x y = x * y
aop Div x y = x `div` y

comop :: ComOp -> Integer -> Integer -> Integer
comop LE x y | x <= y    = 1
             | otherwise = 0
comop Gr x y | x > y     = 1
             | otherwise =0
comop Equal x y | x == y    = 1
                | otherwise = 0
comop UnE x y | x == y    = 0
              | otherwise = 1


eval :: Program -> Memory
eval (Program stms) = execState (evalStmts stms) empty