module Obs.Typing

import Data.Vect

import Obs.Logging
import Obs.NormalForm
import Obs.Sort
import Obs.Substitution
import Obs.Term

import System

import Text.Bounded
import Text.PrettyPrint.Prettyprinter
import Text.PrettyPrint.Prettyprinter.Render.Terminal

%default total

-- Loggings ----------------------------------------------------------------------

Rename (Logging ann . NormalForm) where
  rename t f = pure $ rename !t f

mismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
mismatch bounds lhs rhs = fatal $
  MkBounded
    (pretty "expected" <++> lhs <+> comma <+> softline <+> pretty "got" <++> rhs)
    False
    bounds

typeMismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
typeMismatch bounds lhs rhs = inScope "type mismatch" $ mismatch bounds lhs rhs

sortMismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
sortMismatch bounds lhs rhs = inScope "sort mismatch" $ mismatch bounds lhs rhs

-- Substitution ----------------------------------------------------------------

mergeName : String -> String -> String
mergeName "" s' = s'
mergeName "_" s' = s'
mergeName s s' = s

wkn : Vect k String -> (Fin n -> Logging ann (NormalForm m)) -> Fin (k + n) -> Logging ann (NormalForm (k + m))
wkn []            f = f
wkn (var :: vars) f =
  add
    (Logging ann . NormalForm)
    (pure $ Ntrl $ Var var FZ)
    (\i => pure $ rename !(wkn vars f i) FS)

covering partial
substCnstr : Constructor n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (Constructor m)
covering partial
substNtrl  : Neutral n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
covering partial
subst      : NormalForm n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
covering partial
subst1     : NormalForm n -> NormalForm (S n) -> Logging ann (NormalForm n)
covering partial
doApp      : NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doFst      : NormalForm n -> Logging ann (NormalForm n)
covering partial
doSnd      : NormalForm n -> Logging ann (NormalForm n)
covering partial
doEqual    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doEqualL   : Nat -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doEqualR   : Nat -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doEqualU   : Nat -> Constructor n -> Constructor n -> Logging ann (NormalForm n)
covering partial
doCastL    : NormalForm n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doCastR    : Constructor n -> NormalForm n -> NormalForm n -> Logging ann (NormalForm n)
covering partial
doCastU    : Constructor n -> Constructor n -> NormalForm n -> Logging ann (NormalForm n)

