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

import Control.Exception
import Control.Monad

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Traversable
import Data.Set (Set)
import Data.Typeable
import Data.Foldable

import Eval

import qualified Presyntax as P

import Syntax

import Systems

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
    `catch` \er -> throwIO $ InSpan s e (UnifyError er)

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

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 (toDNF formula == toDNF 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 (fn /= fm && 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 (VPartial fm@(toFormula -> Just phi) a) = do
  case addFormula phi env of
    [] -> pure $ System (FMap mempty)
    (x:xs) -> do
      tm <- check x exp a
      for_ xs $ \e -> check e exp a
      pure 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 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
    `catch` \er -> throwIO $ InSpan s e (UnifyError er)

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, ty) <- infer env d
  case ty of
    VType -> pure VType
    VTypeω -> pure VTypeω
    _ -> throwIO . UnifyError $ NotSort ty
  let d' = eval env dom
  (rng, rng_t) <-
    infer env { names = Map.insert s (makeValueGluingSub d' s) (names env) } r
  case ty of
    VType -> pure VType
    VTypeω -> pure VTypeω
    _ -> throwIO . UnifyError $ NotSort ty
  pure (Pi s dom rng, rng_t)

infer env (P.Sigma s d r) = do
  (dom, ty) <- infer env d
  rng_t <-
    case ty of
      VType -> pure VType
      VTypeω -> pure VTypeω
      _ -> throwIO . UnifyError $ NotSort ty
  let d' = eval env dom
  rng <- check env { names = Map.insert s (d', VVar s) (names env) } r rng_t
  pure (Sigma s dom rng, rng_t)

infer env P.Type  = pure (Type, 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 =
  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 =
  pure
    ( Lam "r" I $
        Lam "A" Type $
          Partial (Var "r") (Var "A")
    , VPi "I" VI \i ->
        VPi "A" VType (const VTypeω))

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

infer env P.Comp = do
  let
    u_t a r = VPi "i" VI \i -> VPartial r (a @@ i)
    index_t :: Value
    index_t = VPi "_" VI (const VType)

  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 =
  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.GlueTy =
  pure
    ( Lam "A" Type $
        Lam "phi" I $
          Lam "T" (Partial (Var "phi") Type) $
            Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $
              GlueTy (Var "A") (Var "phi") (Var "T") (Var "e")
    , VPi "A" VType \a ->
        VPi "phi" VI \phi ->
          VPi "T" (VPartial phi VType) \t ->
            VPi "e" (VPartialP phi (equiv t a)) \_ ->
              VType
    )

infer env P.Glue =
  pure
    ( Lam "A" Type $
        Lam "phi" I $
          Lam "T" (Partial (Var "phi") Type) $
            Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $
              Lam "t" (PartialP (Var "phi") (Var "T")) $ 
                Lam "a" (Var "A") $
                  Glue (Var "A") (Var "phi") (Var "T") (Var "e") (Var "t") (Var "a")
    , VPi "A" VType \a ->
        VPi "phi" VI \phi ->
          VPi "T" (VPartial phi VType) \t ->
            VPi "e" (VPartialP phi (equiv t a)) \_ ->
              VPi "_" (VPartialP phi t) \_ ->
                VPi "_" a \_ -> VType
    )

infer env P.Unglue =
  pure
    ( Lam "A" Type $
        Lam "phi" I $
          Lam "T" (Partial (Var "phi") Type) $
            Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $
              Lam "g" (GlueTy (Var "A") (Var "phi") (Var "T") (Var "e")) $
                Unglue (Var "A") (Var "phi") (Var "T") (Var "e") (Var "g")
    , VPi "A" VType \a ->
        VPi "phi" VI \phi ->
          VPi "T" (VPartial phi VType) \t ->
            VPi "e" (VPartialP phi (equiv t a)) \e ->
              VPi "_" (VGlue a phi t e) \_ -> a
    )

infer env P.Bool = pure (Bool, VType)
infer env P.Tt   = pure (Tt, VBool)
infer env P.Ff   = pure (Ff, VBool)

infer env P.If =
  pure
    (Lam "P" (Pi "_" Bool Type) $
      Lam "x" (App (Var "P") Tt) $
        Lam "y" (App (Var "P") Ff) $
          Lam "b" Bool $
            If (Var "P") (Var "x") (Var "y") (Var "b")
    , VPi "P" (VPi "_" VBool (const VType)) \p ->
        VPi "_" (p @@ VTrue) \_ ->
          VPi "_" (p @@ VFalse) \_ ->
            VPi "x" VBool \x -> p @@ x)


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ω)