{-# LANGUAGE LambdaCase #-}
module Eval
  ( fv
  , subst
  , normalstep
  , pickFresh
  , reduce
  ) where

import           AST
import qualified Data.Map as Map
import qualified Data.Set as Set

-- | Return the free variables in an expression
fv :: Term -> Set.Set String
fv (TVar x)            = Set.singleton x
fv (TLam x m)          = Set.delete x (fv m)
fv (TApp m1 m2)        = Set.union (fv m1) (fv m2)
fv (TConstr _)         = Set.empty -- good
fv (TLet x m1 m2)      = Set.delete x (Set.union (fv m1) (fv m2)) -- good
fv (TLit _)            = Set.empty -- good
fv (TBinop _ m1 m2)    = Set.union (fv m1) (fv m2) -- good
fv (TLetAnn x _ m1 m2) = Set.delete x (Set.union (fv m1) (fv m2)) -- good
fv (TCase m1 [])       = fv m1
fv (TCase m1 ((PConstr _ pvars, m2):pms)) = Set.union (Set.union (fv m1) (Set.difference l (fv m2))) (fv (TCase m1 pms)) where l = pvset pvars
fv (TCase _ _) = error "PConstr type not handled" -- fv (TCase m1 ((p@(PVar x), m2):pms))
fv (TIf m1 m2 m3)      = Set.union (Set.union (fv m1) (fv m2)) (fv m3)


-- helper method to make list of patterns
pvset :: [Pattern] -> Set.Set String
pvset []            = Set.empty
pvset ((PVar x):ls) = Set.insert x (pvset ls)
pvset _             = error "did not handle"


-- | Substitute n for x in e, avoiding name capture
--    subst n x e     e[x := n]
subst :: Term -> String -> Term -> Term
subst n x (TVar y)
    | x == y                    = n
    | otherwise                 = TVar y
subst n x (TApp m1 m2)          = TApp (subst n x m1) (subst n x m2)
subst n x (TLam y m)
    | x == y                    = TLam y m
    | Set.notMember y (fv n)    = TLam y (subst n x m)
    | otherwise                 = TLam y2 (subst n x m')
        where y2 = pickFresh (Set.union (fv n) (fv m)) y
              m' = subst (TVar y2) y m
subst _ _ (TConstr y)           = TConstr y -- good
subst _ _ (TLit y)              = TLit y -- good
subst n x (TLet y m1 m2) -- same as lambda
    | x == y                    = TLet y m1 m2
    | Set.notMember y (fv n)    = TLet y (subst n x m1) (subst n x m2)
    | otherwise                 = TLet y2 (subst n x m1') (subst n x m2')
        where y2 = pickFresh (Set.union (Set.union (fv n) (fv m1)) (fv m2)) y
              m1'= subst (TVar y2) y m1
              m2'= subst (TVar y2) y m2
subst n x (TLetAnn y a m1 m2) -- same as lambda
    | x == y                    = TLetAnn y a m1 m2
    | Set.notMember y (fv n)    = TLetAnn y a (subst n x m1) (subst n x m2)
    | otherwise                 = TLetAnn y2 a (subst n x m1') (subst n x m2')
        where y2 = pickFresh (Set.union (Set.union (fv n) (fv m1)) (fv m2)) y
              m1'= subst (TVar y2) y m1
              m2'= subst (TVar y2) y m2
subst n x (TBinop b m1 m2)      = TBinop b (subst n x m1) (subst n x m2) --just like tapp
subst n x (TIf m1 m2 m3)        = TIf (subst n x m1) (subst n x m2) (subst n x m3) -- just like tapp
-- TODO - check TCase
subst n x (TCase m1 ps)         = TCase (subst n x m1) (subsP n x ps) -- same as lambda, might have to rename patter vars (check this)

pvrep :: Term -> [Pattern] -> Term -> ([Pattern], Term)
pvrep _ [] m2 = ([], m2)
pvrep n ((PVar y):ys) m2
    | Set.notMember y (fv n)    = (PVar y:ys', m2')
    | otherwise                 = (PVar y2:ys'', m2''')
        where y2 = pickFresh (Set.union (fv n) (fv m2)) y
              (ys'', m2''') = pvrep n ys m2''
              m2''= subst (TVar y2) y m2
              (ys', m2') = pvrep n ys m2
pvrep _ _ _ = error "pattern type not implemented"

subsP :: Term -> String -> [(Pattern, Term)] -> [(Pattern, Term)]
subsP _ _ [] = []
subsP n x ((PConstr s pvs, m2):pms) = (PConstr s pvs2, m2''):subsP n x pms
    where (pvs2, m2') = pvrep n pvs m2
          m2'' = subst n x m2'
subsP _ _ _ = error "pattern type not implemented"



type Env = Map.Map String Term

-- | Take a single step in normal order reduction or return Nothing
normalstep :: Term -> Env -> (Maybe Term, Env)

normalstep (TVar x) env     = (Map.lookup x env, env)

normalstep (TConstr _) env  = (Nothing, env)
normalstep (TLit _) env     = (Nothing, env)

normalstep (TApp (TLam x m) n) env      = (Just (subst n x m), env)

normalstep (TLetAnn x _ m1 m2) env      = (Just $ TLet x m1 m2, env) -- drop annotation
normalstep (TLet x m1 m2)  env          = (Just $ subst m1 x m2, Map.insert x m1 env)

normalstep (TLam x m) env = case normalstep m env of
   (Nothing, envb) -> (Nothing, envb)
   (Just mb, envb) -> (Just (TLam x mb), envb)

normalstep (TApp m n) env = case normalstep m env of
   (Just m2, envb) -> (Just $ TApp m2 n, envb)
   (Nothing, envb) -> case normalstep n envb of
       (Just n2, envc) -> (Just $ TApp m n2, envc)
       (Nothing, envc) -> (Nothing, envc)

normalstep (TBinop b m1 m2) env = case normalstep m1 env of
    (Just m1b, envb) -> (Just $ TBinop b m1b m2, envb)
    (Nothing, envb) -> case normalstep m2 envb of
      (Nothing, envc)  -> (evalBinop b m1 m2, envc)
      (Just m2b, envc) -> (Just $ TBinop b m1 m2b, envc)

normalstep (TIf m1 m2 m3) env = case normalstep m1 env of
   (Nothing, envb) -> case m1 of
      TLit (LBool b) -> if b then (Just m2, envb) else (Just m3, envb)
      _              -> error "error: if condition is not a boolean"
   (Just m1b, envb) -> (Just (TIf m1b m2 m3), envb)

normalstep (TCase t pl) env = case auxCheckPList pl of
        Just m -> (Just m, env)
        Nothing -> case normalstep t env of
            (Nothing, envb) -> (Nothing, envb)
            (Just tb, envb) -> (Just (TCase tb pl), envb)

    where
        auxCheckPList :: [(Pattern, Term)] -> Maybe Term
        auxCheckPList [] = Nothing
        auxCheckPList ((p, m):ps) = case checkPattern t p of
            Nothing -> auxCheckPList ps
            Just l  -> Just (foldr (\(x, y) z -> subst y x z) m l)


checkPattern :: Term -> Pattern -> Maybe [(String, Term)]
checkPattern t (PConstr s' args)         = case extractTermList t s' of
    Nothing -> Nothing
    Just l  -> Just (zip vars l)
        where vars = map (\case
                                PVar x -> x
                                _      -> error "pattern type not implemented"
                ) args
checkPattern _ _ = error "pattern type not implemented"

extractTermList :: Term -> String -> Maybe [Term]
extractTermList (TConstr s) s' | s == s' = Just []
extractTermList (TApp m1 m2) s' = case extractTermList m1 s' of
    Nothing -> Nothing
    Just l  -> Just (l ++ [m2])
extractTermList _ _ = Nothing


-- auxPattern :: Term -> [Pattern] -> Maybe [(String, Term)]
-- auxPattern _ [] = Just []
-- auxPattern t [PVar x] = Just [(x, t)]
-- auxPattern (TApp t1 t2) ((PVar x):ps) = case auxPattern t2 ps of
--     Nothing -> Nothing
--     Just l  -> Just ((x, t1):l)
-- auxPattern _ _ = Nothing


evalBinop :: Binop -> Term -> Term -> Maybe Term
evalBinop Add (TLit (LInt i1)) (TLit (LInt i2)) =  Just $ TLit (LInt (i1 + i2))
evalBinop Sub (TLit (LInt i1)) (TLit (LInt i2)) =  Just $ TLit (LInt (i1 - i2))
evalBinop Mul (TLit (LInt i1)) (TLit (LInt i2)) =  Just $ TLit (LInt ( i1 * i2))
evalBinop Equal m1 m2                           = Just $ TLit $ LBool (m1 == m2) -- todo: better notion of equality
evalBinop _ _ _                                 = Nothing


-- | Return a "fresh" name not already in the set.
-- Tries x' then x'', etc.
pickFresh :: Set.Set String -> String -> String
pickFresh s = pickFresh'
  where pickFresh' n | n `Set.notMember` s = n
        pickFresh' n                       = pickFresh' $ n ++ "'"


reduce :: Term -> Term
reduce t = aux t Map.empty
    where
    aux tb env = case normalstep tb env of
        (Nothing, _)    -> tb
        (Just tc, envb) -> aux tc envb


-- | Repeatedly apply a function to transform a value, returning the list
-- of steps it took.  The result list starts with the given initial value
-- repeatedly :: (a -> Maybe a) -> a -> [a]
-- repeatedly f = repeatedly'
--     where repeatedly' x = x : maybe [] repeatedly' (f x)

-- -- | Print out the series of normal order reduction steps
-- printnormal :: String -> IO ()
-- printnormal s = let Program (_, t) = parse s in
--   mapM_ print $ repeatedly (\t' -> fst $ normalstep t' Map.empty) t


-- printLast :: String -> IO ()
-- printLast s = let Program (_, t) = parse s in
--     let arr = repeatedly (\t' -> fst $ normalstep t' Map.empty) t in
--         print $ last arr
