module Quox.Reduce

import Quox.No
import Quox.Syntax
import Quox.Definition
import Quox.Displace
import Quox.Typing.Context
import Quox.Typing.Error
import Data.SnocVect
import Data.Maybe
import Data.List
import Control.Eff

%default total


-- [fixme] rename this to Whnf and the interface to CanWhnf
public export
Whnf : List (Type -> Type)
Whnf = [NameGen, Except Error]

export
runWhnfWith : NameSuf -> Eff Whnf a -> (Either Error a, NameSuf)
runWhnfWith suf act = extract $ runStateAt GEN suf $ runExcept act

export
runWhnf : Eff Whnf a -> Either Error a
runWhnf = fst . runWhnfWith 0


public export
0 RedexTest : TermLike -> Type
RedexTest tm = {d, n : Nat} -> Definitions -> tm d n -> Bool

public export
interface CanWhnf (0 tm : TermLike) (0 isRedex : RedexTest tm) | tm
where
  whnf : {d, n : Nat} -> (defs : Definitions) ->
         (ctx : WhnfContext d n) ->
         tm d n -> Eff Whnf (Subset (tm d n) (No . isRedex defs))

public export %inline
whnf0 : {d, n : Nat} -> {0 isRedex : RedexTest tm} -> CanWhnf tm isRedex =>
        (defs : Definitions) -> WhnfContext d n -> tm d n -> Eff Whnf (tm d n)
whnf0 defs ctx t = fst <$> whnf defs ctx t

public export
0 IsRedex, NotRedex : {isRedex : RedexTest tm} -> CanWhnf tm isRedex =>
                      Definitions -> Pred (tm d n)
IsRedex  defs = So . isRedex defs
NotRedex defs = No . isRedex defs

public export
0 NonRedex : (tm : TermLike) -> {isRedex : RedexTest tm} ->
             CanWhnf tm isRedex => (d, n : Nat) -> (defs : Definitions) -> Type
NonRedex tm d n defs = Subset (tm d n) (NotRedex defs)

public export %inline
nred : {0 isRedex : RedexTest tm} -> (0 _ : CanWhnf tm isRedex) =>
       (t : tm d n) -> (0 nr : NotRedex defs t) => NonRedex tm d n defs
nred t = Element t nr


public export %inline
isLamHead : Elim {} -> Bool
isLamHead (Ann (Lam {}) (Pi {}) _) = True
isLamHead (Coe {})                 = True
isLamHead _                        = False

public export %inline
isDLamHead : Elim {} -> Bool
isDLamHead (Ann (DLam {}) (Eq {}) _) = True
isDLamHead (Coe {})                  = True
isDLamHead _                         = False

public export %inline
isPairHead : Elim {} -> Bool
isPairHead (Ann (Pair {}) (Sig {}) _) = True
isPairHead (Coe {})                   = True
isPairHead _                          = False

public export %inline
isTagHead : Elim {} -> Bool
isTagHead (Ann (Tag {}) (Enum {}) _) = True
isTagHead (Coe {})                   = True
isTagHead _                          = False

public export %inline
isNatHead : Elim {} -> Bool
isNatHead (Ann (Zero {}) (Nat {}) _) = True
isNatHead (Ann (Succ {}) (Nat {}) _) = True
isNatHead (Coe {})                   = True
isNatHead _                          = False

public export %inline
isBoxHead : Elim {} -> Bool
isBoxHead (Ann (Box {}) (BOX {}) _) = True
isBoxHead (Coe {})                  = True
isBoxHead _                         = False

public export %inline
isE : Term {} -> Bool
isE (E {}) = True
isE _      = False

public export %inline
isAnn : Elim {} -> Bool
isAnn (Ann {}) = True
isAnn _        = False

||| true if a term is syntactically a type.
public export %inline
isTyCon : Term {} -> Bool
isTyCon (TYPE  {}) = True
isTyCon (Pi    {}) = True
isTyCon (Lam   {}) = False
isTyCon (Sig   {}) = True
isTyCon (Pair  {}) = False
isTyCon (Enum  {}) = True
isTyCon (Tag   {}) = False
isTyCon (Eq    {}) = True
isTyCon (DLam  {}) = False
isTyCon (Nat   {}) = True
isTyCon (Zero  {}) = False
isTyCon (Succ  {}) = False
isTyCon (BOX   {}) = True
isTyCon (Box   {}) = False
isTyCon (E     {}) = False
isTyCon (CloT  {}) = False
isTyCon (DCloT {}) = False

||| true if a term is syntactically a type, or a neutral.
public export %inline
isTyConE : Term {} -> Bool
isTyConE s = isTyCon s || isE s

||| true if a term is syntactically a type.
public export %inline
isAnnTyCon : Elim {} -> Bool
isAnnTyCon (Ann ty (TYPE {}) _) = isTyCon ty
isAnnTyCon _                    = False

public export %inline
isK : Dim d -> Bool
isK (K {}) = True
isK _      = False


