module Quox.Whnf.Coercion

import Quox.Whnf.Interface
import Quox.Whnf.ComputeElimType
import Quox.Whnf.TypeCase

%default total



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 (N body)) =
  S names $ N $ E $ Coe ty p q body loc
coeScoped ty p q loc (S names (Y body)) =
  ST names $ E $ Coe (weakDS s ty) p q body loc
where
  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


parameters {auto _ : CanWhnf Term Interface.isRedexT}
           {auto _ : CanWhnf Elim Interface.isRedexE}
           {d, n : Nat} (defs : Definitions) (ctx : WhnfContext d n)
  ||| reduce a function application `App (Coe ty p q val) s loc`
  export 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`
  export 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`
  export 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`
  export 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



  ||| pushes a coercion inside a whnf-ed term
  export covering
  pushCoe : 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 i ty p q s loc {snf} =
    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