{-# LANGUAGE NamedFieldPuns #-}

module Eval where

import Control.Monad.Error.Class
import qualified Data.Map.Strict as M
import Errors
import Expr
import Relude.Extra.Lens

data Definition = Def {_ty :: Expr, _val :: Expr}

tyL :: Lens' Definition Expr
tyL = lens _ty setter
  where
    setter (Def{_val}) new = Def{_val, _ty = new}

valL :: Lens' Definition Expr
valL = lens _val setter
  where
    setter (Def{_ty}) new = Def{_val = new, _ty}

type Env = Map Text Definition

emptyEnv :: Env
emptyEnv = M.empty

showEnvEntry :: Text -> Definition -> Text
showEnvEntry k Def{_ty = t} = k <> " : " <> prettyT t

dumpEnv :: Env -> IO ()
dumpEnv = void . M.traverseWithKey ((putTextLn .) . showEnvEntry)

-- substitute s for k *AND* decrement indices; only use after reducing a redex.
subst :: Integer -> Expr -> Expr -> Expr
subst k s (Var x n)
    | k == n = s
    | n > k = Var x (n - 1)
    | otherwise = Var x n
subst _ _ (Free s) = Free s
subst _ _ (Axiom s) = Axiom s
subst _ _ Star = Star
subst _ _ (Level i) = Level i
subst k s (App m n) = App (subst k s m) (subst k s n)
subst k s (Abs x m n) = Abs x (subst k s m) (subst (k + 1) (incIndices s) n)
subst k s (Pi x m n) = Pi x (subst k s m) (subst (k + 1) (incIndices s) n)
subst k s (Let x t v b) = Let x t (subst k s v) (subst (k + 1) (incIndices s) b)
subst k s (Sigma x m n) = Sigma x (subst k s m) (subst (k + 1) (incIndices s) n)
subst k s (Pair m n) = Pair (subst k s m) (subst k s n)
subst k s (Pi1 x) = Pi1 (subst k s x)
subst k s (Pi2 x) = Pi2 (subst k s x)

envLookupVal :: Text -> ReaderT Env Result Expr
envLookupVal n = asks ((_val <$>) . M.lookup n) >>= maybe (throwError $ UnboundVariable n) pure

envLookupTy :: Text -> ReaderT Env Result Expr
envLookupTy n = asks ((_ty <$>) . M.lookup n) >>= maybe (throwError $ UnboundVariable n) pure

reduce :: Expr -> ReaderT Env Result Expr
reduce (App (Abs _ _ v) n) = pure $ subst 0 n v
reduce (App m n) = App <$> reduce m <*> reduce n
reduce (Abs x t v) = Abs x <$> reduce t <*> reduce v
reduce (Pi x t v) = Pi x <$> reduce t <*> reduce v
reduce (Free n) = envLookupVal n
reduce (Let _ _ v b) = pure $ subst 0 v b
reduce (Sigma x a b) = Sigma x <$> reduce a <*> reduce b
reduce (Pair a b) = Pair <$> reduce a <*> reduce b
reduce (Pi1 (Pair a _)) = pure a
reduce (Pi2 (Pair _ b)) = pure b
reduce (Pi1 x) = Pi1 <$> reduce x
reduce (Pi2 x) = Pi2 <$> reduce x
reduce e = pure e

normalize :: Expr -> ReaderT Env Result Expr
normalize e = do
    e' <- reduce e
    if e == e'
        then pure e
        else normalize e'

-- reduce until β reducts or let simplifications are impossible in head position
whnf :: Expr -> ReaderT Env Result Expr
whnf (App (Abs _ _ v) n) = whnf $ subst 0 n v
whnf (App m n) = do
    m' <- whnf m
    if m == m'
        then pure $ App m n
        else whnf $ App m' n
whnf (Free n) = envLookupVal n >>= whnf
whnf (Let _ _ v b) = whnf $ subst 0 v b
whnf (Pi1 (Pair a _)) = pure a
whnf (Pi2 (Pair _ b)) = pure b
whnf (Pi1 x) = do
    x' <- whnf x
    if x == x'
        then pure $ Pi1 x
        else whnf $ Pi1 x'
whnf (Pi2 x) = do
    x' <- whnf x
    if x == x'
        then pure $ Pi2 x
        else whnf $ Pi2 x'
whnf e = pure e

betaEquiv :: Expr -> Expr -> ReaderT Env Result Bool
betaEquiv e1 e2
    | e1 == e2 = pure True
    | otherwise = do
        e1' <- whnf e1
        e2' <- whnf e2
        case (e1', e2') of
            (Var _ n1, Var _ n2) -> pure $ n1 == n2
            (Free n, Free m) -> pure $ n == m
            (Free n, e) -> envLookupVal n >>= betaEquiv e
            (e, Free n) -> envLookupVal n >>= betaEquiv e
            (Axiom n1, Axiom n2) -> pure $ n1 == n2
            (Star, Star) -> pure True
            (Level i, Level j) -> pure $ i == j
            (App m1 n1, App m2 n2) -> (&&) <$> betaEquiv m1 m2 <*> betaEquiv n1 n2
            (Abs _ t1 v1, Abs _ t2 v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2 -- i want idiom brackets
            (Pi _ t1 v1, Pi _ t2 v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2
            (Let _ _ v b, e) -> betaEquiv (subst 0 v b) e
            (e, Let _ _ v b) -> betaEquiv (subst 0 v b) e
            (Sigma _ a b, Sigma _ a' b') -> (&&) <$> betaEquiv a a' <*> betaEquiv b b'
            (Pair a b, Pair a' b') -> (&&) <$> betaEquiv a a' <*> betaEquiv b b'
            (Pi1 x, Pi1 x') -> betaEquiv x x'
            (Pi2 x, Pi2 x') -> betaEquiv x x'
            _ -> pure False -- remaining cases impossible, false, or redundant

betaEquiv' :: Expr -> Expr -> Expr -> ReaderT Env Result ()
betaEquiv' ctxt e1 e2 = unlessM (betaEquiv e1 e2) $ throwError $ NotEquivalent e1 e2 ctxt

checkBeta :: Env -> Expr -> Expr -> Expr -> Result ()
checkBeta env e1 e2 ctxt = case runReaderT (betaEquiv e1 e2) env of
    Left err -> Left err
    Right False -> Left $ NotEquivalent e1 e2 ctxt
    Right True -> Right ()

isSortPure :: Expr -> Bool
isSortPure (Level _) = True
isSortPure _ = False

isSort :: Expr -> ReaderT Env Result Bool
isSort = fmap isSortPure . whnf