{-# OPTIONS_GHC -fglasgow-exts -fallow-overlapping-instances #-}

module HW5 where

import qualified Data.List as L
import qualified Data.Generics as G
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.State
import Control.Monad.Error


----------------------------------------------------------------------------------
--
-- Terms
--
data Term = Var String       
          | String :\ Term   
          | Term :$: Term

          | Con Integer
          | Term :+: Term
          | Term :-: Term

          | Count
          | Out Term

          | Let [Decl] Term
          | LetRec [Decl] Term

          | Term :*: Term
          | Fst Term
          | Snd Term

          | TT
          | FF
          | Not Term
          | Term :==: Term
          | If Term Term Term

          | Skip
          | Term :& Term
          | While Term Term

          | Term ::: Term
          | Nil 
          | LCase Term Term Term

          | MkRef Term
          | Read Term
          | ReadV String
          | Term := Term
          | String :!= Term
  deriving (G.Data, G.Typeable, Show)

data Decl = String :=: Term | String :!=: Term 
  deriving (G.Data, G.Typeable, Show)

infix 1 :=:
infix 1 :!=:
infixr 2 :\
infixr 3 :&
infix 4 :!=
infix 4 :=
infixl 5 :$:
infixl 7 :+:
infixr 8 :::


----------------------------------------------------------------------------------
--
-- Finding free variables of a term
--

evT :: G.GenericQ [String] -> Term -> [String]
evT f e = 
    case e of
        (x :\ e') -> go e' L.\\ [x]
        (Let decl e') -> go decl ++ (go e' L.\\ map declV decl)
        (LetRec decl e') -> declVs ++ (go e' L.\\ map declV decl) where
            declVs = concat (map g decl)
            g (x :=: e) = go e L.\\ [x]
            g (x :!=: e) = go e L.\\ [x]
        _ -> f e ++ (concat $ G.gmapQ go e)
  where
    declV x = case x of v :=: _ -> v; v :!=: _ -> v
    go x = L.nub $ G.mkQ [] (evT f) x

freeVars e = L.nub $ evT (G.mkQ [] goTerm) e where
    goTerm (Var x) = [x]
    goTerm (ReadV x) = [x]
    goTerm (x :!= _) = [x]
    goTerm _ = []


----------------------------------------------------------------------------------
--
-- Values
--
data Value = Wrong
             | Num Integer
             | Fun (Value -> M Value)
             | Unit
             | List [Value]
             | Bl Bool
             | Ref Int
             | Pair (Value, Value)
instance Eq Value where
    Num x == Num y = x == y
    Unit == Unit = True
    List x == List y = x == y
    Bl x == Bl y = x == y
    Ref x == Ref y = x == y
    Pair x == Pair y = x == y
    _ == _ = False
instance Show Value where
    show Wrong = "<wrong>"
    show (Num  n) = show n
    show (Fun _) = "<function>"
    show Unit = "()"
    show (List x) = show x
    show (Bl x) = show x 
    show (Ref x) = "<ref "++(show x)++">"
    show (Pair x) = show x


----------------------------------------------------------------------------------
--
-- Types for interpreter
--
data MyState = MyState {step::Integer, refs::[Value]} deriving Show
initSt = MyState{step = 0, refs = []}

type Env = [(String, Value)]

instance Show (Either String Value) where
    show (Right v) = "Success: "++show v
    show (Left err) = "Error: "++err
instance Show ((Either String Value, String), MyState) where
    show ((v,o),s) = ("("++(show v)++" | Output: "++o++" | "++(show s)++")\n")
   
--
-- Function to run an interp computation 
--
test :: Term -> ((Either String Value, String), MyState)
test e = run initSt [] $ interp e

-- !!!!!!!!!!
--
--
-- Supply definition of M so that test has given type. 
--
type M a = 

--
-- Supply definition of run
--
run :: MyState -> Env -> M Value -> ((Either String Value, String), MyState)
run st env x = 
--
--
-- !!!!!!!!!!


----------------------------------------------------------------------------------
--
-- Helper functions for interpreter
--
tick :: M ()
tick = modify $ \s -> s{step = (step s) + 1}

pair proj (Pair v) = do tick; return$  proj v
pair _ v = throwError $ "expecting pair, found "++(show v)

not' (Bl v) = return $ Bl $ not v
not' x =  throwError $ "expected bool, found "++(show x)

if' (Bl v) x y = do tick; if v then x else y
if' v _ _ = throwError $ "expected bool, found "++(show v)

cons x (List y) = do tick; return $ List $ x:y
cons x y = throwError $ "expected list, found "++(show y)

lcase (List []) x y = do tick; x
lcase (List (h:t)) x y = do
    tick
    v1 <- y 
    case v1 of 
        Fun f1 -> do
            v3 <- f1 h
            case v3 of
                Fun f2 -> f2 (List t)
                f2 -> throwError $ "expecting function, found "++(show f2)
        f1 -> throwError $ "expecting function, found "++(show f1)
lcase x _ _ = throwError $ "expecting list, found "++(show x)

getRef :: Value -> M Value
getRef (Ref x) = gets $ \s -> (refs s)!!x
getRef x = throwError $ "expecting ref, found "++(show x)

setRef :: Value -> Value -> M Value
setRef (Ref x) v = do modify (\s -> s{refs = go x (refs s)}); return Unit where
    go 0 (_:vs) = v:vs
    go n (v:vs) = v : go (n-1) vs
setRef x _ = throwError $ "expecting ref, found "++(show x)

newRef :: Value -> M Value
newRef v = do
    s <- get
    let rfs = refs s
    put $ s{refs = rfs++[v]}
    return $ Ref $ length rfs

getVar :: Env -> String -> M Value
getVar ((x,v):xs) y = if x == y then return v else getVar xs y
getVar [] y = throwError $ "unbound variable "++y

arith :: (Integer -> Integer -> Integer) -> Value -> Value -> M Value
arith f (Num m) (Num n) = do tick; return $ Num $ f m n
arith f m n = throwError $ "expected numbers, found "++(show m)++" and "++(show n)

apply :: Value -> Value -> M Value
apply (Fun v1) v2 = do tick; v1 v2
apply f _ = throwError $ "expected function, found "++(show f)


----------------------------------------------------------------------------------
--
-- Interpreter
--
interp :: Term -> M Value

interp (Var x) = do
    env <- ask
    getVar env x
interp (x :\ e) = do
    env <- ask
    return $ Fun $ \v -> local (\_ -> (x,v):env) $ interp e
interp (e1 :$: e2) = do
    v1 <- interp e1 
    v2 <- interp e2 
    apply v1 v2

interp (Con n) = return $ Num n
interp (e1 :+: e2) = do
    v1 <- interp e1 
    v2 <- interp e2 
    arith (+) v1 v2
interp (e1 :-: e2) = do
    v1 <- interp e1 
    v2 <- interp e2 
    arith (-) v1 v2
interp Count = do 
    n <- gets step 
    return $ Num n

interp (Out e) = do 
    v <- interp e  
    tell $ (show v)++"; " 
    return Unit

interp (e1 :*: e2) = do
    v1 <- interp e1
    v2 <- interp e2
    return $ Pair (v1,v2)
interp (Fst e) = do 
    v <- interp e 
    pair fst v
interp (Snd e) = do 
    v <- interp e  
    pair snd v

interp TT = return $ Bl True
interp FF = return $ Bl False
interp (Not e) = do
    v <- interp e
    not' v
interp (e1 :==: e2) = do
    v1 <- interp e1
    v2 <- interp e2
    return $ Bl $ v1 == v2
interp (If e1 e2 e3) = do
    v1 <- interp e1
    if' v1 (interp e2) (interp e3)

interp Skip = return Unit
interp (e1 :& e2) = do 
    interp e1
    interp e2 

interp Nil = return $ List []
interp (e1 ::: e2) = do
    v1 <- interp e1
    v2 <- interp e2
    cons v1 v2
interp (LCase e1 e2 e3) = do
    v1 <- interp e1
    lcase v1 (interp e2) (interp e3)

interp (MkRef e) = do
    v <- interp e
    newRef v
interp (Read e) = do
    v <- interp e 
    getRef v
interp (ReadV x) = do
    v <- interp (Var x) 
    getRef v
interp (e1 := e2) = do
    v1 <- interp e1
    v2 <- interp e2
    setRef v1 v2
interp (x :!= e2) = do
    v1 <- interp (Var x)
    v2 <- interp e2
    setRef v1 v2

-- !!!!!!!!!!
--
-- Supply desugarings of Let and LetRec for which:
--    all declarations are independent
--    no variables are captured
--    malformed LetRec declarations (i.e. not of form x :=: v :\ e)
--      trigger an interpreter error.
--
interp (Let decls e) = 

interp (LetRec decls e) = 
--
--
-- !!!!!!!!!!

interp (While e1 e2) =
    let vs = freeVars (e1 :*: e2)
        name x = if x `elem` vs then name (x ++ "'") else x
        w = name "w"
        x = name "x"
    in interp $ LetRec [w :=: x :\ If e1 (e2 :& Var w :$: Skip) Skip] $ Var w :$: Skip

----------------------------------------------------------------------------------
--
-- Example terms
--

i = "x" :\ Var "x"
k = "x" :\ "y" :\ Var "x"
s = "f" :\ "g" :\ "x" :\ Var "f" :$: Var "x" :$: (Var "g" :$: Var "x")
omega = ("x" :\ Var "x" :$: Var "x") :$: ("x" :\ Var "x" :$: Var "x")


search = 
  LetRec ["srch" :=: "l" :\ "p" :\ "f" :\ "g" :\
             LCase (Var "l") 
                   (Var "g" :$: Skip) 
                   ("h" :\ "t" :\ 
                        If(Var "p" :$: Var "h") 
                           (Var "f" :$: Var "h") 
                           (Var "srch" :$: Var "t" :$: Var "p" :$: Var "f" :$: Var "g")
                   )
         ]
         (Var "srch")
          

memo = "f" :\ Let ["r" :!=: Nil] ("n" :\ 
           search :$: (ReadV "r") :$: ("x" :\ Fst (Var "x") :==: Var "n") :$: ("x" :\ Snd (Var "x")) :$:
                      ("_" :\ Let ["y" :=: Var "f" :$: Var "n"] 
                                  ("r" :!= (Var "n" :*: Var "y") ::: ReadV "r" :& Var "y")
                      )
           )

fibImp = "x" :\ "f" :\ 
    If (ReadV "x" :==: Con 0)
       ("f" :!= (Con 0))
       (Let ["g" :!=: Con 0] $ 
            "f" :!= Con 1 :& 
            While (Not $ Con 1 :==: ReadV "x")
                  (Let ["t" :=: ReadV "g"] $ 
                        "g" :!= ReadV "f" :&  
                        "f" :!= ReadV "f" :+: Var "t" :&
                        "x" :!= ReadV "x" :-: Con 1
                  ) :&
            Out (ReadV "f")
       )

fib = LetRec ["f" :=: ("n" :\ If (Var "n" :==: Con 0) (Con 0) 
                                 (If (Var "n" :==: Con 1) (Con 1) 
                                     ((Var "f" :$: (Var "n" :-: Con 1)) :+: (Var "f" :$: (Var "n" :-: Con 2))) 
                                 )
                      )
             ] $ Var "f"

memorec :: Term
memorec = 
  "f" :\ Let ["r" :!=: Nil] (
    LetRec ["mf" :=: ("n" :\ 
           search :$: (ReadV "r") :$: ("x" :\ Fst (Var "x") :==: Var "n") :$: ("x" :\ Snd (Var "x")) :$:
                      ("_" :\ Let ["y" :=: Var "f" :$: Var "mf" :$: Var "n"] 
                                  ("r" :!= (Var "n" :*: Var "y") ::: ReadV "r" :& Var "y")
                      )
           )
           ] 
           (Var "mf")
    )

fibmr = memorec :$: ("f" :\ "n" :\ 
    If (Var "n" :==: Con 0) (Con 0) 
           (If (Var "n" :==: Con 1) (Con 1) 
               ((Var "f" :$: (Var "n" :-: Con 1)) :+: (Var "f" :$: (Var "n" :-: Con 2))) 
           )
   )

plus = "v" :\ "x" :\ "y" :\ "v" :!= ReadV "v" :+: Con 1 :& Var "x" :+: Var "y"

fib' = "x" :\ Let ["c" :!=: Con 0] (
      LetRec ["f" :=: ("n" :\ If (Var "n" :==: Con 0) (Con 0) 
                                 (If (Var "n" :==: Con 1) (Con 1) 
                                     (plus :$: Var "c" :$: 
                                       (Var "f" :$: (Var "n" :-: Con 1)) :$: (Var "f" :$: (Var "n" :-: Con 2))
                                     )
                                 )
                      )
             ]
             ((Var "f" :$: Var "x") :*: ReadV "c")
      )
 
fibmr' = "x" :\ Let ["c" :!=: Con 0] (
  (memorec :$: ("f" :\ "n" :\ 
                  If (Var "n" :==: Con 0) (Con 0) 
                         (If (Var "n" :==: Con 1) 
                             (Con 1) 
                             (plus :$: Var "c" :$: (Var "f" :$: (Var "n" :-: Con 1)) :$: (Var "f" :$: (Var "n" :-: Con 2))) 
                         )
               ) :$: 
               (Var "x")
   ) :*: ReadV "c"
 )

fibImp' = "x" :\ "f" :\ Let ["c" :!=: Con 0] (
    If (ReadV "x" :==: Con 0)
       ("f" :!= (Con 0))
       (Let ["g" :!=: Con 0] $ 
            "f" :!= Con 1 :& 
            While (Not $ Con 1 :==: ReadV "x")
                  (Let ["t" :=: ReadV "g"] $ 
                        "g" :!= ReadV "f" :&  
                        "f" :!= plus :$: Var "c" :$: ReadV "f" :$: Var "t" :&
                        "x" :!= ReadV "x" :-: Con 1
                  ) :&
            Out (ReadV "f" :*: ReadV "c")
       )
  )