{-# LANGUAGE TupleSections #-}

module Eval where

import Control.Monad.Except (MonadError (..))
import Control.Monad.Reader
import qualified Data.Map as M
import Data.Text (Text)
import Errors
import Expr

type Env = M.Map Text Expr

-- substitute s for k *AND* decrement indices; only use after reducing a redex.
subst :: Integer -> Expr -> Expr -> Expr
subst k s (Var n x)
    | k == n = s
    | n > k = Var (n - 1) x
    | otherwise = Var n x
subst _ _ (Free s) = Free s
subst _ _ Star = Star
subst _ _ Square = Square
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)

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

-- reduce until β reducts 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) = envLookup n
whnf e = pure e

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) = envLookup n
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'

-- naive beta equivalence check
betaEquiv :: Expr -> Expr -> ReaderT Env Result Bool
betaEquiv e1 e2 = (==) <$> normalize e1 <*> normalize e2

-- this slightly smarter beta equivalence check is a little buggy,
-- failing to notice that `add one one` and `two` are beta equivalent in the
-- example file

-- 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 k1 _, Var k2 _) -> pure $ k1 == k2
--             (Free n, Free m) -> pure $ n == m
--             (Free n, e) -> envLookup n >>= betaEquiv e
--             (e, Free n) -> envLookup n >>= betaEquiv e
--             (Star, Star) -> pure True
--             (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
--             _ -> pure False -- remaining cases impossible or false

checkBeta :: Env -> Expr -> Expr -> Result Bool
checkBeta env e1 e2 = runReaderT (betaEquiv e1 e2) env

isSortPure :: Expr -> Bool
isSortPure Star = True
isSortPure Square = True
isSortPure _ = False

isSort :: Expr -> ReaderT Env Result (Bool, Expr)
isSort s = (,s) . isSortPure <$> normalize s