{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveAnyClass #-}
module Elab where

import Control.Exception

import qualified Data.Map.Strict as Map
import Data.Typeable

import qualified Presyntax as P
import Syntax
import Eval
import Control.Monad
import Systems
import Data.Traversable
import qualified Data.Set as Set
import Data.Set (Set)
import Data.Foldable

data TypeError
  = NotInScope String
  | UnifyError UnifyError
  | WrongFaces Value [([Value], Value, Elab.TypeError)]
  | InSpan     (Int, Int) (Int, Int) Elab.TypeError
  | IncompleteSystem P.Formula P.Formula
  | IncompatibleFaces (P.Formula, Term) (P.Formula, Term) Elab.TypeError
  | InvalidSystem (Set Face) (Set Face)
  deriving (Show, Typeable, Exception)

check :: Env -> P.Exp -> Value -> IO Term
check env (P.Span s e exp) wants =
  check env exp wants
    `catch` \case
      InSpan s e err -> throwIO $ InSpan s e err
      err -> throwIO $ InSpan s e err

check env exp (VSub a fm@(toFormula -> Just phi) el) = do
  tm <- check env exp a
  for (addFormula phi env) \env -> 
    let tm' = eval env tm
     in unifyTC env tm' el
  pure (InclSub (quote a) (quote fm) (quote el) tm)

check env (P.Lam s b) expected = do
  expc <- isPiOrPathType expected
  case expc of
    Left (_, d, r) -> do -- function
      bd <- check env { names = Map.insert s (makeValueGluingSub d s) (names env) } b (r (VVar s))
      pure (Lam s (quote d) bd)
    Right (a, x, y) -> do
      bd <- check env { names = Map.insert s (VI, VVar s) (names env) } b (a @@ VVar s)

      let t = Lam s I bd
          t' = eval env t

      checkBoundary env [s] t'
        [ ([VI0], x)
        , ([VI1], y)
        ]

      pure (PathI (quote a) (quote x) (quote y) s bd)

check env (P.Let v t d b) expected = do
  ty <- check env t VType
  let ty' = eval env ty
  d <- check env d ty'
  let d' = eval env d
  b <- check env { names = Map.insert v (ty', d') (names env) }  b expected
  pure (Let v (quote ty') d b)

check env (P.Pair a b) expected = do
  (_, fst, snd) <- isSigmaType expected
  a <- check env a fst
  let a' = eval env a
  b <- check env b (snd a')
  pure (Pair a b)

check env (P.Partial fs) ty = do
  let formula = orFormula (map fst fs)
  (extent, ty) <- isPartialType formula ty
  let ourFaces    = Systems.faces formula
      extentFaces = Systems.faces extent

  unless (formula == extent) $
    throwIO $ IncompleteSystem formula extent

  let range = formulaToTm $ toDNF formula

  rangeTm <- check env range VI
  let rangeTy = eval env rangeTm

  ts <- for fs $ \(fn, tn) -> do
    tms <- for (addFormula fn env) \env -> check env tn ty
    pure (fn, head tms)

  fmap (System . FMap . Map.fromList) $ for ts \(fn, tn) -> do
    for ts \(fm, tm) -> do
      when (possible (fn `P.And` fm)) do
        for_ (addFormula (fn `P.And` fm) env) $ \env ->
          unifyTC (env) (eval env tn) (eval env tm)
            `catch` \e -> throwIO (IncompatibleFaces (fn, tn) (fm, tm) e)
    pure (fn, tn)

check env exp expected = do
  (term, actual) <- infer env exp
  unifyTC env actual expected
  pure term

makeValueGluingSub :: Value -> String -> (Value, Value)
makeValueGluingSub ty@(VSub a phi a0) s = (ty, VOfSub a phi a0 (VVar s))
makeValueGluingSub ty s = (ty, VVar s)

addFormula :: P.Formula -> Env -> [Env]
addFormula (P.And x y) = addFormula x >=> addFormula y
addFormula (P.Or x y) = (++) <$> addFormula x <*> addFormula y
addFormula P.Top = pure
addFormula P.Bot = const []
addFormula (P.Is0 x) = \env -> pure env{ names = Map.insert x (VI, VI0) (names env) }
addFormula (P.Is1 x) = \env -> pure env{ names = Map.insert x (VI, VI1) (names env) }

unifyTC :: Env -> Value -> Value -> IO ()
unifyTC env a b = unify env a b `catch` \e -> const (throwIO (UnifyError (Mismatch a b))) (e :: UnifyError)

checkBoundary :: Env -> [String] -> Value -> [([Value], Value)] -> IO ()
checkBoundary env ns f = finish <=< go where
  go :: [([Value], Value)] -> IO [([Value], Value, Elab.TypeError)]
  go [] = pure []
  go ((ixs, vl):faces) = do
    let env' = foldr (\(x, t) env -> env { names = Map.insert x t (names env) }) env (zip ns (zip (repeat VI) ixs))
    t <- try $ unifyTC env' (foldl (@@) f ixs) vl
    case t of
      Right _ -> go faces
      Left e -> ((ixs, vl, e):) <$> go faces
  
  finish [] = pure ()
  finish xs = throwIO $ WrongFaces f xs

infer :: Env -> P.Exp -> IO (Term, Value)
infer env (P.Span s e exp) =
  infer env exp
    `catch` \case
      InSpan s e err -> throwIO $ InSpan s e err
      err -> throwIO $ InSpan s e err

infer env (P.Var s) =
  case Map.lookup s (names env) of
    Just (t, _) -> pure (Var s, t)
    Nothing     -> throwIO (NotInScope s)

infer env (P.App f x) = do
  (fun, ty) <- infer env f
  funt <- isPiOrPathType ty
  case funt of
    Left (_, dom, rng) -> do
      arg <- check env x dom
      let arg' = eval env arg
      pure (App fun arg, rng arg')
    Right (a, ai0, ai1) -> do
      arg <- check env x VI
      let arg' = eval env arg
      pure (PathP (quote a) (quote ai0) (quote ai1) fun arg, a @@ arg')

infer env (P.Pi s d r) = do
  dom <- check env d VType
  let d' = eval env dom
  rng <- check env { names = Map.insert s (d', VVar s) (names env) } r VType
  pure (Pi s dom rng, VType)

infer env (P.Sigma s d r) = do
  dom <- check env d VType
  let d' = eval env dom
  rng <- check env { names = Map.insert s (d', VVar s) (names env) } r VType
  pure (Sigma s dom rng, VType)

infer env P.Type = pure (Type, VType)
infer env P.I    = pure (I,    VType)
infer env P.I0   = pure (I0,   VI)
infer env P.I1   = pure (I1,   VI)

infer env (P.Cut e t) = do
  t <- check env t VType
  let t' = eval env t
  (, t') <$> check env e t'

infer env (P.IAnd x y) = do
  x <- check env x VI
  y <- check env y VI
  pure (IAnd x y, VI)

infer env (P.IOr x y) = do
  x <- check env x VI
  y <- check env y VI
  pure (IOr x y, VI)

infer env P.Path = do
  pure
    ( Lam "A" (quote index_t) $
        Lam "x" (App (Var "A") I0) $
          Lam "y" (App (Var "A") I1) $
            Path (Var "A") (Var "x") (Var "y")
    , VPi "A" index_t \a ->
        VPi "x" (a @@ VI0) \_ ->
          VPi "y" (a @@ VI1) (const VType))

infer env P.PartialT = do
  pure
    ( Lam "r" I $
        Lam "A" Type $
          Partial (Var "r") (Var "A")
    , VPi "I" VI \i ->
        VPi "A" VType (const VType))

infer env P.Comp = do
  let u_t a r = VPi "i" VI \i -> VPartial r (a @@ i)

  pure
    ( Lam "A" (quote index_t) $
        Lam "phi" I $
          Lam "u" (quote (u_t (VVar "A") (VVar "r"))) $
            Lam "a0" (Sub (App (Var "A") I0) (Var "phi") (App (Var "u") I0)) $
              Comp (Var "A") (Var "phi") (Var "u") (Var "a0")
    , VPi "A" index_t \a ->
        VPi "phi" VI \phi ->
          VPi "u" (u_t a phi) \u ->
            VPi "_" (VSub (a @@ VI0) phi (u @@ VI0)) \_ ->
              a @@ VI1
    )

infer env P.SubT = do
  pure
    ( Lam "A" Type $
        Lam "phi" I $
          Lam "u" (Partial (Var "phi") (Var "A")) $
            Sub (Var "A") (Var "phi") (Var "u")
    , VPi "A" VType \a ->
        VPi "phi" VI \phi ->
          VPi "_" (VPartial phi a) (const VType)
    )

infer env (P.INot x) = (, VI) . INot <$> check env x VI
infer env P.Lam{}    = error "can't infer type for lambda"

infer env (P.Let v t d b) = do
  ty <- check env t VType
  let ty' = eval env ty
  d <- check env d ty'
  let d' = eval env d
  (b, t) <- infer env{ names = Map.insert v (ty', d') (names env) } b
  pure (Let v ty d b, t)

infer env (P.Proj1 x) = do
  (t, ty) <- infer env x
  (_, d, _) <- isSigmaType ty
  pure (Proj1 t, d)
  
infer env (P.Proj2 x) = do
  (t, ty) <- infer env x
  let t' = eval env t
  (_, _, r) <- isSigmaType ty
  pure (Proj2 t, r (proj1 t'))

formulaToTm :: P.Formula -> P.Exp
formulaToTm (P.Is0 i) = P.INot (P.Var i)
formulaToTm (P.Is1 i) = P.Var i
formulaToTm (P.And x y) = P.IAnd (formulaToTm x) (formulaToTm y)
formulaToTm (P.Or x y) = P.IOr (formulaToTm x) (formulaToTm y)
formulaToTm P.Top = P.I1
formulaToTm P.Bot = P.I0

checkFormula :: Env -> P.Formula -> IO ()
checkFormula env P.Top = pure ()
checkFormula env P.Bot = pure ()
checkFormula env (P.And x y) = checkFormula env x *> checkFormula env y
checkFormula env (P.Or x y)  = checkFormula env x *> checkFormula env y
checkFormula env (P.Is0 x) =
  case Map.lookup x (names env) of
    Just (ty, _) -> unifyTC env ty VI
    Nothing -> throwIO (NotInScope x)
checkFormula env (P.Is1 x) =
  case Map.lookup x (names env) of
    Just (ty, _) -> unifyTC env ty VI
    Nothing -> throwIO (NotInScope x)

index_t :: Value
index_t = VPi "_" VI (const VType)