module Quox.Whnf.Main

import Quox.Whnf.Interface
import Quox.Whnf.ComputeElimType
import Quox.Whnf.TypeCase
import Quox.Whnf.Coercion
import Quox.Pretty
import Quox.Displace
import Data.SnocVect

%default total


export covering CanWhnf Term Interface.isRedexT
export covering CanWhnf Elim Interface.isRedexE


-- the String is what to call the "s" argument in logs (maybe "s", or "e")
private %inline
whnfDefault :
  {0 isRedex : RedexTest tm} ->
  (CanWhnf tm isRedex, Located2 tm) =>
  String ->
  (forall d, n. WhnfContext d n -> tm d n -> Eff Pretty LogDoc) ->
  (defs : Definitions) ->
  (ctx : WhnfContext d n) ->
  (sg : SQty) ->
  (s : tm d n) ->
  Eff Whnf (Subset (tm d n) (No . isRedex defs ctx sg))
whnfDefault name ppr defs ctx sg s = do
  sayMany "whnf" s.loc
    [10 :> "whnf",
     95 :> hsep ["ctx =", runPretty $ prettyWhnfContext ctx],
     95 :> hsep ["sg =", runPretty $ prettyQty sg.qty],
     10 :> hsep [text name, "=", runPretty $ ppr ctx s]]
  res <- whnfNoLog defs ctx sg s
  say "whnf" 11 s.loc $ hsep ["whnf ⇝", runPretty $ ppr ctx res.fst]
  pure res