mutual
  public export
  isRedexE : RedexTest Elim
  isRedexE defs (F {x, _}) {d, n} =
    isJust $ lookupElim x defs {d, n}
  isRedexE _ (B {}) = False
  isRedexE defs (App {fun, _}) =
    isRedexE defs fun || isLamHead fun
  isRedexE defs (CasePair {pair, _}) =
    isRedexE defs pair || isPairHead pair
  isRedexE defs (CaseEnum {tag, _}) =
    isRedexE defs tag || isTagHead tag
  isRedexE defs (CaseNat {nat, _}) =
    isRedexE defs nat || isNatHead nat
  isRedexE defs (CaseBox {box, _}) =
    isRedexE defs box || isBoxHead box
  isRedexE defs (DApp {fun, arg, _}) =
    isRedexE defs fun || isDLamHead fun || isK arg
  isRedexE defs (Ann {tm, ty, _}) =
    isE tm || isRedexT defs tm || isRedexT defs ty
  isRedexE defs (Coe {val, _}) =
    isRedexT defs val || not (isE val)
  isRedexE defs (Comp {ty, r, _}) =
    isRedexT defs ty || isK r
  isRedexE defs (TypeCase {ty, ret, _}) =
    isRedexE defs ty || isRedexT defs ret || isAnnTyCon ty
  isRedexE _ (CloE  {}) = True
  isRedexE _ (DCloE {}) = True

  public export
  isRedexT : RedexTest Term
  isRedexT _    (CloT  {}) = True
  isRedexT _    (DCloT {}) = True
  isRedexT defs (E {e, _}) = isAnn e || isRedexE defs e
  isRedexT _    _          = False


public export
tycaseRhs : (k : TyConKind) -> TypeCaseArms d n ->
            Maybe (ScopeTermN (arity k) d n)
tycaseRhs k arms = lookupPrecise k arms

public export
tycaseRhsDef : Term d n -> (k : TyConKind) -> TypeCaseArms d n ->
               ScopeTermN (arity k) d n
tycaseRhsDef def k arms = fromMaybe (SN def) $ tycaseRhs k arms

public export
tycaseRhs0 : (k : TyConKind) -> TypeCaseArms d n ->
             (0 eq : arity k = 0) => Maybe (Term d n)
tycaseRhs0 k arms {eq} with (tycaseRhs k arms) | (arity k)
  tycaseRhs0 k arms {eq = Refl} | res | 0 = map (.term) res

public export
tycaseRhsDef0 : Term d n -> (k : TyConKind) -> TypeCaseArms d n ->
                (0 eq : arity k = 0) => Term d n
tycaseRhsDef0 def k arms = fromMaybe def $ tycaseRhs0 k arms



private
weakDS : (by : Nat) -> DScopeTerm d n -> DScopeTerm d (by + n)
weakDS by (S names (Y body)) = S names $ Y $ weakT by body
weakDS by (S names (N body)) = S names $ N $ weakT by body

private
dweakS : (by : Nat) -> ScopeTerm d n -> ScopeTerm (by + d) n
dweakS by (S names (Y body)) = S names $ Y $ dweakT by body
dweakS by (S names (N body)) = S names $ N $ dweakT by body

private
coeScoped : {s : Nat} -> DScopeTerm d n -> Dim d -> Dim d -> Loc ->
            ScopeTermN s d n -> ScopeTermN s d n
coeScoped ty p q loc (S names (Y body)) =
  S names $ Y $ E $ Coe (weakDS s ty) p q body loc
coeScoped ty p q loc (S names (N body)) =
  S names $ N $ E $ Coe ty p q body loc


export covering
CanWhnf Term Reduce.isRedexT

export covering
CanWhnf Elim Reduce.isRedexE

parameters {d, n : Nat} (defs : Definitions) (ctx : WhnfContext d n)
  ||| performs the minimum work required to recompute the type of an elim.
  |||
  ||| ⚠ **assumes the elim is already typechecked.** ⚠
  export covering
  computeElimType : (e : Elim d n) -> (0 ne : No (isRedexE defs e)) =>
                    Eff Whnf (Term d n)
  computeElimType (F {x, u, loc}) = do
    let Just def = lookup x defs | Nothing => throw $ NotInScope loc x
    pure $ displace u def.type
  computeElimType (B {i, _}) = pure $ ctx.tctx !! i
  computeElimType (App {fun = f, arg = s, loc}) {ne} = do
    Pi {arg, res, _} <- whnf0 defs ctx =<< computeElimType f {ne = noOr1 ne}
      | t => throw $ ExpectedPi loc ctx.names t
    pure $ sub1 res $ Ann s arg loc
  computeElimType (CasePair {pair, ret, _}) = pure $ sub1 ret pair
  computeElimType (CaseEnum {tag, ret, _}) = pure $ sub1 ret tag
  computeElimType (CaseNat {nat, ret, _}) = pure $ sub1 ret nat
  computeElimType (CaseBox {box, ret, _}) = pure $ sub1 ret box
  computeElimType (DApp {fun = f, arg = p, loc}) {ne} = do
    Eq {ty, _} <- whnf0 defs ctx =<< computeElimType f {ne = noOr1 ne}
      | t => throw $ ExpectedEq loc ctx.names t
    pure $ dsub1 ty p
  computeElimType (Ann {ty, _}) = pure ty
  computeElimType (Coe {ty, q, _}) = pure $ dsub1 ty q
  computeElimType (Comp {ty, _}) = pure ty
  computeElimType (TypeCase {ret, _}) = pure ret

