module Quox.Syntax.Term.Base

import public Quox.Syntax.Var
import public Quox.Syntax.Shift
import public Quox.Syntax.Subst
import public Quox.Syntax.Qty
import public Quox.Syntax.Dim
import public Quox.Syntax.Term.TyConKind
import public Quox.Name
import public Quox.Context

import Quox.Pretty

import public Data.DPair
import Data.List
import Data.Maybe
import Data.Nat
import public Data.So
import Data.String
import public Data.SortedMap
import public Data.SortedMap.Dependent
import public Data.SortedSet
import Derive.Prelude

%default total
%language ElabReflection

%hide TT.Name


public export
TermLike : Type
TermLike = Nat -> Nat -> Type

public export
TSubstLike : Type
TSubstLike = Nat -> Nat -> Nat -> Type

public export
Universe : Type
Universe = Nat

public export
TagVal : Type
TagVal = String


public export
data ScopedBody : Nat -> (Nat -> Type) -> Nat -> Type where
  Y : (body : f (s + n)) -> ScopedBody s f n
  N : (body : f n)       -> ScopedBody s f n
%name ScopedBody body

export %inline %hint
EqScopedBody : (forall n. Eq (f n)) => Eq (ScopedBody s f n)
EqScopedBody = deriveEq

export %inline %hint
ShowScopedBody : (forall n. Show (f n)) => Show (ScopedBody s f n)
ShowScopedBody = deriveShow

||| a scoped term with names
public export
record Scoped (s : Nat) (f : Nat -> Type) (n : Nat) where
  constructor S
  names : NContext s
  body  : ScopedBody s f n
%name Scoped body

export %inline
(forall n. Eq (f n)) => Eq (Scoped s f n) where
  s == t = s.body == t.body

export %inline %hint
ShowScoped : (forall n. Show (f n)) => Show (Scoped s f n)
ShowScoped = deriveShow