covering
CanWhnf Elim Interface.isRedexE where
  whnf = whnfDefault "e" $ \ctx, e => prettyElim ctx.dnames ctx.tnames e

  whnfNoLog defs ctx sg (F x u loc) with (lookupElim0 x u defs) proof eq
    _ | Just y  = whnf defs ctx sg $ setLoc loc $ injElim ctx y
    _ | Nothing = pure $ Element (F x u loc) $ rewrite eq in Ah

  whnfNoLog defs ctx sg (B i loc) with (ctx.tctx !! i) proof eq1
    _ | l with (l.term) proof eq2
      _ | Just y  = whnf defs ctx sg $ Ann y l.type loc
      _ | Nothing = pure $ Element (B i loc) $ rewrite eq1 in rewrite eq2 in Ah

  -- ((λ x ⇒ t) ∷ (π.x : A) → B) s ⇝ t[s∷A/x] ∷ B[s∷A/x]
  whnfNoLog defs ctx sg (App f s appLoc) = do
    Element f fnf <- whnf defs ctx sg 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 sg $ Ann (sub1 body s) (sub1 res s) appLoc
        Coe ty p q val _ => piCoe defs ctx sg 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]
  --
  -- 0 · case e return p ⇒ C of { (a, b) ⇒ u } ⇝
  -- u[fst e/a, snd e/b] ∷ C[e/p]
  whnfNoLog defs ctx sg (CasePair pi pair ret body caseLoc) = do
    Element pair pairnf <- whnf defs ctx sg 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 sg $ Ann (subN body [< fst, snd]) (sub1 ret pair) caseLoc
        Coe ty p q val _ => do
          sigCoe defs ctx sg pi ty p q val ret body caseLoc
      Right np =>
        case sg `decEq` SZero of
          Yes Refl =>
            whnf defs ctx SZero $
              Ann (subN body [< Fst pair caseLoc, Snd pair caseLoc])
                  (sub1 ret pair)
                  caseLoc
          No n0 =>
            pure $ Element (CasePair pi pair ret body caseLoc)
                           (pairnf `orNo` np `orNo` notYesNo n0)

  -- fst ((s, t) ∷ (x : A) × B) ⇝ s ∷ A
  whnfNoLog defs ctx sg (Fst pair fstLoc) = do
    Element pair pairnf <- whnf defs ctx sg pair
    case nchoose $ isPairHead pair of
      Left _ => case pair of
        Ann (Pair {fst, snd, _}) (Sig {fst = tfst, snd = tsnd, _}) pairLoc =>
          whnf defs ctx sg $ Ann fst tfst pairLoc
        Coe ty p q val _ => do
          fstCoe defs ctx sg ty p q val fstLoc
      Right np =>
        pure $ Element (Fst pair fstLoc) (pairnf `orNo` np)

  -- snd ((s, t) ∷ (x : A) × B) ⇝ t ∷ B[(s ∷ A)/x]
  whnfNoLog defs ctx sg (Snd pair sndLoc) = do
    Element pair pairnf <- whnf defs ctx sg pair
    case nchoose $ isPairHead pair of
      Left _ => case pair of
        Ann (Pair {fst, snd, _}) (Sig {fst = tfst, snd = tsnd, _}) pairLoc =>
          whnf defs ctx sg $ Ann snd (sub1 tsnd (Ann fst tfst fst.loc)) sndLoc
        Coe ty p q val _ => do
          sndCoe defs ctx sg ty p q val sndLoc
      Right np =>
        pure $ Element (Snd pair sndLoc) (pairnf `orNo` np)

  -- case 'a ∷ {a,…} return p ⇒ C of { 'a ⇒ u } ⇝
  -- u ∷ C['a∷{a,…}/p]
  whnfNoLog defs ctx sg (CaseEnum pi tag ret arms caseLoc) = do
    Element tag tagnf <- whnf defs ctx sg 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 sg $ 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 sg $
            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]
  whnfNoLog defs ctx sg (CaseNat pi piIH nat ret zer suc caseLoc) = do
    Element nat natnf <- whnf defs ctx sg nat
    case nchoose $ isNatHead nat of
      Left _ =>
        let ty = sub1 ret nat in
        case nat of
          Ann (Nat 0 _) (NAT _) _ =>
            whnf defs ctx sg $ Ann zer ty zer.loc
          Ann (Nat (S n) succLoc) (NAT natLoc) _ =>
            let nn = Ann (Nat n succLoc) (NAT natLoc) succLoc
                tm = subN suc [< nn, CaseNat pi piIH nn ret zer suc caseLoc]
            in
            whnf defs ctx sg $ Ann tm ty caseLoc
          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 sg $ Ann tm ty caseLoc
          Coe ty p q val _ =>
            -- same deal as Enum
            whnf defs ctx sg $
              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]
  whnfNoLog defs ctx sg (CaseBox pi box ret body caseLoc) = do
    Element box boxnf <- whnf defs ctx sg 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 sg $ Ann (sub1 body (Ann val bty val.loc)) ty caseLoc
        Coe ty p q val _ =>
          boxCoe defs ctx sg 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‹𝑘/𝑗›
  whnfNoLog defs ctx sg (DApp f p appLoc) = do
    Element f fnf <- whnf defs ctx sg f
    case nchoose $ isDLamHead f of
      Left _ => case f of
        Ann (DLam {body, _}) (Eq {ty, l, r, _}) _ =>
          whnf defs ctx sg $
            Ann (endsOr (setLoc appLoc l) (setLoc appLoc r) (dsub1 body p) p)
                (dsub1 ty p) appLoc
        Coe ty p' q' val _ =>
          eqCoe defs ctx sg ty p' q' val p appLoc
      Right ndlh => case p of
        K e _ => do
          Eq {l, r, ty, _} <- computeWhnfElimType0 defs ctx sg f
            | ty => throw $ ExpectedEq ty.loc ctx.names ty
          whnf defs ctx sg $
            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
  whnfNoLog defs ctx sg (Ann s a annLoc) = do
    Element s snf <- whnf defs ctx sg 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 SZero a
        pure $ Element (Ann s a annLoc) (ne `orNo` snf `orNo` anf)

  whnfNoLog defs ctx sg (Coe sty@(S [< i] ty) p q val coeLoc) =
    -- reduction if A‹0/𝑖› = A‹1/𝑖› lives in Equal
    case p `decEqv` q of
      -- coe (𝑖 ⇒ A) @p @p s ⇝ (s ∷ A‹p/𝑖›)
      Yes _   => whnf defs ctx sg $ Ann val (dsub1 sty p) coeLoc
      No  npq => do
        let ty = getTerm ty
        Element ty tynf <- whnf defs (extendDim i ctx) SZero ty
        case nchoose $ canPushCoe sg ty val of
          Left  pc  => pushCoe defs ctx sg i ty p q val coeLoc
          Right npc => pure $ Element (Coe (SY [< i] ty) p q val coeLoc)
                                      (tynf `orNo` npc `orNo` notYesNo npq)

  whnfNoLog defs ctx sg (Comp ty p q val r zero one compLoc) =
    case p `decEqv` q of
      -- comp [A] @p @p s @r { ⋯ } ⇝ s ∷ A
      Yes y   => whnf defs ctx sg $ Ann val ty compLoc
      No  npq => case r of
        -- comp [A] @p @q s @0 { 0 𝑗 ⇒ t₀; ⋯ } ⇝ t₀‹q/𝑗› ∷ A
        K Zero _ => whnf defs ctx sg $ Ann (dsub1 zero q) ty compLoc
        -- comp [A] @p @q s @1 { 1 𝑗 ⇒ t₁; ⋯ } ⇝ t₁‹q/𝑗› ∷ A
        K One  _ => whnf defs ctx sg $ Ann (dsub1 one  q) ty compLoc
        B {}     => pure $ Element (Comp ty p q val r zero one compLoc)
                                   (notYesNo npq `orNo` Ah)

  whnfNoLog defs ctx sg (TypeCase ty ret arms def tcLoc) =
    case sg `decEq` SZero of
      Yes Refl => do
        Element ty  tynf  <- whnf defs ctx SZero ty
        Element ret retnf <- whnf defs ctx SZero 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)
      No _ =>
        throw $ ClashQ tcLoc sg.qty Zero

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

covering
CanWhnf Term Interface.isRedexT where
  whnf = whnfDefault "e" $ \ctx, s => prettyTerm ctx.dnames ctx.tnames s

  whnfNoLog _ _ _ t@(TYPE    {}) = pure $ nred t
  whnfNoLog _ _ _ t@(IOState {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Pi      {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Lam     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Sig     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Pair    {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Enum    {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Tag     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Eq      {}) = pure $ nred t
  whnfNoLog _ _ _ t@(DLam    {}) = pure $ nred t
  whnfNoLog _ _ _ t@(NAT     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Nat     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(STRING  {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Str     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(BOX     {}) = pure $ nred t
  whnfNoLog _ _ _ t@(Box     {}) = pure $ nred t

  whnfNoLog _ _ _ (Succ p loc) =
    case nchoose $ isNatConst p of
      Left _ => case p of
        Nat p _               => pure $ nred $ Nat (S p) loc
        E (Ann (Nat p _) _ _) => pure $ nred $ Nat (S p) loc
      Right nc => pure $ nred $ Succ p loc

  whnfNoLog defs ctx sg (Let _ rhs body _) =
    whnf defs ctx sg $ sub1 body rhs

  -- s ∷ A ⇝ s   (in term context)
  whnfNoLog defs ctx sg (E e) = do
    Element e enf <- whnf defs ctx sg 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

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