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

%default total


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


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 : TermLike 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  : (tm : Term d from) -> (th : Lazy (TSubst d from to)) ->
            Term d to
    ||| dimension closure/suspended substitution
    DCloT : (tm : Term dfrom n) -> (th : Lazy (DSubst dfrom dto)) ->
            Term dto n

  ||| first argument `d` is dimension scope size, second `n` is term scope size
  public export
  data Elim : TermLike 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  : (el : Elim d from) -> (th : Lazy (TSubst d from to)) ->
            Elim d to
    ||| dimension closure/suspended substitution
    DCloE : (el : Elim dfrom n) -> (th : Lazy (DSubst dfrom dto)) ->
            Elim dto n

  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)


  ||| 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

  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

  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


%name Term s, t, r
%name Elim e, f
%name Scoped body
%name ScopedBody body

||| 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
Pi_ : (qty : Qty) -> (x : BaseName) ->
      (arg : Term d n) -> (res : Term d (S n)) -> Term d n
Pi_ {qty, x, arg, res} = Pi {qty, arg, res = S [< x] $ Y 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
Sig_ : (x : BaseName) -> (fst : Term d n) ->
       (snd : Term d (S n)) -> Term d n
Sig_ {x, fst, snd} = Sig {fst, snd = S [< x] $ Y 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
Eq_ : (i : BaseName) -> (ty : Term (S d) n) ->
      (l, r : Term d n) -> Term d n
Eq_ {i, ty, l, r} = Eq {ty = S [< i] $ Y 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
typeCase1 : 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
typeCase1 ty ret k ns body {def} =
  typeCase ty ret [(k ** SY ns body)] def