parameters {d, n : Nat} (defs : Definitions) (ctx : WhnfContext (S d) n)
  ||| for π.(x : A) → B, returns (A, B);
  ||| for an elim returns a pair of type-cases that will reduce to that;
  ||| for other intro forms error
  private covering
  tycasePi : (t : Term (S d) n) -> (0 tnf : No (isRedexT defs t)) =>
             Eff Whnf (Term (S d) n, ScopeTerm (S d) n)
  tycasePi (Pi {arg, res, _}) = pure (arg, res)
  tycasePi (E e) {tnf} = do
    ty <- computeElimType defs ctx e @{noOr2 tnf}
    let loc  = e.loc
        narg = mnb "Arg"; nret = mnb "Ret"
        arg  = E $ typeCase1Y e ty KPi [< !narg, !nret] (BVT 1 loc) loc
        res' = typeCase1Y e (Arr Zero arg ty loc) KPi [< !narg, !nret]
                 (BVT 0 loc) loc
        res  = SY [< !narg] $ E $ App (weakE 1 res') (BVT 0 loc) loc
    pure (arg, res)
  tycasePi t = throw $ ExpectedPi t.loc ctx.names t

  ||| for (x : A) × B, returns (A, B);
  ||| for an elim returns a pair of type-cases that will reduce to that;
  ||| for other intro forms error
  private covering
  tycaseSig : (t : Term (S d) n) -> (0 tnf : No (isRedexT defs t)) =>
              Eff Whnf (Term (S d) n, ScopeTerm (S d) n)
  tycaseSig (Sig {fst, snd, _}) = pure (fst, snd)
  tycaseSig (E e) {tnf} = do
    ty <- computeElimType defs ctx e @{noOr2 tnf}
    let loc  = e.loc
        nfst = mnb "Fst"; nsnd = mnb "Snd"
        fst  = E $ typeCase1Y e ty KSig [< !nfst, !nsnd] (BVT 1 loc) loc
        snd' = typeCase1Y e (Arr Zero fst ty loc) KSig [< !nfst, !nsnd]
                 (BVT 0 loc) loc
        snd  = SY [< !nfst] $ E $ App (weakE 1 snd') (BVT 0 loc) loc
    pure (fst, snd)
  tycaseSig t = throw $ ExpectedSig t.loc ctx.names t

  ||| for [π. A], returns A;
  ||| for an elim returns a type-case that will reduce to that;
  ||| for other intro forms error
  private covering
  tycaseBOX : (t : Term (S d) n) -> (0 tnf : No (isRedexT defs t)) =>
              Eff Whnf (Term (S d) n)
  tycaseBOX (BOX {ty, _}) = pure ty
  tycaseBOX (E e) {tnf} = do
    ty <- computeElimType defs ctx e @{noOr2 tnf}
    pure $ E $ typeCase1Y e ty KBOX [< !(mnb "Ty")] (BVT 0 e.loc) e.loc
  tycaseBOX t = throw $ ExpectedBOX t.loc ctx.names t

  ||| for Eq [i ⇒ A] l r, returns (A‹0/i›, A‹1/i›, A, l, r);
  ||| for an elim returns five type-cases that will reduce to that;
  ||| for other intro forms error
  private covering
  tycaseEq : (t : Term (S d) n) -> (0 tnf : No (isRedexT defs t)) =>
             Eff Whnf (Term (S d) n, Term (S d) n, DScopeTerm (S d) n,
                       Term (S d) n, Term (S d) n)
  tycaseEq (Eq {ty, l, r, _}) = pure (ty.zero, ty.one, ty, l, r)
  tycaseEq (E e) {tnf} = do
    ty <- computeElimType defs ctx e @{noOr2 tnf}
    let loc   = e.loc
        names = traverse' (\x => mnb x) [< "A0", "A1", "A", "L", "R"]
        a0    = E $ typeCase1Y e ty KEq !names (BVT 4 loc) loc
        a1    = E $ typeCase1Y e ty KEq !names (BVT 3 loc) loc
        a'    = typeCase1Y e (Eq0 ty a0 a1 loc) KEq !names (BVT 2 loc) loc
        a     = SY [< !(mnb "i")] $ E $ DApp (dweakE 1 a') (B VZ loc) loc
        l     = E $ typeCase1Y e a0 KEq !names (BVT 1 loc) loc
        r     = E $ typeCase1Y e a1 KEq !names (BVT 0 loc) loc
    pure (a0, a1, a, l, r)
  tycaseEq t = throw $ ExpectedEq t.loc ctx.names t

-- new block because the functions below might pass a different ctx
-- into the ones above
parameters {d, n : Nat} (defs : Definitions) (ctx : WhnfContext d n)
  ||| reduce a function application `App (Coe ty p q val) s loc`
  private covering
  piCoe : (ty : DScopeTerm d n) -> (p, q : Dim d) ->
          (val, s : Term d n) -> Loc ->
          Eff Whnf (Subset (Elim d n) (No . isRedexE defs))
  piCoe sty@(S [< i] ty) p q val s loc = do
    -- (coe [i ⇒ π.(x : A) → B] @p @q t) s ⇝
    -- coe [i ⇒ B[𝒔‹i›/x] @p @q ((t ∷ (π.(x : A) → B)‹p/i›) 𝒔‹p›)
    --   where 𝒔‹j› ≔ coe [i ⇒ A] @q @j s
    --
    -- type-case is used to expose A,B if the type is neutral
    let ctx1 = extendDim i ctx
    Element ty tynf <- whnf defs ctx1 ty.term
    (arg, res) <- tycasePi defs ctx1 ty
    let s0   = CoeT i arg q p s s.loc
        body = E $ App (Ann val (ty // one p) val.loc) (E s0) loc
        s1   = CoeT i (arg // (BV 0 i.loc ::: shift 2)) (weakD 1 q) (BV 0 i.loc)
                    (s // shift 1) s.loc
    whnf defs ctx $ CoeT i (sub1 res s1) p q body loc

  ||| reduce a pair elimination `CasePair pi (Coe ty p q val) ret body loc`
  private covering
  sigCoe : (qty : Qty) ->
           (ty : DScopeTerm d n) -> (p, q : Dim d) -> (val : Term d n) ->
           (ret : ScopeTerm d n) -> (body : ScopeTermN 2 d n) -> Loc ->
           Eff Whnf (Subset (Elim d n) (No . isRedexE defs))
  sigCoe qty sty@(S [< i] ty) p q val ret body loc = do
    -- caseπ (coe [i ⇒ (x : A) × B] @p @q s) return z ⇒ C of { (a, b) ⇒ e }
    -- ⇝
    -- caseπ s ∷ ((x : A) × B)‹p/i› return z ⇒ C
    -- of { (a, b) ⇒
    --   e[(coe [i ⇒ A] @p @q a)/a,
    --     (coe [i ⇒ B[(coe [j ⇒ A‹j/i›] @p @i a)/x]] @p @q b)/b] }
    --
    -- type-case is used to expose A,B if the type is neutral
    let ctx1 = extendDim i ctx
    Element ty tynf <- whnf defs ctx1 ty.term
    (tfst, tsnd) <- tycaseSig defs ctx1 ty
    let [< x, y] = body.names
        a' = CoeT i (weakT 2 tfst) p q (BVT 1 noLoc) x.loc
        tsnd' = tsnd.term //
          (CoeT i (weakT 2 $ tfst // (B VZ noLoc ::: shift 2))
            (weakD 1 p) (B VZ noLoc) (BVT 1 noLoc) y.loc ::: shift 2)
        b' = CoeT i tsnd' p q (BVT 0 noLoc) y.loc
    whnf defs ctx $ CasePair qty (Ann val (ty // one p) val.loc) ret
      (ST body.names $ body.term // (a' ::: b' ::: shift 2)) loc

  ||| reduce a dimension application `DApp (Coe ty p q val) r loc`
  private covering
  eqCoe : (ty : DScopeTerm d n) -> (p, q : Dim d) -> (val : Term d n) ->
          (r : Dim d) -> Loc ->
          Eff Whnf (Subset (Elim d n) (No . isRedexE defs))
  eqCoe sty@(S [< j] ty) p q val r loc = do
    -- (coe [j ⇒ Eq [i ⇒ A] L R] @p @q eq) @r
    -- ⇝
    -- comp [j ⇒ A‹r/i›] @p @q (eq ∷ (Eq [i ⇒ A] L R)‹p/j›)
    --   @r { 0 j ⇒ L; 1 j ⇒ R }
    let ctx1 = extendDim j ctx
    Element ty tynf <- whnf defs ctx1 ty.term
    (a0, a1, a, s, t) <- tycaseEq defs ctx1 ty
    let a'   = dsub1 a (weakD 1 r)
        val' = E $ DApp (Ann val (ty // one p) val.loc) r loc
    whnf defs ctx $ CompH j a' p q val' r j s j t loc

  ||| reduce a pair elimination `CaseBox pi (Coe ty p q val) ret body`
  private covering
  boxCoe : (qty : Qty) ->
           (ty : DScopeTerm d n) -> (p, q : Dim d) -> (val : Term d n) ->
           (ret : ScopeTerm d n) -> (body : ScopeTerm d n) -> Loc ->
           Eff Whnf (Subset (Elim d n) (No . isRedexE defs))
  boxCoe qty sty@(S [< i] ty) p q val ret body loc = do
    -- caseπ (coe [i ⇒ [ρ. A]] @p @q s) return z ⇒ C of { [a] ⇒ e }
    -- ⇝
    -- caseπ s ∷ [ρ. A]‹p/i› return z ⇒ C
    -- of { [a] ⇒ e[(coe [i ⇒ A] p q a)/a] }
    let ctx1 = extendDim i ctx
    Element ty tynf <- whnf defs ctx1 ty.term
    ta <- tycaseBOX defs ctx1 ty
    let a' = CoeT i (weakT 1 ta) p q (BVT 0 noLoc) body.name.loc
    whnf defs ctx $ CaseBox qty (Ann val (ty // one p) val.loc) ret
      (ST body.names $ body.term // (a' ::: shift 1)) loc


||| reduce a type-case applied to a type constructor
private covering
reduceTypeCase : {d, n : Nat} -> (defs : Definitions) -> WhnfContext d n ->
                 (ty : Term d n) -> (u : Universe) -> (ret : Term d n) ->
                 (arms : TypeCaseArms d n) -> (def : Term d n) ->
                 (0 _ : So (isTyCon ty)) => Loc ->
                 Eff Whnf (Subset (Elim d n) (No . isRedexE defs))
reduceTypeCase defs ctx ty u ret arms def loc = case ty of
  -- (type-case ★ᵢ ∷ _ return Q of { ★ ⇒ s; ⋯ }) ⇝ s ∷ Q
  TYPE {} =>
    whnf defs ctx $ Ann (tycaseRhsDef0 def KTYPE arms) ret loc

  -- (type-case π.(x : A) → B ∷ ★ᵢ return Q of { (a → b) ⇒ s; ⋯ }) ⇝
  -- s[(A ∷ ★ᵢ)/a, ((λ x ⇒ B) ∷ 0.A → ★ᵢ)/b] ∷ Q
  Pi {arg, res, loc = piLoc, _} =>
    let arg' = Ann arg (TYPE u noLoc) arg.loc
        res' = Ann (Lam res res.loc)
                   (Arr Zero arg (TYPE u noLoc) arg.loc) res.loc
    in
    whnf defs ctx $
      Ann (subN (tycaseRhsDef def KPi arms) [< arg', res']) ret loc

  -- (type-case (x : A) × B ∷ ★ᵢ return Q of { (a × b) ⇒ s; ⋯ }) ⇝
  -- s[(A ∷ ★ᵢ)/a, ((λ x ⇒ B) ∷ 0.A → ★ᵢ)/b] ∷ Q
  Sig {fst, snd, loc = sigLoc, _} =>
    let fst' = Ann fst (TYPE u noLoc) fst.loc
        snd' = Ann (Lam snd snd.loc)
                   (Arr Zero fst (TYPE u noLoc) fst.loc) snd.loc
    in
    whnf defs ctx $
      Ann (subN (tycaseRhsDef def KSig arms) [< fst', snd']) ret loc

  -- (type-case {⋯} ∷ _ return Q of { {} ⇒ s; ⋯ }) ⇝ s ∷ Q
  Enum {} =>
    whnf defs ctx $ Ann (tycaseRhsDef0 def KEnum arms) ret loc

  -- (type-case Eq [i ⇒ A] L R ∷ ★ᵢ return Q
  --  of { Eq a₀ a₁ a l r ⇒ s; ⋯ }) ⇝
  -- s[(A‹0/i› ∷ ★ᵢ)/a₀, (A‹1/i› ∷ ★ᵢ)/a₁,
  --   ((δ i ⇒ A) ∷ Eq [★ᵢ] A‹0/i› A‹1/i›)/a,
  --   (L ∷ A‹0/i›)/l, (R ∷ A‹1/i›)/r] ∷ Q
  Eq {ty = a, l, r, loc = eqLoc, _} =>
    let a0 = a.zero; a1 = a.one in
    whnf defs ctx $ Ann
      (subN (tycaseRhsDef def KEq arms)
         [< Ann a0 (TYPE u noLoc) a.loc, Ann a1 (TYPE u noLoc) a.loc,
            Ann (DLam a a.loc) (Eq0 (TYPE u noLoc) a0 a1 a.loc) a.loc,
            Ann l a0 l.loc, Ann r a1 r.loc])
      ret loc

  -- (type-case ℕ ∷ _ return Q of { ℕ ⇒ s; ⋯ }) ⇝ s ∷ Q
  Nat {} =>
    whnf defs ctx $ Ann (tycaseRhsDef0 def KNat arms) ret loc

  -- (type-case [π.A] ∷ ★ᵢ return Q of { [a] ⇒ s; ⋯ }) ⇝ s[(A ∷ ★ᵢ)/a] ∷ Q
  BOX {ty = a, loc = boxLoc, _} =>
    whnf defs ctx $ Ann
      (sub1 (tycaseRhsDef def KBOX arms) (Ann a (TYPE u noLoc) a.loc))
      ret loc


||| pushes a coercion inside a whnf-ed term
private covering
pushCoe : {d, n : Nat} -> (defs : Definitions) -> WhnfContext d n ->
          BindName ->
          (ty : Term (S d) n) -> (0 tynf : No (isRedexT defs ty)) =>
          Dim d -> Dim d ->
          (s : Term d n) -> (0 snf : No (isRedexT defs s)) => Loc ->
          Eff Whnf (NonRedex Elim d n defs)
pushCoe defs ctx i ty p q s loc =
  if p == q then whnf defs ctx $ Ann s (ty // one q) loc else
  case s of
    -- (coe [_ ⇒ ★ᵢ] @_ @_ ty) ⇝ (ty ∷ ★ᵢ)
    TYPE {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    Pi   {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    Sig  {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    Enum {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    Eq   {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    Nat  {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc
    BOX  {} => pure $ nred $ Ann s (TYPE !(unwrapTYPE ty) ty.loc) loc

    -- just η expand it. then whnf for App will handle it later
    -- this is how @xtt does it
    --
    -- (coe [i ⇒ A] @p @q (λ x ⇒ s)) ⇝
    -- (λ y ⇒ (coe [i ⇒ A] @p @q (λ x ⇒ s)) y) ∷ A‹q/i› ⇝ ⋯
    lam@(Lam {body, _}) => do
      let lam'  = CoeT i ty p q lam loc
          term' = LamY !(fresh body.name)
                    (E $ App (weakE 1 lam') (BVT 0 noLoc) loc) loc
          type' = ty // one q
      whnf defs ctx $ Ann term' type' loc

    -- (coe [i ⇒ (x : A) × B] @p @q (s, t)) ⇝
    -- (coe [i ⇒ A] @p @q s,
    --  coe [i ⇒ B[(coe [j ⇒ A‹j/i›] @p @i s)/x]] @p @q t)
    --   ∷ (x : A‹q/i›) × B‹q/i›
    --
    -- can't use η here because... it doesn't exist
    Pair {fst, snd, loc = pairLoc} => do
      let Sig {fst = tfst, snd = tsnd, loc = sigLoc} = ty
        | _ => throw $ ExpectedSig ty.loc (extendDim i ctx.names) ty
      let fst'  = E $ CoeT i tfst p q fst fst.loc
          tfst' = tfst // (B VZ noLoc ::: shift 2)
          tsnd' = sub1 tsnd $
                  CoeT !(fresh i) tfst' (weakD 1 p) (B VZ noLoc)
                    (dweakT 1 fst) fst.loc
          snd'  = E $ CoeT i tsnd' p q snd snd.loc
      pure $
        Element (Ann (Pair fst' snd' pairLoc)
                     (Sig (tfst // one q) (tsnd // one q) sigLoc) loc) Ah

    -- η expand, like for Lam
    --
    -- (coe [i ⇒ A] @p @q (δ j ⇒ s)) ⇝
    -- (δ k ⇒ (coe [i ⇒ A] @p @q (δ j ⇒ s)) @k) ∷ A‹q/i› ⇝ ⋯
    dlam@(DLam {body, _}) => do
      let dlam' = CoeT i ty p q dlam loc
          term' = DLamY !(mnb "j")
                    (E $ DApp (dweakE 1 dlam') (B VZ noLoc) loc) loc
          type' = ty // one q
      whnf defs ctx $ Ann term' type' loc

    -- (coe [_ ⇒ {⋯}] @_ @_ t) ⇝ (t ∷ {⋯})
    Tag {tag, loc = tagLoc} => do
      let Enum {cases, loc = enumLoc} = ty
        | _ => throw $ ExpectedEnum ty.loc (extendDim i ctx.names) ty
      pure $ Element (Ann (Tag tag tagLoc) (Enum cases enumLoc) loc) Ah

    -- (coe [_ ⇒ ℕ] @_ @_ n) ⇝ (n ∷ ℕ)
    Zero {loc = zeroLoc} => do
      pure $ Element (Ann (Zero zeroLoc) (Nat ty.loc) loc) Ah
    Succ {p = pred, loc = succLoc} => do
      pure $ Element (Ann (Succ pred succLoc) (Nat ty.loc) loc) Ah

    -- (coe [i ⇒ [π.A]] @p @q [s]) ⇝
    -- [coe [i ⇒ A] @p @q s] ∷ [π. A‹q/i›]
    Box {val, loc = boxLoc} => do
      let BOX {qty, ty = a, loc = tyLoc} = ty
        | _ => throw $ ExpectedBOX ty.loc (extendDim i ctx.names) ty
      pure $ Element
        (Ann (Box (E $ CoeT i a p q val val.loc) boxLoc)
             (BOX qty (a // one q) tyLoc) loc)
        Ah

    E e => pure $ Element (CoeT i ty p q (E e) e.loc) (snf `orNo` Ah)
where
  unwrapTYPE : Term (S d) n -> Eff Whnf Universe
  unwrapTYPE (TYPE {l, _}) = pure l
  unwrapTYPE ty = throw $ ExpectedTYPE ty.loc (extendDim i ctx.names) ty


export covering
CanWhnf Elim Reduce.isRedexE where
  whnf defs ctx (F x u loc) with (lookupElim x defs) proof eq
    _ | Just y  = whnf defs ctx $ setLoc loc $ displace u y
    _ | Nothing = pure $ Element (F x u loc) $ rewrite eq in Ah

  whnf _ _ (B i loc) = pure $ nred $ B i loc

  -- ((λ x ⇒ t) ∷ (π.x : A) → B) s ⇝ t[s∷A/x] ∷ B[s∷A/x]
  whnf defs ctx (App f s appLoc) = do
    Element f fnf <- whnf defs ctx f
    case nchoose $ isLamHead f of
      Left _ => case f of
        Ann (Lam {body, _}) (Pi {arg, res, _}) floc =>
          let s = Ann s arg s.loc in
          whnf defs ctx $ Ann (sub1 body s) (sub1 res s) appLoc
        Coe ty p q val _ => piCoe defs ctx ty p q val s appLoc
      Right nlh => pure $ Element (App f s appLoc) $ fnf `orNo` nlh

  -- case (s, t) ∷ (x : A) × B return p ⇒ C of { (a, b) ⇒ u } ⇝
  -- u[s∷A/a, t∷B[s∷A/x]] ∷ C[(s, t)∷((x : A) × B)/p]
  whnf defs ctx (CasePair pi pair ret body caseLoc) = do
    Element pair pairnf <- whnf defs ctx pair
    case nchoose $ isPairHead pair of
      Left _ => case pair of
        Ann (Pair {fst, snd, _}) (Sig {fst = tfst, snd = tsnd, _}) pairLoc =>
          let fst = Ann fst tfst fst.loc
              snd = Ann snd (sub1 tsnd fst) snd.loc
          in
          whnf defs ctx $ Ann (subN body [< fst, snd]) (sub1 ret pair) caseLoc
        Coe ty p q val _ => do
          sigCoe defs ctx pi ty p q val ret body caseLoc
      Right np =>
        pure $ Element (CasePair pi pair ret body caseLoc) $ pairnf `orNo` np

  -- case 'a ∷ {a,…} return p ⇒ C of { 'a ⇒ u } ⇝
  -- u ∷ C['a∷{a,…}/p]
  whnf defs ctx (CaseEnum pi tag ret arms caseLoc) = do
    Element tag tagnf <- whnf defs ctx tag
    case nchoose $ isTagHead tag of
      Left _ => case tag of
        Ann (Tag t _) (Enum ts _) _ =>
          let ty = sub1 ret tag in
          case lookup t arms of
            Just arm => whnf defs ctx $ Ann arm ty arm.loc
            Nothing  => throw $ MissingEnumArm caseLoc t (keys arms)
        Coe ty p q val _ =>
          -- there is nowhere an equality can be hiding inside an enum type
          whnf defs ctx $
            CaseEnum pi (Ann val (dsub1 ty q) val.loc) ret arms caseLoc
      Right nt =>
        pure $ Element (CaseEnum pi tag ret arms caseLoc) $ tagnf `orNo` nt

  -- case zero ∷ ℕ return p ⇒ C of { zero ⇒ u; … } ⇝
  --   u ∷ C[zero∷ℕ/p]
  --
  -- case succ n ∷ ℕ return p ⇒ C of { succ n', π.ih ⇒ u; … } ⇝
  --   u[n∷ℕ/n', (case n ∷ ℕ ⋯)/ih] ∷ C[succ n ∷ ℕ/p]
  whnf defs ctx (CaseNat pi piIH nat ret zer suc caseLoc) = do
    Element nat natnf <- whnf defs ctx nat
    case nchoose $ isNatHead nat of
      Left _ =>
        let ty = sub1 ret nat in
        case nat of
          Ann (Zero _) (Nat _) _ =>
            whnf defs ctx $ Ann zer ty zer.loc
          Ann (Succ n succLoc) (Nat natLoc) _ =>
            let nn = Ann n (Nat natLoc) succLoc
                tm = subN suc [< nn, CaseNat pi piIH nn ret zer suc caseLoc]
            in
            whnf defs ctx $ Ann tm ty caseLoc
          Coe ty p q val _ =>
            -- same deal as Enum
            whnf defs ctx $
              CaseNat pi piIH (Ann val (dsub1 ty q) val.loc) ret zer suc caseLoc
      Right nn => pure $
        Element (CaseNat pi piIH nat ret zer suc caseLoc) $ natnf `orNo` nn

  -- case [t] ∷ [π.A] return p ⇒ C of { [x] ⇒ u } ⇝
  --   u[t∷A/x] ∷ C[[t] ∷ [π.A]/p]
  whnf defs ctx (CaseBox pi box ret body caseLoc) = do
    Element box boxnf <- whnf defs ctx box
    case nchoose $ isBoxHead box of
      Left _ => case box of
        Ann (Box val boxLoc) (BOX q bty tyLoc) _ =>
          let ty = sub1 ret box in
          whnf defs ctx $ Ann (sub1 body (Ann val bty val.loc)) ty caseLoc
        Coe ty p q val _ =>
          boxCoe defs ctx pi ty p q val ret body caseLoc
      Right nb =>
        pure $ Element (CaseBox pi box ret body caseLoc) $ boxnf `orNo` nb

  -- e : Eq (𝑗 ⇒ A) t u ⊢ e @0 ⇝ t ∷ A‹0/𝑗›
  -- e : Eq (𝑗 ⇒ A) t u ⊢ e @1 ⇝ u ∷ A‹1/𝑗›
  --
  -- ((δ 𝑖 ⇒ s) ∷ Eq (𝑗 ⇒ A) t u) @𝑘 ⇝ s‹𝑘/𝑖› ∷ A‹𝑘/𝑗›
  whnf defs ctx (DApp f p appLoc) = do
    Element f fnf <- whnf defs ctx f
    case nchoose $ isDLamHead f of
      Left _ => case f of
        Ann (DLam {body, _}) (Eq {ty, l, r, _}) _ =>
          whnf defs ctx $
            Ann (endsOr (setLoc appLoc l) (setLoc appLoc r) (dsub1 body p) p)
                (dsub1 ty p) appLoc
        Coe ty p' q' val _ =>
          eqCoe defs ctx ty p' q' val p appLoc
      Right ndlh => case p of
        K e _ => do
          Eq {l, r, ty, _} <- whnf0 defs ctx =<< computeElimType defs ctx f
            | ty => throw $ ExpectedEq ty.loc ctx.names ty
          whnf defs ctx $
            ends (Ann (setLoc appLoc l) ty.zero appLoc)
                 (Ann (setLoc appLoc r) ty.one  appLoc) e
        B {} => pure $ Element (DApp f p appLoc) $ fnf `orNo` ndlh `orNo` Ah

  -- e ∷ A ⇝ e
  whnf defs ctx (Ann s a annLoc) = do
    Element s snf <- whnf defs ctx s
    case nchoose $ isE s of
      Left  _  => let E e = s in pure $ Element e $ noOr2 snf
      Right ne => do
        Element a anf <- whnf defs ctx a
        pure $ Element (Ann s a annLoc) $ ne `orNo` snf `orNo` anf

  whnf defs ctx (Coe (S _ (N ty)) _ _ val coeLoc) =
    whnf defs ctx $ Ann val ty coeLoc
  whnf defs ctx (Coe (S [< i] ty) p q val coeLoc) = do
    Element ty  tynf  <- whnf defs (extendDim i ctx) ty.term
    Element val valnf <- whnf defs ctx val
    pushCoe defs ctx i ty p q val coeLoc

  whnf defs ctx (Comp ty p q val r zero one compLoc) =
    -- comp [A] @p @p s { ⋯ } ⇝ s ∷ A
    if p == q then whnf defs ctx $ Ann val ty compLoc else
    case nchoose (isK r) of
      -- comp [A] @p @q s @0 { 0 j ⇒ t; ⋯ } ⇝ t‹q/j› ∷ A
      -- comp [A] @p @q s @1 { 1 j ⇒ t; ⋯ } ⇝ t‹q/j› ∷ A
      Left y => case r of
        K Zero _ => whnf defs ctx $ Ann (dsub1 zero q) ty compLoc
        K One  _ => whnf defs ctx $ Ann (dsub1 one  q) ty compLoc
      Right nk => do
        Element ty tynf <- whnf defs ctx ty
        pure $ Element (Comp ty p q val r zero one compLoc) $ tynf `orNo` nk

  whnf defs ctx (TypeCase ty ret arms def tcLoc) = do
    Element ty  tynf  <- whnf defs ctx ty
    Element ret retnf <- whnf defs ctx ret
    case nchoose $ isAnnTyCon ty of
      Left  y  =>
        let Ann ty (TYPE u _) _ = ty in
        reduceTypeCase defs ctx ty u ret arms def tcLoc
      Right nt => pure $
        Element (TypeCase ty ret arms def tcLoc) (tynf `orNo` retnf `orNo` nt)

  whnf defs ctx (CloE  (Sub el th)) = whnf defs ctx $ pushSubstsWith' id th el
  whnf defs ctx (DCloE (Sub el th)) = whnf defs ctx $ pushSubstsWith' th id el

export covering
CanWhnf Term Reduce.isRedexT where
  whnf _ _ t@(TYPE {}) = pure $ nred t
  whnf _ _ t@(Pi   {}) = pure $ nred t
  whnf _ _ t@(Lam  {}) = pure $ nred t
  whnf _ _ t@(Sig  {}) = pure $ nred t
  whnf _ _ t@(Pair {}) = pure $ nred t
  whnf _ _ t@(Enum {}) = pure $ nred t
  whnf _ _ t@(Tag  {}) = pure $ nred t
  whnf _ _ t@(Eq   {}) = pure $ nred t
  whnf _ _ t@(DLam {}) = pure $ nred t
  whnf _ _ t@(Nat  {}) = pure $ nred t
  whnf _ _ t@(Zero {}) = pure $ nred t
  whnf _ _ t@(Succ {}) = pure $ nred t
  whnf _ _ t@(BOX  {}) = pure $ nred t
  whnf _ _ t@(Box  {}) = pure $ nred t

  -- s ∷ A ⇝ s   (in term context)
  whnf defs ctx (E e) = do
    Element e enf <- whnf defs ctx e
    case nchoose $ isAnn e of
      Left  _  => let Ann {tm, _} = e in pure $ Element tm $ noOr1 $ noOr2 enf
      Right na => pure $ Element (E e) $ na `orNo` enf

  whnf defs ctx (CloT  (Sub tm th)) = whnf defs ctx $ pushSubstsWith' id th tm
  whnf defs ctx (DCloT (Sub tm th)) = whnf defs ctx $ pushSubstsWith' th id tm