{-# LANGUAGE LambdaCase #-}
module Algo where

import Data.Maybe (fromMaybe)
import Data.Map (Map)
import qualified Data.Map as M

import Control.Monad.RWS
import Control.Monad.Except

import Data.List

type Id = String 

data Exp
  = Var Id
  | App Exp Exp
  | Abs Id Exp
  | Let Id Exp Exp Exp
  | Pi Id Exp Exp
  | Type 
  deriving Show

data Val
  = VGen Int 
  | VApp Val Val
  | VType
  | VClos Env Exp 
  deriving Show

type Env = Map Id Val

data Err
  = Lookup
  | ExpectedPi
  | ExpectedType
  | NoInfer
  | AppError
  deriving Show

data Result = Result (Either Err Bool) [String]

instance Show Result where
  show (Result r s) = header <> "\n" <> intercalate "\n" s
    where header = case r of
                     Left e -> show e
                     Right True  -> "Checked successfully!"
                     Right False -> "Types do not match!"

type Conv  = ExceptT Err (RWS Env [String] Int)
type Check = ExceptT Err (RWS (Env, Env) [String] Int)

li :: Conv a -> Check a
li conv = do
  state   <- get
  (env,_) <- ask
  let (a, s, w) = runRWS (runExceptT conv) env state
  put s
  tell w
  case a of
    Left e  -> throwError e
    Right v -> pure v

maybeConv :: MonadError e m => e -> Maybe a -> m a
maybeConv _ (Just a) = pure a
maybeConv e Nothing  = throwError e

lookupEnv :: Id -> Conv Val
lookupEnv x = ask >>= maybeConv Lookup . M.lookup x

insertEnv :: Id -> Val -> Conv a -> Conv a
insertEnv x v = local (M.insert x v)

withEnv :: Env -> Conv a -> Conv a
withEnv = local . const

localCh :: (Env -> Env, Env -> Env) -> Check a -> Check a
localCh (f,g) = local h
  where h (a,b) = (f a, g b)

fresh :: (Num i, MonadState i m) => m i
fresh = do
  i <- get
  put (i+1)
  pure i

app :: Val -> Val -> Conv Val
app u v = case u of
            VClos env (Abs x e) -> insertEnv x v (eval e)
            _                   -> pure (VApp u v)

eval :: Exp -> Conv Val
eval = \case
          Var x         -> lookupEnv x

          App e1 e2     -> do
            e1' <- eval e1
            e2' <- eval e2
            app e1' e2'

          Let x e1 _ e3 -> do
            e1' <- eval e1
            insertEnv x e1' (eval e3)

          Type          -> pure (VType)
          e             -> ask >>= pure . (flip VClos) e

whnf :: Val -> Conv Val
whnf = \case
          VApp u w    -> do
            u' <- whnf u
            w' <- whnf w
            app u' w'
          VClos env e -> withEnv env (eval e)
          v           -> pure v

eqVal :: Val -> Val -> Conv Bool
eqVal u1 u2 = do
  u1' <- whnf u1
  u2' <- whnf u2
  tell ["checking equality of terms '" <> show u1 <> "' and '" <> show u2 <> "'."]
  tell ["with value representations '" <> show u1' <> "' and '" <> show u2' <> "'."]
  case (u1', u2') of
    (VType, VType) -> pure True
    (VApp t1 w1, VApp t2 w2) -> (&&) <$> eqVal t1 t2 <*> eqVal w1 w2
    (VGen k1, VGen k2) -> pure (k1 == k2)

    (VClos env1 (Abs x1 e1), VClos env2 (Abs x2 e2)) -> do
      v <- VGen <$> fresh
      eqVal (VClos (M.insert x1 v env1) e1)
            (VClos (M.insert x2 v env2) e2)

    (VClos env1 (Pi x1 a1 b1), VClos env2 (Pi x2 a2 b2)) -> do
      v <- VGen <$> fresh
      (&&) <$> eqVal (VClos env1 a1)
                     (VClos env2 a2)
           <*> eqVal (VClos (M.insert x1 v env1) b1)
                     (VClos (M.insert x1 v env2) b2)
    _ -> pure False

checkType :: Exp -> Check Bool
checkType = (flip checkExp) VType

checkExp :: Exp -> Val -> Check Bool
checkExp e v' = do
  v <- li (whnf v')
  case e of
    Abs x n ->
      case v of
        VClos env (Pi y a b) -> do
          v <- VGen <$> fresh
          localCh (M.insert x v, M.insert x (VClos env a))
            (checkExp n (VClos (M.insert y v env) b))

        _ -> throwError ExpectedPi

    Pi x a b ->
      case v of
        VType -> do
          v <- VGen <$> fresh
          rho <- fst <$> ask
          (&&) <$> checkType a <*> localCh (M.insert x v, M.insert x (VClos rho a)) (checkType b)

        _ -> throwError ExpectedType

    Let x e1 e2 e3 -> do
      etr <- li (eval e1)
      ety <- li (eval e2)
      (&&) <$> checkType e2 <*> localCh (M.insert x etr, M.insert x ety) (checkExp e3 v)

    _ -> do 
      e' <- inferExp e
      li (eqVal v e')

inferExp :: Exp -> Check Val
inferExp = \case
  Var id -> maybeConv Lookup . M.lookup id =<< snd <$> ask
  App e1 e2 -> 
    inferExp e1 >>= li . whnf >>= \case
      VClos env (Pi x a b) -> do
        rho <- fst <$> ask
        bool <- checkExp e2 (VClos env a)
        if bool
           then pure (VClos (M.insert x (VClos rho e2) env) b)
           else throwError AppError
      _ -> throwError ExpectedPi
  Type -> pure VType
  _    -> throwError NoInfer

typecheck :: Exp -> Exp -> Result
typecheck m a = resolve ((&&) <$> checkType a <*> checkExp m (VClos M.empty a))
  where resolve action = let (r,_,w) = (runRWS (runExceptT action) (M.empty, M.empty) 0)
                         in Result r w

test1 :: Result
test1 = typecheck (Abs "A" (Abs "x" (Var "x")))
                  (Pi "A" Type (Pi "x" (Var "A") (Var "A")))

test2 :: Result
test2 = typecheck (Abs "A" (Abs "B" (Abs "f" (Abs "x" (App (Var "f") (Var "x"))))))
                  (Pi "A" Type 
                          (Pi "B" (Pi "_" (Var "A") Type)
                                  (Pi "f" (Pi "a" (Var "A") (App (Var "B") (Var "a"))) 
                                          (Pi "x" (Var "A") (App (Var "B") (Var "x"))))))