module Quox.Whnf.TypeCase

import Quox.Whnf.Interface
import Quox.Whnf.ComputeElimType
import Data.SnocVect

%default total


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

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

private
tycaseRhs0 : (k : TyConKind) -> TypeCaseArms d n ->
             (0 eq : arity k = 0) => Maybe (Term d n)
tycaseRhs0 k arms = map (.term0) $ rewrite sym eq in tycaseRhs k arms

private
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


parameters {auto _ : CanWhnf Term Interface.isRedexT}
           {auto _ : CanWhnf Elim Interface.isRedexE}
           (defs : Definitions) (ctx : WhnfContext 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
  export covering
  tycasePi : (t : Term d n) -> (0 tnf : No (isRedexT defs ctx SZero t)) =>
             Eff Whnf (Term d n, ScopeTerm d n)
  tycasePi (Pi {arg, res, _}) = pure (arg, res)
  tycasePi (E e) {tnf} = do
    ty <- computeElimType defs ctx SZero e {ne = noOr2 tnf}
    let loc  = e.loc
        narg = mnb "Arg" loc; nret = mnb "Ret" loc
        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  = ST [< !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
  export covering
  tycaseSig : (t : Term d n) -> (0 tnf : No (isRedexT defs ctx SZero t)) =>
              Eff Whnf (Term d n, ScopeTerm d n)
  tycaseSig (Sig {fst, snd, _}) = pure (fst, snd)
  tycaseSig (E e) {tnf} = do
    ty <- computeElimType defs ctx SZero e {ne = noOr2 tnf}
    let loc  = e.loc
        nfst = mnb "Fst" loc; nsnd = mnb "Snd" loc
        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  = ST [< !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
  export covering
  tycaseBOX : (t : Term d n) -> (0 tnf : No (isRedexT defs ctx SZero t)) =>
              Eff Whnf (Term d n)
  tycaseBOX (BOX {ty, _}) = pure ty
  tycaseBOX (E e) {tnf} = do
    ty <- computeElimType defs ctx SZero e {ne = noOr2 tnf}
    pure $ E $ typeCase1Y e ty KBOX [< !(mnb "Ty" e.loc)] (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
  export covering
  tycaseEq : (t : Term d n) -> (0 tnf : No (isRedexT defs ctx SZero t)) =>
             Eff Whnf (Term d n, Term d n, DScopeTerm d n, Term d n, Term d n)
  tycaseEq (Eq {ty, l, r, _}) = pure (ty.zero, ty.one, ty, l, r)
  tycaseEq (E e) {tnf} = do
    ty <- computeElimType defs ctx SZero e {ne = noOr2 tnf}
    let loc   = e.loc
        names = traverse' (\x => mnb x loc) [< "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     = DST [< !(mnb "i" loc)] $ 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

  ||| reduce a type-case applied to a type constructor
  |||
  ||| `reduceTypeCase A i Q arms def _` reduces an expression
  ||| `type-case A ∷ ★ᵢ return Q of { arms; _ ⇒ def }`
  export covering
  reduceTypeCase : (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 ctx SZero))
  reduceTypeCase ty u ret arms def loc = case ty of
    -- (type-case ★ᵢ ∷ _ return Q of { ★ ⇒ s; ⋯ }) ⇝ s ∷ Q
    TYPE {} =>
      whnf defs ctx SZero $ Ann (tycaseRhsDef0 def KTYPE arms) ret loc

    -- (type-case IOState ∷ _ return Q of { IOState ⇒ s; ⋯ }) ⇝ s ∷ Q
    IOState {} =>
      whnf defs ctx SZero $ Ann (tycaseRhsDef0 def KIOState 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 SZero $
        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 SZero $
        Ann (subN (tycaseRhsDef def KSig arms) [< fst', snd']) ret loc

    -- (type-case {⋯} ∷ _ return Q of { {} ⇒ s; ⋯ }) ⇝ s ∷ Q
    Enum {} =>
      whnf defs ctx SZero $ 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 SZero $ 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 SZero $ Ann (tycaseRhsDef0 def KNat arms) ret loc

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

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