module Quox.Equal

import Quox.BoolExtra
import public Quox.Typing
import public Control.Monad.Either
import public Control.Monad.Reader
import Data.Maybe


public export
record CmpContext where
  constructor MkCmpContext
  mode : EqMode


public export
0 HasCmpContext : (Type -> Type) -> Type
HasCmpContext = MonadReader CmpContext

public export
0 CanEqual : (q : Type) -> (Type -> Type) -> Type
CanEqual q m = (HasErr q m, HasCmpContext m)



private %inline
mode : HasCmpContext m => m EqMode
mode = asks mode

parameters {auto _ : CanEqual q m} (ctx : EqContext q n)
  private %inline
  clashT : Term q 0 n -> Term q 0 n -> Term q 0 n -> m a
  clashT ty s t = throwError $ ClashT ctx !mode ty s t

  private %inline
  clashTy : Term q 0 n -> Term q 0 n -> m a
  clashTy s t = throwError $ ClashTy ctx !mode s t

  private %inline
  clashE : Elim q 0 n -> Elim q 0 n -> m a
  clashE e f = throwError $ ClashE ctx !mode e f

  private %inline
  wrongType : Term q 0 n -> Term q 0 n -> m a
  wrongType ty s = throwError $ WrongType ctx ty s


||| true if a term is syntactically a type, or is neutral.
|||
||| this function *doesn't* push substitutions, because its main use is as a
||| `So` argument to skip cases that are already known to be nonsense. and
||| the substitutions have already been pushed.
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 (E     {}) = True
isTyCon (CloT  {}) = False
isTyCon (DCloT {}) = False

public export %inline
sameTyCon : (s, t : Term q d n) ->
            (0 ts : So (isTyCon s)) => (0 tt : So (isTyCon t)) =>
            Bool
sameTyCon (TYPE {}) (TYPE {}) = True
sameTyCon (TYPE {}) _         = False
sameTyCon (Pi   {}) (Pi   {}) = True
sameTyCon (Pi   {}) _         = False
sameTyCon (Sig  {}) (Sig  {}) = True
sameTyCon (Sig  {}) _         = False
sameTyCon (Enum {}) (Enum {}) = True
sameTyCon (Enum {}) _         = False
sameTyCon (Eq   {}) (Eq   {}) = True
sameTyCon (Eq   {}) _         = False
sameTyCon Nat       Nat       = True
sameTyCon Nat       _         = False
sameTyCon (E    {}) (E    {}) = True
sameTyCon (E    {}) _         = False


