module Quox.Reduce

import Quox.No
import Quox.Syntax
import Quox.Definition
import Data.SnocVect
import Data.Maybe
import Data.List

%default total


||| errors that might happen if you pass an ill typed expression into
||| whnf. don't do that please
public export
data WhnfError = MissingEnumArm TagVal (List TagVal)


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

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

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

public export
0 NonRedex : (tm : TermLike) -> {isRedex : RedexTest tm} ->
             Whnf tm isRedex err =>
             (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 _ : Whnf tm isRedex err) =>
       (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 (Lam {} :# Pi {}) = True
isLamHead _                 = False

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

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

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

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

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

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

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

||| true if a term is syntactically a type.
public export %inline
isTyCon : (t : 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 : (t : Term {}) -> Bool
isTyConE s = isTyCon s || isE s

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


mutual
  public export
  isRedexE : RedexTest Elim
  isRedexE defs (F x) {d, n} =
    isJust $ lookupElim x defs {d, n}
  isRedexE _ (B _) = False
  isRedexE defs (f :@ _) =
    isRedexE defs f || isLamHead f
  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 (f :% _) =
    isRedexE defs f || isDLamHead f
  isRedexE defs (t :# a) =
    isE t || isRedexT defs t || isRedexT defs a
  isRedexE defs (TypeCase {ty, _}) =
    isRedexE defs ty || 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

mutual
  export covering
  Whnf Elim Reduce.isRedexE WhnfError where
    whnf defs (F x) with (lookupElim x defs) proof eq
      _ | Just y  = whnf defs y
      _ | Nothing = pure $ Element (F x) $ rewrite eq in Ah

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

    -- ((λ x ⇒ t) ∷ (π.x : A) → B) s ⇝ t[s∷A/x] ∷ B[s∷A/x]
    whnf defs (f :@ s) = do
      Element f fnf <- whnf defs f
      case nchoose $ isLamHead f of
        Left _ =>
          let Lam body :# Pi {arg, res, _} = f
              s = s :# arg
          in
          whnf defs $ sub1 body s :# sub1 res s
        Right nlh => pure $ Element (f :@ s) $ 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 (CasePair pi pair ret body) = do
      Element pair pairnf <- whnf defs pair
      case nchoose $ isPairHead pair of
        Left _ =>
          let Pair {fst, snd} :# Sig {fst = tfst, snd = tsnd, _} = pair
              fst = fst :# tfst
              snd = snd :# sub1 tsnd fst
          in
          whnf defs $ subN body [< fst, snd] :# sub1 ret pair
        Right np =>
          pure $ Element (CasePair pi pair ret body)
                         (pairnf `orNo` np)

    -- case 'a ∷ {a,…} return p ⇒ C of { 'a ⇒ u } ⇝
    -- u ∷ C['a∷{a,…}/p]
    whnf defs (CaseEnum pi tag ret arms) = do
      Element tag tagnf <- whnf defs tag
      case nchoose $ isTagHead tag of
        Left t =>
          let Tag t :# Enum ts = tag
              ty = sub1 ret tag
          in
          case lookup t arms of
            Just arm => whnf defs $ arm :# ty
            Nothing  => Left $ MissingEnumArm t (keys arms)
        Right nt =>
          pure $ Element (CaseEnum pi tag ret arms) $ 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 (CaseNat pi piIH nat ret zer suc) = do
      Element nat natnf <- whnf defs nat
      case nchoose $ isNatHead nat of
        Left _ =>
          let ty = sub1 ret nat in
          case nat of
            Zero :# Nat => whnf defs (zer :# ty)
            Succ n :# Nat =>
              let nn = n :# Nat
                  tm = subN suc [< nn, CaseNat pi piIH nn ret zer suc]
              in
              whnf defs $ tm :# ty
        Right nn =>
          pure $ Element (CaseNat pi piIH nat ret zer suc) $ natnf `orNo` nn

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

    -- ((δ 𝑖 ⇒ s) ∷ Eq [𝑗 ⇒ A] t u) @0 ⇝ t ∷ A‹0/𝑗›
    -- ((δ 𝑖 ⇒ s) ∷ Eq [𝑗 ⇒ A] t u) @1 ⇝ u ∷ A‹1/𝑗›
    -- ((δ 𝑖 ⇒ s) ∷ Eq [𝑗 ⇒ A] t u) @𝑘 ⇝ s‹𝑘/𝑖› ∷ A‹𝑘/𝑗›
    --   (if 𝑘 is a variable)
    whnf defs (f :% p) = do
      Element f fnf <- whnf defs f
      case nchoose $ isDLamHead f of
        Left _ =>
          let DLam body :# Eq {ty = ty, l, r, _} = f
              body = endsOr l r (dsub1 body p) p
          in
          whnf defs $ body :# dsub1 ty p
        Right ndlh =>
          pure $ Element (f :% p) $ fnf `orNo` ndlh

    -- e ∷ A ⇝ e
    whnf defs (s :# a) = do
      Element s snf <- whnf defs 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 a
          pure $ Element (s :# a) $ ne `orNo` snf `orNo` anf

    -- (type-case ★ᵢ ∷ _ return Q of { ★ ⇒ s; ⋯ }) ⇝ s ∷ Q
    --
    -- (type-case π.(x : A) → B ∷ ★ᵢ return Q of { (a → b) ⇒ s; ⋯ }) ⇝
    -- s[(A ∷ ★ᵢ)/a, ((λ x ⇒ B) ∷ 0.A → ★ᵢ)/b] ∷ ★ᵢ
    --
    -- (type-case (x : A) × B ∷ ★ᵢ return Q of { (a × b) ⇒ s; ⋯ }) ⇝
    -- s[(A ∷ ★ᵢ)/a, ((λ x ⇒ B) ∷ 0.A → ★ᵢ)/b] ∷ ★ᵢ
    --
    -- (type-case {⋯} ∷ _ return Q of { {} ⇒ s; ⋯ }) ⇝ s ∷ Q
    --
    -- (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
    --
    -- (type-case ℕ ∷ _ return Q of { ℕ ⇒ s; ⋯ }) ⇝ s ∷ Q
    --
    -- (type-case [π.A] ∷ ★ᵢ return Q of { [a] ⇒ s; ⋯ }) ⇝ s[(A ∷ ★ᵢ)/a] ∷ Q
    whnf defs (TypeCase {ty, ret, univ, pi, sig, enum, eq, nat, box}) = do
      Element ty tynf <- whnf defs ty
      case nchoose $ isAnnTyCon ty of
        Left y =>
          let ty :# TYPE u = ty in
          case ty of
            TYPE _ =>
              whnf defs $ univ :# ret
            Pi _ arg res =>
              let arg = arg :# TYPE u
                  res = toLam res :# Arr Zero (TYPE u) (TYPE u)
              in
              whnf defs $ subN pi [< arg, res] :# ret
            Sig fst snd =>
              let fst = fst :# TYPE u
                  snd = toLam snd :# Arr Zero (TYPE u) (TYPE u)
              in
              whnf defs $ subN sig [< fst, snd] :# ret
            Enum _ =>
              whnf defs $ enum :# ret
            Eq a l r =>
              let a0' = a.zero; a1' = a.one
                  a0 = a0' :# TYPE u
                  a1 = a1' :# TYPE u
                  a  = toDLam a :# Eq0 (TYPE u) a0' a1'
                  l  = l :# a0'
                  r  = r :# a1'
              in
              whnf defs $ subN eq [< a0, a1, a, l, r] :# ret
            Nat =>
              whnf defs $ nat :# ret
            BOX _ s =>
              whnf defs $ sub1 box (s :# TYPE u) :# ret
        Right nt => pure $
          Element (TypeCase {ty, ret, univ, pi, sig, enum, eq, nat, box}) $
            tynf `orNo` nt
    where
      toLam : {s : Nat} -> ScopeTermN s d n -> Term d n
      toLam (S names body) = names :\\ body.term

      toDLam : {s : Nat} -> DScopeTermN s d n -> Term d n
      toDLam (S names body) = names :\\% body.term

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

  export covering
  Whnf Term Reduce.isRedexT WhnfError 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 _    Nat      = pure $ nred Nat
    whnf _    Zero     = pure $ nred Zero
    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 (E e) = do
      Element e enf <- whnf defs e
      case nchoose $ isAnn e of
        Left  _  => let tm :# _ = e in pure $ Element tm $ noOr1 $ noOr2 enf
        Right na => pure $ Element (E e) $ na `orNo` enf

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