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.Name
-- import public Quox.OPE

import Quox.Pretty

import public Data.DPair
import Data.List
import Data.Maybe
import Data.Nat
import public Data.So
import Data.String
import Data.Vect
import public Data.SortedMap
import public Data.SortedSet

%default total


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

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

public export
0 Universe : Type
Universe = Nat

public export
0 TagVal : Type
TagVal = String

infixl 8 :#
infixl 9 :@, :%
mutual
  public export
  0 TSubst : TSubstLike
  TSubst q d = Subst $ Elim q d

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

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

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

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

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

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

    ||| term closure/suspended substitution
    CloT  : (tm : Term q d from) -> (th : Lazy (TSubst q d from to)) ->
            Term q d to
    ||| dimension closure/suspended substitution
    DCloT : (tm : Term q dfrom n) -> (th : Lazy (DSubst dfrom dto)) ->
            Term q 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 q d n
    ||| bound variable
    B : (i : Var n) -> Elim q d n

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

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

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

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

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

    ||| term closure/suspended substitution
    CloE  : (el : Elim q d from) -> (th : Lazy (TSubst q d from to)) ->
            Elim q d to
    ||| dimension closure/suspended substitution
    DCloE : (el : Elim q dfrom n) -> (th : Lazy (DSubst dfrom dto)) ->
            Elim q dto n

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

  ||| a scoped term with names
  public export
  record Scoped (s : Nat) (f : Nat -> Type) (n : Nat) where
    constructor S
    names : Vect s BaseName
    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
  0 ScopeTermN, DScopeTermN : Nat -> TermLike
  ScopeTermN  s q d n = Scoped s (Term q d) n
  DScopeTermN s q d n = Scoped s (\d => Term q d n) d

  public export
  0 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 "_") . N

||| scope which uses its binders
public export %inline
SY : Vect s BaseName -> 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 : q) -> (x : BaseName) ->
      (arg : Term q d n) -> (res : Term q d (S n)) -> Term q d n
Pi_ {qty, x, arg, res} = Pi {qty, arg, res = S [x] $ Y res}

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

||| more convenient Sig
public export %inline
Sig_ : (x : BaseName) -> (fst : Term q d n) ->
       (snd : Term q d (S n)) -> Term q d n
Sig_ {x, fst, snd} = Sig {fst, snd = S [x] $ Y snd}

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

||| more convenient Eq
public export %inline
Eq_ : (i : BaseName) -> (ty : Term q (S d) n) ->
      (l, r : Term q d n) -> Term q 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 q d n) -> Term q 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 q 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 q 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 q d n
BVT i = E $ BV i