parameters (defs : Definitions' q g)
  ||| true if a type is known to be a subsingleton purely by its form.
  ||| a subsingleton is a type with only zero or one possible values.
  ||| equality/subtyping accepts immediately on values of subsingleton types.
  |||
  ||| * a function type is a subsingleton if its codomain is.
  ||| * a pair type is a subsingleton if both its elements are.
  ||| * all equality types are subsingletons because uip is admissible by
  |||   boundary separation.
  ||| * an enum type is a subsingleton if it has zero or one tags.
  public export
  isSubSing : HasErr q m => {n : Nat} -> Term q 0 n -> m Bool
  isSubSing ty = do
    Element ty nc <- whnfT defs ty
    case ty of
      TYPE _          => pure False
      Pi   {res, _}   => isSubSing res.term
      Lam  {}         => pure False
      Sig  {fst, snd} => isSubSing fst `andM` isSubSing snd.term
      Pair {}         => pure False
      Enum tags       => pure $ length (SortedSet.toList tags) <= 1
      Tag  {}         => pure False
      Eq   {}         => pure True
      DLam {}         => pure False
      Nat             => pure False
      Zero            => pure False
      Succ {}         => pure False
      E    (s :# _)   => isSubSing s
      E    _          => pure False


export
ensureTyCon : HasErr q m =>
              (ctx : EqContext q n) -> (t : Term q 0 n) -> m (So (isTyCon t))
ensureTyCon ctx t = case nchoose $ isTyCon t of
  Left  y => pure y
  Right n => throwError $ NotType (toTyContext ctx) (t // shift0 ctx.dimLen)

parameters (defs : Definitions' q _) {auto _ : (CanEqual q m, IsQty q)}
  mutual
    namespace Term
      ||| `compare0 ctx ty s t` compares `s` and `t` at type `ty`, according to
      ||| the current variance `mode`.
      |||
      ||| ⚠ **assumes that `s`, `t` have already been checked against `ty`**. ⚠
      export covering %inline
      compare0 : EqContext q n -> (ty, s, t : Term q 0 n) -> m ()
      compare0 ctx ty s t =
        wrapErr (WhileComparingT ctx !mode ty s t) $ do
          let Val n = ctx.termLen
          Element ty nty <- whnfT defs ty
          Element s  ns  <- whnfT defs s
          Element t  nt  <- whnfT defs t
          tty <- ensureTyCon ctx ty
          compare0' ctx ty s t

      ||| converts an elim "Γ ⊢ e" to "Γ, x ⊢ e x", for comparing with
      ||| a lambda "Γ ⊢ λx ⇒ t" that has been converted to "Γ, x ⊢ t".
      private %inline
      toLamBody : Elim q d n -> Term q d (S n)
      toLamBody e = E $ weakE e :@ BVT 0

      private covering
      compare0' : EqContext q n ->
                  (ty, s, t : Term q 0 n) ->
                  (0 nty : NotRedex defs ty) => (0 tty : So (isTyCon ty)) =>
                  (0 ns  : NotRedex defs s)  => (0 nt  : NotRedex defs t) =>
                  m ()
      compare0' ctx (TYPE _) s t = compareType ctx s t

      compare0' ctx ty@(Pi {qty, arg, res}) s t {n} = local {mode := Equal} $
        case (s, t) of
          --           Γ, x : A ⊢ s = t : B
          -- -------------------------------------------
          --  Γ ⊢ (λ x ⇒ s) = (λ x ⇒ t) : (π·x : A) → B
          (Lam b1, Lam b2) => compare0 ctx' res.term b1.term b2.term

          --       Γ, x : A ⊢ s = e x : B
          -- -----------------------------------
          --  Γ ⊢ (λ x ⇒ s) = e : (π·x : A) → B
          (E e,   Lam b) => eta e b
          (Lam b, E e)   => eta e b

          (E e, E f) => Elim.compare0 ctx e f

          (Lam _, t) => wrongType ctx ty t
          (E _,   t) => wrongType ctx ty t
          (s,     _) => wrongType ctx ty s
        where
          ctx' : EqContext q (S n)
          ctx' = extendTy qty res.name arg ctx

          eta : Elim q 0 n -> ScopeTerm q 0 n -> m ()
          eta e (S _ (Y b)) = compare0 ctx' res.term (toLamBody e) b
          eta e (S _ (N _)) = clashT ctx ty s t

      compare0' ctx ty@(Sig {fst, snd, _}) s t = local {mode := Equal} $
        case (s, t) of
          --  Γ ⊢ s₁ = t₁ : A     Γ ⊢ s₂ = t₂ : B{s₁/x}
          -- --------------------------------------------
          --     Γ ⊢ (s₁, t₁) = (s₂,t₂) : (x : A) × B
          --
          -- [todo] η for π ≥ 0 maybe
          (Pair sFst sSnd, Pair tFst tSnd) => do
            compare0 ctx fst sFst tFst
            compare0 ctx (sub1 snd (sFst :# fst)) sSnd tSnd

          (E e, E f) => Elim.compare0 ctx e f

          (Pair {}, E _)     => clashT ctx ty s t
          (E _,     Pair {}) => clashT ctx ty s t

          (Pair {}, t) => wrongType ctx ty t
          (E _,     t) => wrongType ctx ty t
          (s,       _) => wrongType ctx ty s

      compare0' ctx ty@(Enum tags) s t = local {mode := Equal} $
        case (s, t) of
          -- --------------------
          --  Γ ⊢ `t = `t : {ts}
          --
          -- t ∈ ts is in the typechecker, not here, ofc
          (Tag t1, Tag t2) => unless (t1 == t2) $ clashT ctx ty s t
          (E e,    E f)    => Elim.compare0 ctx e f

          (Tag _, E _)   => clashT ctx ty s t
          (E _,   Tag _) => clashT ctx ty s t

          (Tag _, t) => wrongType ctx ty t
          (E _,   t) => wrongType ctx ty t
          (s,     _) => wrongType ctx ty s

      compare0' _ (Eq {}) _ _ =
        -- ✨ uip ✨
        --
        -- Γ ⊢ e = f : Eq [i ⇒ A] s t
        pure ()

      compare0' ctx Nat s t = local {mode := Equal} $
        case (s, t) of
          -- ---------------
          --  Γ ⊢ 0 = 0 : ℕ
          (Zero, Zero) => pure ()

          --      Γ ⊢ m = n : ℕ
          -- -------------------------
          --  Γ ⊢ succ m = succ n : ℕ
          (Succ m, Succ n) => compare0 ctx Nat m n

          (E e, E f) => Elim.compare0 ctx e f

          (Zero,   Succ _) => clashT ctx Nat s t
          (Zero,   E _)    => clashT ctx Nat s t
          (Succ _, Zero)   => clashT ctx Nat s t
          (Succ _, E _)    => clashT ctx Nat s t
          (E _,    Zero)   => clashT ctx Nat s t
          (E _,    Succ _) => clashT ctx Nat s t

          (Zero,   t) => wrongType ctx Nat t
          (Succ _, t) => wrongType ctx Nat t
          (E _,    t) => wrongType ctx Nat t
          (s,      _) => wrongType ctx Nat s

      compare0' ctx ty@(E _) s t = do
        -- a neutral type can only be inhabited by neutral values
        -- e.g. an abstract value in an abstract type, bound variables, …
        E e <- pure s | _ => wrongType ctx ty s
        E f <- pure t | _ => wrongType ctx ty t
        Elim.compare0 ctx e f

    ||| compares two types, using the current variance `mode` for universes.
    ||| fails if they are not types, even if they would happen to be equal.
    export covering %inline
    compareType : EqContext q n -> (s, t : Term q 0 n) -> m ()
    compareType ctx s t = do
      let Val n = ctx.termLen
      Element s ns <- whnfT defs s
      Element t nt <- whnfT defs t
      ts <- ensureTyCon ctx s
      tt <- ensureTyCon ctx t
      st <- either pure (const $ clashTy ctx s t) $
            nchoose $ sameTyCon s t
      compareType' ctx s t

    private covering
    compareType' : EqContext q n -> (s, t : Term q 0 n) ->
                   (0 ns : NotRedex defs s) => (0 ts : So (isTyCon s)) =>
                   (0 nt : NotRedex defs t) => (0 tt : So (isTyCon t)) =>
                   (0 st : So (sameTyCon s t)) =>
                   m ()
    -- equality is the same as subtyping, except with the
    -- "≤" in the TYPE rule being replaced with "="
    compareType' ctx (TYPE k) (TYPE l) =
      --        𝓀 ≤ ℓ
      -- ----------------------
      --  Γ ⊢ Type 𝓀 <: Type ℓ
      expectModeU !mode k l

    compareType' ctx (Pi {qty = sQty, arg = sArg, res = sRes, _})
                     (Pi {qty = tQty, arg = tArg, res = tRes, _}) = do
      --   Γ ⊢ A₁ :> A₂    Γ, x : A₁ ⊢ B₁ <: B₂
      -- ----------------------------------------
      --  Γ ⊢ (π·x : A₁) → B₁ <: (π·x : A₂) → B₂
      expectEqualQ sQty tQty
      local {mode $= flip} $ compareType ctx sArg tArg -- contra
      compareType (extendTy zero sRes.name sArg ctx) sRes.term tRes.term

    compareType' ctx (Sig {fst = sFst, snd = sSnd, _})
                     (Sig {fst = tFst, snd = tSnd, _}) = do
      --  Γ ⊢ A₁ <: A₂    Γ, x : A₁ ⊢ B₁ <: B₂
      -- --------------------------------------
      --   Γ ⊢ (x : A₁) × B₁ <: (x : A₂) × B₂
      compareType ctx sFst tFst
      compareType (extendTy zero sSnd.name sFst ctx) sSnd.term tSnd.term

    compareType' ctx (Eq {ty = sTy, l = sl, r = sr, _})
                     (Eq {ty = tTy, l = tl, r = tr, _}) = do
      --            Γ ⊢ A₁‹ε/i› <: A₂‹ε/i›
      --  Γ ⊢ l₁ = l₂ : A₁‹𝟎/i›    Γ ⊢ r₁ = r₂ : A₁‹𝟏/i›
      -- ------------------------------------------------
      --    Γ ⊢ Eq [i ⇒ A₁] l₁ r₂ <: Eq [i ⇒ A₂] l₂ r₂
      compareType (extendDim sTy.name Zero ctx) sTy.zero tTy.zero
      compareType (extendDim sTy.name One  ctx) sTy.one  tTy.one
      local {mode := Equal} $ do
        Term.compare0 ctx sTy.zero sl tl
        Term.compare0 ctx sTy.one  sr tr

    compareType' ctx s@(Enum tags1) t@(Enum tags2) = do
      -- ------------------
      --  Γ ⊢ {ts} <: {ts}
      --
      -- no subtyping based on tag subsets, since that would need
      -- a runtime coercion
      unless (tags1 == tags2) $ clashTy ctx s t

    compareType' ctx Nat Nat =
      -- ------------
      --  Γ ⊢ ℕ <: ℕ
      pure ()

    compareType' ctx (E e) (E f) = do
      -- no fanciness needed here cos anything other than a neutral
      -- has been inlined by whnf
      Elim.compare0 ctx e f

    ||| performs the minimum work required to recompute the type of an elim.
    |||
    ||| ⚠ **assumes the elim is already typechecked.** ⚠
    private covering
    computeElimType : EqContext q n -> (e : Elim q 0 n) ->
                      (0 ne : NotRedex defs e) ->
                      m (Term q 0 n)
    computeElimType ctx (F x) _ = do
      defs <- lookupFree' defs x
      pure $ injectT ctx defs.type
    computeElimType ctx (B i) _ = pure $ ctx.tctx !! i
    computeElimType ctx (f :@ s) ne = do
      (_, arg, res) <- expectPiE defs ctx !(computeElimType ctx f (noOr1 ne))
      pure $ sub1 res (s :# arg)
    computeElimType ctx (CasePair {pair, ret, _}) _ = pure $ sub1 ret pair
    computeElimType ctx (CaseEnum {tag, ret, _}) _ = pure $ sub1 ret tag
    computeElimType ctx (CaseNat {nat, ret, _}) _ = pure $ sub1 ret nat
    computeElimType ctx (f :% p) ne = do
      (ty, _, _) <- expectEqE defs ctx !(computeElimType ctx f (noOr1 ne))
      pure $ dsub1 ty p
    computeElimType ctx (_ :# ty) _ = pure ty

    private covering
    replaceEnd : EqContext q n ->
                 (e : Elim q 0 n) -> DimConst -> (0 ne : NotRedex defs e) ->
                 m (Elim q 0 n)
    replaceEnd ctx e p ne = do
      (ty, l, r) <- expectEqE defs ctx !(computeElimType ctx e ne)
      pure $ ends l r p :# dsub1 ty (K p)

    namespace Elim
      -- [fixme] the following code ends up repeating a lot of work in the
      -- computeElimType calls. the results should be shared better

      ||| compare two eliminations according to the given variance `mode`.
      |||
      ||| ⚠ **assumes that they have both been typechecked, and have
      ||| equal types.** ⚠
      export covering %inline
      compare0 : EqContext q n -> (e, f : Elim q 0 n) -> m ()
      compare0 ctx e f =
        wrapErr (WhileComparingE ctx !mode e f) $ do
          let Val n = ctx.termLen
          Element e ne <- whnfT defs e
          Element f nf <- whnfT defs f
          -- [fixme] there is a better way to do this "isSubSing" stuff for sure
          unless !(isSubSing defs !(computeElimType ctx e ne)) $
            compare0' ctx e f ne nf

      private covering
      compare0' : EqContext q n ->
                  (e, f : Elim q 0 n) ->
                  (0 ne : NotRedex defs e) -> (0 nf : NotRedex defs f) ->
                  m ()
      -- replace applied equalities with the appropriate end first
      -- e.g. e : Eq [i ⇒ A] s t ⊢ e 𝟎 = s : A‹𝟎/i›
      --
      -- [todo] maybe have typed whnf and do this (and η???) there instead
      compare0' ctx (e :% K p) f ne nf =
        compare0 ctx !(replaceEnd ctx e p $ noOr1 ne) f
      compare0' ctx e (f :% K q) ne nf =
        compare0 ctx e !(replaceEnd ctx f q $ noOr1 nf)

      compare0' ctx e@(F x) f@(F y) _ _ = unless (x == y) $ clashE ctx e f
      compare0' ctx e@(F _) f _ _ = clashE ctx e f

      compare0' ctx e@(B i) f@(B j) _ _ = unless (i == j) $ clashE ctx e f
      compare0' ctx e@(B _) f _ _ = clashE ctx e f

      compare0' ctx (e :@ s) (f :@ t) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          (_, arg, _) <- expectPiE defs ctx !(computeElimType ctx e (noOr1 ne))
          Term.compare0 ctx arg s t
      compare0' ctx e@(_ :@ _) f _ _ = clashE ctx e f

      compare0' ctx (CasePair epi e eret ebody)
                    (CasePair fpi f fret fbody) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          ety <- computeElimType ctx e (noOr1 ne)
          compareType (extendTy zero eret.name ety ctx) eret.term fret.term
          (fst, snd) <- expectSigE defs ctx ety
          let [< x, y] = ebody.names
          Term.compare0 (extendTyN [< (epi, x, fst), (epi, y, snd.term)] ctx)
                        (substCasePairRet ety eret)
                        ebody.term fbody.term
          expectEqualQ epi fpi
      compare0' ctx e@(CasePair {}) f _ _ = clashE ctx e f

      compare0' ctx (CaseEnum epi e eret earms)
                    (CaseEnum fpi f fret farms) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          ety <- computeElimType ctx e (noOr1 ne)
          compareType (extendTy zero eret.name ety ctx) eret.term fret.term
          for_ !(expectEnumE defs ctx ety) $ \t =>
            compare0 ctx (sub1 eret $ Tag t :# ety)
              !(lookupArm t earms) !(lookupArm t farms)
          expectEqualQ epi fpi
        where
          lookupArm : TagVal -> CaseEnumArms q d n -> m (Term q d n)
          lookupArm t arms = case lookup t arms of
            Just arm => pure arm
            Nothing  => throwError $ TagNotIn t (fromList $ keys arms)
      compare0' ctx e@(CaseEnum {}) f _ _ = clashE ctx e f

      compare0' ctx (CaseNat epi epi' e eret ezer esuc)
                    (CaseNat fpi fpi' f fret fzer fsuc) ne nf =
        local {mode := Equal} $ do
          compare0 ctx e f
          ety <- computeElimType ctx e (noOr1 ne)
          compareType (extendTy zero eret.name ety ctx) eret.term fret.term
          compare0 ctx (sub1 eret (Zero :# Nat)) ezer fzer
          let [< p, ih] = esuc.names
          compare0 (extendTyN [< (epi, p, Nat), (epi', ih, eret.term)] ctx)
                   (weakT eret.term)
                   esuc.term fsuc.term
          expectEqualQ epi  fpi
          expectEqualQ epi' fpi'
      compare0' ctx e@(CaseNat {}) f _ _ = clashE ctx e f

      compare0' ctx (s :# a) (t :# b) _ _ =
          Term.compare0 ctx !(bigger a b) s t
        where
          bigger : forall a. a -> a -> m a
          bigger l r = asks mode <&> \case Super => l; _ => r

      compare0' ctx   (s :# a) f        _ _ = Term.compare0 ctx a s (E f)
      compare0' ctx e          (t :# b) _ _ = Term.compare0 ctx b (E e) t
      compare0' ctx e@(_ :# _) f        _ _ = clashE ctx e f


parameters {auto _ : (HasDefs' q _ m, HasErr q m, IsQty q)}
           (ctx : TyContext q d n)
  -- [todo] only split on the dvars that are actually used anywhere in
  -- the calls to `splits`

  parameters (mode : EqMode)
    namespace Term
      export covering
      compare : (ty, s, t : Term q d n) -> m ()
      compare ty s t = do
        defs <- ask
        runReaderT {m} (MkCmpContext {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            let ectx = makeEqContext ctx th in
            compare0 defs ectx (ty // th) (s // th) (t // th)

      export covering
      compareType : (s, t : Term q d n) -> m ()
      compareType s t = do
        defs <- ask
        runReaderT {m} (MkCmpContext {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            let ectx = makeEqContext ctx th in
            compareType defs ectx (s // th) (t // th)

    namespace Elim
      ||| you don't have to pass the type in but the arguments must still be
      ||| of the same type!!
      export covering %inline
      compare : (e, f : Elim q d n) -> m ()
      compare e f = do
        defs <- ask
        runReaderT {m} (MkCmpContext {mode}) $
          for_ (splits ctx.dctx) $ \th =>
            let ectx = makeEqContext ctx th in
            compare0 defs ectx (e // th) (f // th)

  namespace Term
    export covering %inline
    equal, sub, super : (ty, s, t : Term q d n) -> m ()
    equal = compare Equal
    sub   = compare Sub
    super = compare Super

    export covering %inline
    equalType, subtype, supertype : (s, t : Term q d n) -> m ()
    equalType = compareType Equal
    subtype   = compareType Sub
    supertype = compareType Super

  namespace Elim
    export covering %inline
    equal, sub, super : (e, f : Elim q d n) -> m ()
    equal = compare Equal
    sub   = compare Sub
    super = compare Super