substCnstr (Sort s) f = pure $ Sort s
substCnstr (Pi s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Pi s s' var a b)
substCnstr (Lambda var t) f = do
  t <- subst t (wkn [var] f)
  pure (Lambda var t)
substCnstr (Sigma s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Sigma s s' var a b)
substCnstr (Pair t u) f = do
  t <- subst t f
  u <- subst u f
  pure (Pair t u)
substCnstr Top f = pure $ Top
substCnstr Bottom  f = pure $ Bottom

substNtrl (Var var i) f = do
  Ntrl (Var var' j) <- f i
    | t => pure t
  pure (Ntrl (Var (mergeName var' var) j))
substNtrl (App t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doApp t u
substNtrl (Fst t) f = do
  t <- substNtrl t f
  doFst t
substNtrl (Snd t) f = do
  t <- substNtrl t f
  doSnd t
substNtrl Absurd f = pure $ Ntrl Absurd
substNtrl (Equal a t u) f = do
  a <- substNtrl a f
  t <- subst t f
  u <- subst u f
  doEqual a t u
substNtrl (EqualL i t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doEqualL i t u
substNtrl (EqualR i t u) f = do
  t <- substCnstr t f
  u <- substNtrl u f
  doEqualR i t u
substNtrl (EqualU i t u) f = do
  t <- substCnstr t f
  u <- substCnstr u f
  doEqualU i t u
substNtrl (CastL a b t) f = do
  a <- substNtrl a f
  b <- subst b f
  t <- subst t f
  doCastL a b t
substNtrl (CastR a b t) f = do
  a <- substCnstr a f
  b <- substNtrl b f
  t <- subst t f
  doCastR a b t
substNtrl (CastU a b t) f = do
  a <- substCnstr a f
  b <- substCnstr b f
  t <- subst t f
  doCastU a b t

subst (Ntrl t)  f = substNtrl t f
subst (Cnstr t) f = map Cnstr $ substCnstr t f
subst Irrel     f = pure Irrel

subst1 t u = subst u (add (Logging ann . NormalForm) (pure t) (pure . point))

doApp (Ntrl t) u = pure $ Ntrl (App t u)
doApp Irrel u = pure $ Irrel
doApp (Cnstr (Lambda var t)) u = subst t (add (Logging ann . NormalForm) (pure u) (pure . point))
doApp (Cnstr t) u = inScope "bug" $ inScope "wrong constructor in apply" $ fatal t

doFst (Ntrl t) = pure $ Ntrl (Fst t)
doFst Irrel = pure $ Irrel
doFst (Cnstr (Pair t u)) = pure $ t
doFst (Cnstr t) = inScope "bug" $ inScope "wrong constructor in fst" $ fatal t

doSnd (Ntrl t) = pure $ Ntrl (Snd t)
doSnd Irrel = pure $ Irrel
doSnd (Cnstr (Pair t u)) = pure $ u
doSnd (Cnstr t) = inScope "bug" $ inScope "wrong constructor in snd" $ fatal t

doEqual (Ntrl a) t u = pure $ Ntrl (Equal a t u)
doEqual Irrel t u = inScope "bug" $ inScope "wrong type over equal" $ fatal "Irrel"
doEqual (Cnstr (Sort Prop)) t u =
  pure $
  Cnstr (Sigma Prop Prop "_"
    (Cnstr $ Pi Prop Prop "_" t (weaken 1 u))
    (weaken 1 $ Cnstr $ Pi Prop Prop "_" u (weaken 1 t)))
doEqual (Cnstr (Sort (Set k))) t u = doEqualL k t u
doEqual (Cnstr (Pi s s' var a b)) t u = do
  eqLhs <- doApp (weaken 1 t) (Ntrl $ Var var FZ)
  eqRhs <- doApp (weaken 1 u) (Ntrl $ Var var FZ)
  eq <- doEqual b eqLhs eqRhs -- b in Set because Pi in Set.
  pure $ Cnstr (Pi s Prop var a eq)
doEqual (Cnstr (Sigma s s' var a b)) t u = do
  t1 <- doFst t
  u1 <- doFst u
  eq1 <- case s of
    Prop => pure $ Cnstr Top
    (Set i) => doEqual a t1 u1
  eq2 <- case s' of
    Prop => pure $ Cnstr Top
    (Set i) => do
      bt <- subst1 t1 b
      bu <- subst1 u1 b
      t2 <- doSnd t
      t2 <- doCastL bt bu t2
      u2 <- doSnd u
      doEqual bu t2 u2
  pure $ Cnstr (Sigma Prop Prop "_" eq1 (weaken 1 eq2))
doEqual (Cnstr Top) t u = pure $ Cnstr Top
doEqual (Cnstr Bottom) t u = pure $ Cnstr Top

doEqualL i (Ntrl t) u = pure $ Ntrl (EqualL i t u)
doEqualL i Irrel u = inScope "bug" $ inScope "wrong type under equalL" $ fatal "Irrel"
doEqualL i (Cnstr t) u = doEqualR i t u

doEqualR i t (Ntrl u) = pure $ Ntrl (EqualR i t u)
doEqualR i t Irrel = inScope "bug" $ inScope "wrong type under equalR" $ fatal "Irrel"
doEqualR i t (Cnstr u) = doEqualU i t u

doEqualU i (Sort s) (Sort s') = pure $ Cnstr Top -- have suc s = i = suc s', and suc injective
doEqualU i (Pi s s' _ a b) (Pi l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i (Sigma s s' _ a b) (Sigma l l' var a' b') = case (s == s' && l == l') of
  False => pure $ Cnstr Bottom
  True => do
    eqLhs <- doEqual (cast s) a' a
    let x = Ntrl $ Var var FZ
    b <- subst b (add (Logging ann . NormalForm) (doCastL (weaken 1 a') (weaken 1 a) x) (pure . Ntrl . Var "" . FS))
    eqRhs <- doEqual (cast s') b b'
    pure $ Cnstr (Sigma Prop Prop "" eqLhs $ weaken 1 $ Cnstr (Pi s Prop var a' eqRhs))
doEqualU i t u = pure $ Ntrl (EqualU i t u) -- assumption: only logical values reach this far

doCastL (Ntrl a) b t = pure $ Ntrl (CastL a b t)
doCastL Irrel b t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastL (Cnstr a) b t = doCastR a b t

doCastR a (Ntrl b) t = pure $ Ntrl (CastR a b t)
doCastR a Irrel t = inScope "bug" $ inScope "wrong type for cast" $ fatal "Irrel"
doCastR a (Cnstr b) t = doCastU a b t

doCastU (Sort s) (Sort s') t = pure t
doCastU (Pi s s' _ a b) (Pi l l' var a' b') t = do
  let x' = Ntrl $ Var var FZ
  let x = doCastL (weaken 1 a') (weaken 1 a) x'
  b <- subst b (add (Logging ann . NormalForm) x (pure . Ntrl . Var "" . FS))
  b' <- subst b' (add (Logging ann . NormalForm) (pure x') (pure . Ntrl . Var "" . FS))
  fx <- doApp (weaken 1 t) !x
  cast <- doCastL b b' fx
  pure $ Cnstr (Lambda var cast)
doCastU (Sigma s s' _ a b) (Sigma l l' var a' b') t = do
  t1 <- doFst t
  t2 <- doSnd t
  t1' <- doCastL a a' t1
  b <- subst1 t1 b
  b' <- subst1 t1' b'
  t2' <- doCastL b b' t2
  pure $ Cnstr (Pair t1' t2')

doCastU Top Top t = pure Irrel
doCastU Bottom Bottom t = pure Irrel
doCastU a b t = pure $ Ntrl (CastU a b t)

-- Conversion ------------------------------------------------------------------

-- invariant: all definitions fully expanded
-- invariant: m |- t, u <- a : s
covering partial
convert : (t, u, a : NormalForm n) -> Sort -> Logging ann Bool
-- In sort Prop
convert Irrel Irrel a Prop = pure True
-- In unknown type in set
convert t u (Ntrl _) (Set k) = pure $ t == u
-- In type Set
convert (Cnstr (Sort s')) (Cnstr (Sort s'')) (Cnstr (Sort _)) (Set _) = pure $ s' == s''
convert (Cnstr (Pi s s' _ a b)) (Cnstr (Pi l l' _ a' b')) (Cnstr (Sort _)) (Set _) =
  pure $
  s == l && s' == l' && !(convert a a' (cast s) (suc s)) && !(convert b b' (cast s') (suc s'))
convert (Cnstr (Sigma s s' _ a b)) (Cnstr (Sigma l l' _ a' b')) (Cnstr (Sort _)) (Set _) =
  pure $
  s == l && s' == l' && !(convert a a' (cast s) (suc s)) && !(convert b b' (cast s') (suc s'))
convert (Cnstr Top) (Cnstr Top) (Cnstr (Sort _)) (Set _) = pure True
convert (Cnstr Bottom) (Cnstr Bottom) (Cnstr (Sort _)) (Set _) = pure True
convert (Ntrl t) u (Cnstr (Sort s)) (Set k) = pure $ Ntrl t == u
convert t (Ntrl u) (Cnstr (Sort s)) (Set k) = pure $ t == Ntrl u
convert t u (Cnstr (Sort s)) (Set k) = pure $ False
-- In type Pi
convert t u (Cnstr (Pi s s' var a b)) (Set k) = do
  t <- doApp (weaken 1 t) (Ntrl $ Var var FZ)
  u <- doApp (weaken 1 u) (Ntrl $ Var var FZ)
  convert t u b s'
-- In type Sigma
convert t u (Cnstr (Sigma s s' var a b)) (Set k) = do
  t1 <- doFst t
  u1 <- doFst t
  True <- convert t1 u1 a s
    | False => pure False
  b <- subst1 t1 b
  t2 <- doSnd t
  u2 <- doSnd t
  convert t2 u2 b s'
-- Default
convert t u a s =
  inScope "invalid conversion" $ fatal $
  fillSep {ann} [prettyPrec App t, prettyPrec App u, prettyPrec App a, prettyPrec App s]


-- Typing Contexts -------------------------------------------------------------

infix 5 ::<

data TyContext : Nat -> Nat -> Type where
  Nil   : TyContext 0 0
  (:<)  : TyContext m n -> NormalForm.Definition n -> TyContext m (S n)
  (::<) : TyContext m n -> (String, NormalForm n, Sort) -> TyContext (S m) (S n)

fromContext : Context n -> TyContext 0 n
fromContext []           = []
fromContext (ctx :< def) = fromContext ctx :< def

countVars : TyContext m n -> Fin (S n)
countVars []          = FZ
countVars (ctx :< y)  = weaken $ countVars ctx
countVars (ctx ::< y) = FS $ countVars ctx

index : TyContext m n -> Fin n -> (NormalForm n, NormalForm n, Sort)
index (ctx :< def)              FZ     = (weaken 1 def.tm, weaken 1 def.ty, def.sort)
index (ctx ::< (var, ty, Prop)) FZ     = (Irrel, weaken 1 ty, Prop)
index (ctx ::< (var, ty, s))    FZ     = (Ntrl $ Var var FZ, weaken 1 ty, s)
index (ctx :< _)                (FS i) = bimap (weaken 1) (mapFst (weaken 1)) $ index ctx i
index (ctx ::< _)               (FS i) = bimap (weaken 1) (mapFst (weaken 1)) $ index ctx i

covering partial
asSubst : (ctx : TyContext m n) -> Fin n -> Logging ann (NormalForm m)
asSubst (ctx :< def)             FZ     = subst def.tm (asSubst ctx)
asSubst (ctx ::< (var, _, Prop)) FZ     = pure Irrel
asSubst (ctx ::< (var, _, _))    FZ     = pure $ Ntrl (Var var FZ)
asSubst (ctx :< def)             (FS i) = asSubst ctx i
asSubst (ctx ::< _)              (FS i) = map (weaken 1) (asSubst ctx i)

-- Checking and Inference ------------------------------------------------------

covering partial
check : TyContext m n -> Term n -> NormalForm n -> Sort -> Logging ann (NormalForm n)
covering partial
infer : TyContext m n -> Term n -> Logging ann (NormalForm n, NormalForm n, Sort)
covering partial
inferType : {default typeMismatch tag : forall a . Bounds -> Doc ann -> Doc ann -> Logging ann a}
  -> TyContext m n -> Term n -> Logging ann (NormalForm n, Sort)

check ctx tm@(Lambda _ _ t) (Cnstr (Pi s s' var a b)) _ = do
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking under lambda with type" <++> pretty a) (fullBounds tm)
  inScope "check" $ trace $ map (\_ => "binding new variable to \{var}") (fullBounds tm)
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking for type" <++> pretty b) (fullBounds tm)
  t <- check (ctx ::< (var, a, s)) t b s'
  case s' of
    Prop => pure Irrel
    _ => pure (Cnstr $ Lambda var t)
check ctx tm@(Lambda _ _ _) ty _ = typeMismatch (fullBounds tm).bounds (pretty "pi") (pretty ty)
check ctx tm@(Pair _ t u) (Cnstr (Sigma s s' var a b)) _ = do
  inScope "check" $ trace $ map (\_ => "checking pair") (fullBounds tm)
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking first for type" <++> pretty a) (fullBounds tm)
  t <- check ctx t a s
  b <- subst1 t b
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking second for type" <++> pretty b) (fullBounds tm)
  u <- check ctx u b s'
  inScope "check" $ trace $ map (\_ => pretty {ann} "pair is well typed") (fullBounds tm)
  case (s, s') of
    (Prop, Prop) => pure Irrel
    _ => pure (Cnstr $ Pair t u)
check ctx tm@(Pair _ _ _) ty _ = typeMismatch (fullBounds tm).bounds (pretty "sigma") (pretty ty)
check ctx tm ty s = do
  let bounded = fullBounds tm
  inScope "check" $ trace $ map (\_ => "checking has fallen through") bounded
  (tm, ty', s') <- infer ctx tm
  inScope "check" $ trace $ map (\_ => pretty {ann} "inferred type is" <++> pretty ty') bounded
  let True = s == s'
    | False => sortMismatch bounded.bounds (pretty s) (pretty s')
  True <- convert !(subst ty $ asSubst ctx) !(subst ty' $ asSubst ctx) (cast s) (suc s)
    | False => typeMismatch bounded.bounds (pretty ty) (pretty ty')
  inScope "check" $ trace $ map (\_ => pretty {ann} "converted" <++> pretty ty' <+> softline <+> pretty "to" <++> pretty ty) bounded
  pure tm

infer ctx tm@(Var b i) = do
  inScope "infer" $ trace $ map (\_ => "encountered variable \{b.val}@\{show i}") (fullBounds tm)
  let (t, a, s) = index ctx i
  inScope "infer" $ trace $ map (\_ => pretty {ann} "variable has type" <++> pretty a) (fullBounds tm)
  pure (t, a, s)
infer ctx (Sort b s) = pure (cast s, cast (suc s), suc (suc s))
infer ctx tm@(Pi _ var a b) = do
  inScope "infer" $ trace $ map (\_ => "encountered Pi type") (fullBounds tm)
  (a, s) <- inferType ctx a
  inScope "infer" $ trace $ map (\_ => pretty {ann} "argument has type" <++> pretty a) (fullBounds tm)
  (b, s') <- inferType (ctx ::< (var, a, s)) b
  inScope "infer" $ trace $ map (\_ => pretty {ann} "result has type" <++> pretty b <+> comma <+> softline <+> pretty "so Pi type has type" <++> pretty (s ~> s')) (fullBounds tm)
  pure (Cnstr (Pi s s' var a b), cast (s ~> s'), suc (s ~> s'))
infer ctx tm@(Lambda _ _ _) = inScope "cannot infer type" $ fatal (fullBounds tm)
infer ctx tm@(App t u) = do
  inScope "infer" $ trace $ map (\_ => "encountered application") (fullBounds tm)
  let bounded = fullBounds t
  (t, ty@(Cnstr (Pi s s' _ a b)), _) <- infer ctx t
    | (_, t, _) => inScope "wrong type to apply" $ fatal (map (\_ => t) bounded)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "function has type" <++> pretty ty) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "checking argument has type" <++> pretty a) (fullBounds tm)
  u <- check ctx u a s
  inScope "infer" $ trace $ map (\_ => "argument is well typed") (fullBounds tm)
  res <- doApp t u
  ty <- subst1 u b
  inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty ty) (fullBounds tm)
  pure (res, ty, s')
infer ctx tm@(Sigma _ var a b) = do
  inScope "infer" $ trace $ map (\_ => "encountered Sigma type") (fullBounds tm)
  (a, s) <- inferType ctx a
  inScope "infer" $ trace $ map (\_ => pretty {ann} "first has type" <++> pretty a) (fullBounds tm)
  (b, s') <- inferType (ctx ::< (var, a, s)) b
  inScope "infer" $ trace $ map (\_ => pretty {ann} "second has type" <++> pretty b <+> comma <+> softline <+> pretty "so Sigma type has type" <++> pretty (lub s s')) (fullBounds tm)
  pure (Cnstr (Sigma s s' var a b), cast (lub s s'), suc (lub s s'))
infer ctx tm@(Pair _ _ _) = inScope "cannot infer type" $ fatal (fullBounds tm)
infer ctx tm@(Fst _ t) = do
  inScope "infer" $ trace $ map (\_ => "encountered first projection") (fullBounds tm)
  let bounded = fullBounds t
  (t, ty@(Cnstr (Sigma s s' _ a b)), _) <- infer ctx t
    | (_, t, _) => inScope "wrong type to fst" $ fatal (map (\_ => t) bounded)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "pair has type" <++> pretty ty) (fullBounds tm)
  res <- doFst t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty a) (fullBounds tm)
  pure (res, a, s)
infer ctx tm@(Snd _ t) = do
  inScope "infer" $ trace $ map (\_ => "encountered first projection") (fullBounds tm)
  let bounded = fullBounds t
  (t, ty@(Cnstr (Sigma s s' _ a b)), _) <- infer ctx t
    | (_, t, _) => inScope "wrong type to fst" $ fatal (map (\_ => t) bounded)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "pair has type" <++> pretty ty) (fullBounds tm)
  t1 <- doFst t
  res <- doSnd t
  b <- subst1 t1 b
  inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty b) (fullBounds tm)
  pure (res, b, s')
infer ctx (Top b) = pure $ (Cnstr Top, cast Prop, Set 0)
infer ctx (Point b) = pure $ (Irrel, Cnstr Top, Prop)
infer ctx (Bottom b) = pure $ (Cnstr Bottom, cast Prop, Set 0)
infer ctx tm@(Absurd b a t) = do
  inScope "infer" $ trace $ map (\_ => "encountered absurd") (fullBounds tm)
  (a, s) <- inferType ctx a
  inScope "infer" $ trace $ map (\_ => pretty {ann} "will fudge to type" <++> pretty a) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "checking for proof of false") (fullBounds tm)
  _ <- check ctx t (Cnstr Bottom) Prop
  inScope "infer" $ trace $ map (\_ => "proof of false is well typed") (fullBounds tm)
  pure (Ntrl Absurd, a, s)
infer ctx tm@(Equal b t u) = do
  inScope "infer" $ trace $ map (\_ => "encountered equal") (fullBounds tm)
  (t, a, s) <- infer ctx t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "lhs has type" <++> pretty a) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking rhs has same type") (fullBounds tm)
  u <- check ctx u a s
  inScope "infer" $ trace $ map (\_ => "rhs is well typed") (fullBounds tm)
  res <- doEqual a t u
  inScope "infer" $ trace $ map (\_ => pretty {ann} "equality computes to" <++> pretty res) (fullBounds tm)
  pure (res, cast Prop, Set 0)
infer ctx tm@(Refl b t) = do
  inScope "infer" $ trace $ map (\_ => "encountered refl") (fullBounds tm)
  (t, a, s) <- infer ctx t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "term has type" <++> pretty a) (fullBounds tm)
  ty <- doEqual a t t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "equality computes to" <++> pretty ty) (fullBounds tm)
  pure (Irrel, ty, Prop)
infer ctx tm@(Transp _ t b u t' e) = do
  inScope "infer" $ trace $ map (\_ => "encountered transp") (fullBounds tm)
  (t, a, s) <- infer ctx t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "start index has type" <++> pretty a) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking end index has same type") (fullBounds tm)
  t' <- check ctx t' a s
  inScope "infer" $ trace $ map (\_ => "end index is well typed") (fullBounds tm)
  let bounded = fullBounds b
  let ty = Cnstr (Pi s (Set 0) "_" a (cast Prop))
  inScope "infer" $ trace $ map (\_ => pretty {ann} "checkout output is in" <++> pretty ty) (fullBounds tm)
  b <- check ctx b ty s
  inScope "infer" $ trace $ map (\_ => pretty {ann} "output is well typed") (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking equality type") (fullBounds tm)
  eq <- doEqual a t t'
  _ <- check ctx e eq Prop
  inScope "infer" $ trace $ map (\_ => "equality is well typed") (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking transformed value") (fullBounds tm)
  inTy <- doApp b t
  _ <- check ctx u inTy Prop
  inScope "infer" $ trace $ map (\_ => "transformed value is well typed") (fullBounds tm)
  outTy <- doApp b t'
  inScope "infer" $ trace $ map (\_ => pretty {ann} "producing value of type" <++> pretty outTy) (fullBounds tm)
  pure (Irrel, outTy, Prop)
infer ctx tm@(Cast _ b e t) = do
  inScope "infer" $ trace $ map (\_ => "encountered cast") (fullBounds tm)
  (t, a, s) <- infer ctx t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "input has type" <++> pretty a) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking output has same sort") (fullBounds tm)
  b <- check ctx b (cast s) (suc s)
  inScope "infer" $ trace $ map (\_ => "output is well sorted") (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => "checking equality type") (fullBounds tm)
  eq <- doEqual (cast s) a b
  _ <- check ctx e eq Prop
  inScope "infer" $ trace $ map (\_ => "equality is well typed") (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "producing value of type" <++> pretty b) (fullBounds tm)
  res <- doCastL a b t
  pure (res, b, s)
infer ctx tm@(CastId _ t) = do
  inScope "infer" $ trace $ map (\_ => "encountered cast identity proof") (fullBounds tm)
  (t, a, s) <- infer ctx t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "term has type" <++> pretty a) (fullBounds tm)
  cast <- doCastL a a t
  eq <- doEqual a cast t
  inScope "infer" $ trace $ map (\_ => pretty {ann} "producing equality type" <++> pretty eq) (fullBounds tm)
  pure (Irrel, eq, Prop)

inferType ctx a = do
  (a, Cnstr (Sort s), _) <- infer ctx a
    | (_, ty, _) => tag (fullBounds a).bounds (pretty "sort") (pretty ty)
  pure (a, s)

-- Checking Definitions and Blocks ---------------------------------------------

covering partial
checkDef : Context n -> Term.Definition n -> Logging ann (NormalForm.Definition n)
checkDef ctx def = do
  let tyBounds = (fullBounds def.ty).bounds
  let tmBounds = (fullBounds def.tm).bounds
  (ty, sort) <-
    inferType {tag = \bounds, lhs, rhs => inScope "invalid declaration" $ mismatch bounds lhs rhs}
      (fromContext ctx) def.ty
  inScope "check" $ debug (pretty {ann} "\{def.name} has type" <++> pretty ty)
  tm <- check (fromContext ctx) def.tm ty sort
  inScope "check" $ debug (pretty {ann} "\{def.name} is well typed with value" <++> pretty tm)
  pure $ MkDefinition {name = def.name, ty, tm, sort}

covering partial
export
checkBlock : Block n -> Logging ann (Context n)
checkBlock [] = pure []
checkBlock (blk :< def) = do
  ctx <- checkBlock blk
  def <- checkDef ctx def
  pure (ctx :< def)