infixl 8 :#
infixl 9 :@, :%
mutual
  public export
  TSubst : TSubstLike
  TSubst d = Subst $ \n => Elim d n

  ||| first argument `d` is dimension scope size;
  ||| second `n` is term scope size
  public export
  data Term : (d, n : Nat) -> Type where
    ||| type of types
    TYPE : (l : Universe) -> Term d n

    ||| function type
    Pi  : (qty : Qty) -> (arg : Term d n) ->
          (res : ScopeTerm d n) -> Term d n
    ||| function term
    Lam : (body : ScopeTerm d n) -> Term d n

    ||| pair type
    Sig : (fst : Term d n) -> (snd : ScopeTerm d n) -> Term d n
    ||| pair value
    Pair : (fst, snd : Term d n) -> Term d n

    ||| enumeration type
    Enum : (cases : SortedSet TagVal) -> Term d n
    ||| enumeration value
    Tag : (tag : TagVal) -> Term d n

    ||| equality type
    Eq : (ty : DScopeTerm d n) -> (l, r : Term d n) -> Term d n
    ||| equality term
    DLam : (body : DScopeTerm d n) -> Term d n

    ||| natural numbers (temporary until 𝐖 gets added)
    Nat  : Term d n
    -- [todo] can these be elims?
    Zero : Term d n
    Succ : (p : Term d n) -> Term d n

    ||| "box" (package a value up with a certain quantity)
    BOX : (qty : Qty) -> (ty : Term d n) -> Term d n
    Box : (val : Term d n) -> Term d n

    ||| elimination
    E : (e : Elim d n) -> Term d n

    ||| term closure/suspended substitution
    CloT  : WithSubst (Term d) (Elim d) n -> Term d n
    ||| dimension closure/suspended substitution
    DCloT : WithSubst (\d => Term d n) Dim d -> Term d n
  %name Term s, t, r

  ||| first argument `d` is dimension scope size, second `n` is term scope size
  public export
  data Elim : (d, n : Nat) -> Type where
    ||| free variable
    F : (x : Name) -> Elim d n
    ||| bound variable
    B : (i : Var n) -> Elim d n

    ||| term application
    (:@) : (fun : Elim d n) -> (arg : Term d n) -> Elim d n

    ||| pair destruction
    |||
    ||| `CasePair 𝜋 𝑒 ([𝑟], 𝐴) ([𝑥, 𝑦], 𝑡)` is
    ||| `𝐜𝐚𝐬𝐞 𝜋 · 𝑒 𝐫𝐞𝐭𝐮𝐫𝐧 𝑟 ⇒ 𝐴 𝐨𝐟 { (𝑥, 𝑦) ⇒ 𝑡 }`
    CasePair : (qty : Qty) -> (pair : Elim d n) ->
               (ret : ScopeTerm d n) ->
               (body : ScopeTermN 2 d n) ->
               Elim d n

    ||| enum matching
    CaseEnum : (qty : Qty) -> (tag : Elim d n) ->
               (ret : ScopeTerm d n) ->
               (arms : CaseEnumArms d n) ->
               Elim d n

    ||| nat matching
    CaseNat : (qty, qtyIH : Qty) -> (nat : Elim d n) ->
              (ret : ScopeTerm d n) ->
              (zero : Term d n) ->
              (succ : ScopeTermN 2 d n) ->
              Elim d n

    ||| unboxing
    CaseBox : (qty : Qty) -> (box : Elim d n) ->
              (ret : ScopeTerm d n) ->
              (body : ScopeTerm d n) ->
              Elim d n

    ||| dim application
    (:%) : (fun : Elim d n) -> (arg : Dim d) -> Elim d n

    ||| type-annotated term
    (:#) : (tm, ty : Term d n) -> Elim d n

    ||| coerce a value along a type equality, or show its coherence
    ||| [@xtt; §2.1.1]
    Coe : (ty : DScopeTerm d n) -> (p, q : Dim d) ->
          (val : Term d n) -> Elim d n

    ||| "generalised composition" [@xtt; §2.1.2]
    Comp : (ty : Term d n) -> (p, q : Dim d) ->
           (val : Term d n) -> (r : Dim d) ->
           (zero, one : DScopeTerm d n) -> Elim d n

    ||| match on types. needed for b.s. of coercions [@xtt; §2.2]
    TypeCase : (ty : Elim d n) -> (ret : Term d n) ->
               (arms : TypeCaseArms d n) -> (def : Term d n) ->
               Elim d n

    ||| term closure/suspended substitution
    CloE  : WithSubst (Elim d) (Elim d) n -> Elim d n
    ||| dimension closure/suspended substitution
    DCloE : WithSubst (\d => Elim d n) Dim d -> Elim d n
  %name Elim e, f

  public export
  CaseEnumArms : TermLike
  CaseEnumArms d n = SortedMap TagVal (Term d n)

  public export
  TypeCaseArms : TermLike
  TypeCaseArms d n = SortedDMap TyConKind (\k => TypeCaseArmBody k d n)

  public export
  TypeCaseArm : TermLike
  TypeCaseArm d n = (k ** TypeCaseArmBody k d n)

  public export
  TypeCaseArmBody : TyConKind -> TermLike
  TypeCaseArmBody k = ScopeTermN (arity k)


  public export
  ScopeTermN, DScopeTermN : Nat -> TermLike
  ScopeTermN  s d n = Scoped s (Term d) n
  DScopeTermN s d n = Scoped s (\d => Term d n) d

  public export
  ScopeTerm, DScopeTerm : TermLike
  ScopeTerm  = ScopeTermN 1
  DScopeTerm = DScopeTermN 1

mutual
  export %hint
  EqTerm : Eq (Term d n)
  EqTerm = assert_total {a = Eq (Term d n)} deriveEq

  export %hint
  EqElim : Eq (Elim d n)
  EqElim = assert_total {a = Eq (Elim d n)} deriveEq

mutual
  export %hint
  ShowTerm : Show (Term d n)
  ShowTerm = assert_total {a = Show (Term d n)} deriveShow

  export %hint
  ShowElim : Show (Elim d n)
  ShowElim = assert_total {a = Show (Elim d n)} deriveShow

||| scope which ignores all its binders
public export %inline
SN : {s : Nat} -> f n -> Scoped s f n
SN = S (replicate s Unused) . N

||| scope which uses its binders
public export %inline
SY : NContext s -> f (s + n) -> Scoped s f n
SY ns = S ns . Y

public export %inline
name : Scoped 1 f n -> BaseName
name (S [< x] _) = x

public export %inline
(.name) : Scoped 1 f n -> BaseName
s.name = name s

||| more convenient Pi
public export %inline
PiY : (qty : Qty) -> (x : BaseName) ->
      (arg : Term d n) -> (res : Term d (S n)) -> Term d n
PiY {qty, x, arg, res} = Pi {qty, arg, res = SY [< x] res}

||| non dependent function type
public export %inline
Arr : (qty : Qty) -> (arg, res : Term d n) -> Term d n
Arr {qty, arg, res} = Pi {qty, arg, res = SN res}

||| more convenient Sig
public export %inline
SigY : (x : BaseName) -> (fst : Term d n) ->
       (snd : Term d (S n)) -> Term d n
SigY {x, fst, snd} = Sig {fst, snd = SY [< x] snd}

||| non dependent pair type
public export %inline
And : (fst, snd : Term d n) -> Term d n
And {fst, snd} = Sig {fst, snd = SN snd}

||| more convenient Eq
public export %inline
EqY : (i : BaseName) -> (ty : Term (S d) n) ->
      (l, r : Term d n) -> Term d n
EqY {i, ty, l, r} = Eq {ty = SY [< i] ty, l, r}

||| non dependent equality type
public export %inline
Eq0 : (ty, l, r : Term d n) -> Term d n
Eq0 {ty, l, r} = Eq {ty = SN ty, l, r}

||| same as `F` but as a term
public export %inline
FT : Name -> Term d n
FT = E . F

||| abbreviation for a bound variable like `BV 4` instead of
||| `B (VS (VS (VS (VS VZ))))`
public export %inline
BV : (i : Nat) -> (0 _ : LT i n) => Elim d n
BV i = B $ V i

||| same as `BV` but as a term
public export %inline
BVT : (i : Nat) -> (0 _ : LT i n) => Term d n
BVT i = E $ BV i

public export
makeNat : Nat -> Term d n
makeNat 0 = Zero
makeNat (S k) = Succ $ makeNat k

public export %inline
enum : List TagVal -> Term d n
enum = Enum . SortedSet.fromList

public export %inline
typeCase : Elim d n -> Term d n ->
           List (TypeCaseArm d n) -> Term d n -> Elim d n
typeCase ty ret arms def = TypeCase ty ret (fromList arms) def

public export %inline
typeCase1Y : Elim d n -> Term d n ->
             (k : TyConKind) -> NContext (arity k) -> Term d (arity k + n) ->
             {default Nat def : Term d n} ->
             Elim d n
typeCase1Y ty ret k ns body {def} =
  typeCase ty ret [(k ** SY ns body)] def