module Quox.Syntax.Term.Base

import public Quox.Var
import public Quox.Scoped
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.Loc
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


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) -> (loc : Loc) -> Term d n

    ||| IO state token. this is a builtin because otherwise #[main] being a
    ||| builtin makes no sense
    IOState : (loc : Loc) -> Term d n

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

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

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

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

    ||| natural numbers (temporary until 𝐖 gets added)
    NAT : (loc : Loc) -> Term d n
    Nat : (val : Nat) -> (loc : Loc) -> Term d n
    Succ : (p : Term d n) -> (loc : Loc) -> Term d n

    ||| strings
    STRING : (loc : Loc) -> Term d n
    Str    : (str : String) -> (loc : Loc) -> Term d n

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

    Let : (qty : Qty) -> (rhs : Elim d n) ->
          (body : ScopeTerm d n) -> (loc : Loc) -> 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, possibly with a displacement (see @crude, or @mugen for a
    ||| more abstract and formalised take)
    |||
    ||| e.g. if f : ★₀ → ★₁, then f¹ : ★₁ → ★₂
    F : (x : Name) -> (u : Universe) -> (loc : Loc) -> Elim d n
    ||| bound variable
    B : (i : Var n) -> (loc : Loc) -> Elim d n

    ||| term application
    App : (fun : Elim d n) -> (arg : Term d n) -> (loc : Loc) -> Elim d n

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

    ||| first element of a pair. only works in non-linear contexts.
    Fst : (pair : Elim d n) -> (loc : Loc) -> Elim d n

    ||| second element of a pair. only works in non-linear contexts.
    Snd : (pair : Elim d n) -> (loc : Loc) -> Elim d n

    ||| enum matching
    CaseEnum : (qty : Qty) -> (tag : Elim d n) ->
               (ret : ScopeTerm d n) ->
               (arms : CaseEnumArms d n) ->
               (loc : Loc) ->
               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) ->
              (loc : Loc) ->
              Elim d n

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

    ||| dim application
    DApp : (fun : Elim d n) -> (arg : Dim d) -> (loc : Loc) -> Elim d n

    ||| type-annotated term
    Ann : (tm, ty : Term d n) -> (loc : Loc) -> 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) -> (loc : Loc) -> 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) -> (loc : Loc) -> 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) ->
               (loc : Loc) ->
               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 $ BN Unused noLoc) . N

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

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

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

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

||| more convenient Lam
public export %inline
LamY : (x : BindName) -> (body : Term d (S n)) -> (loc : Loc) -> Term d n
LamY {x, body, loc} = Lam {body = SY [< x] body, loc}

public export %inline
LamN : (body : Term d n) -> (loc : Loc) -> Term d n
LamN {body, loc} = Lam {body = SN body, loc}

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

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

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

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

||| more convenient DLam
public export %inline
DLamY : (i : BindName) -> (body : Term (S d) n) -> (loc : Loc) -> Term d n
DLamY {i, body, loc} = DLam {body = SY [< i] body, loc}

public export %inline
DLamN : (body : Term d n) -> (loc : Loc) -> Term d n
DLamN {body, loc} = DLam {body = SN body, loc}

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

||| same as `F` but as a term
public export %inline
FT : Name -> Universe -> Loc -> Term d n
FT x u loc = E $ F x u loc

||| same as `B` but as a term
public export %inline
BT : Var n -> (loc : Loc) -> Term d n
BT i loc = E $ B i loc

||| 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) => (loc : Loc) -> Elim d n
BV i loc = B (V i) loc

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

public export %inline
Zero : Loc -> Term d n
Zero = Nat 0

public export %inline
enum : List TagVal -> Loc -> Term d n
enum ts loc = Enum (SortedSet.fromList ts) loc

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

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


export
Located (Elim d n) where
  (F _ _ loc).loc               = loc
  (B _ loc).loc                 = loc
  (App _ _ loc).loc             = loc
  (CasePair _ _ _ _ loc).loc    = loc
  (Fst _ loc).loc               = loc
  (Snd _ loc).loc               = loc
  (CaseEnum _ _ _ _ loc).loc    = loc
  (CaseNat _ _ _ _ _ _ loc).loc = loc
  (CaseBox _ _ _ _ loc).loc     = loc
  (DApp _ _ loc).loc            = loc
  (Ann _ _ loc).loc             = loc
  (Coe _ _ _ _ loc).loc         = loc
  (Comp _ _ _ _ _ _ _ loc).loc  = loc
  (TypeCase _ _ _ _ loc).loc    = loc
  (CloE (Sub e _)).loc          = e.loc
  (DCloE (Sub e _)).loc         = e.loc

export
Located (Term d n) where
  (TYPE _ loc).loc      = loc
  (IOState loc).loc     = loc
  (Pi _ _ _ loc).loc    = loc
  (Lam _ loc).loc       = loc
  (Sig _ _ loc).loc     = loc
  (Pair _ _ loc).loc    = loc
  (Enum _ loc).loc      = loc
  (Tag _ loc).loc       = loc
  (Eq _ _ _ loc).loc    = loc
  (DLam _ loc).loc      = loc
  (NAT loc).loc         = loc
  (Nat _ loc).loc       = loc
  (STRING loc).loc      = loc
  (Str _ loc).loc       = loc
  (Succ _ loc).loc      = loc
  (BOX _ _ loc).loc     = loc
  (Box _ loc).loc       = loc
  (Let _ _ _ loc).loc   = loc
  (E e).loc             = e.loc
  (CloT (Sub t _)).loc  = t.loc
  (DCloT (Sub t _)).loc = t.loc

export
Located1 f => Located (ScopedBody s f n) where
  (Y t).loc = t.loc
  (N t).loc = t.loc

export
Located1 f => Located (Scoped s f n) where
  t.loc = t.body.loc


export
Relocatable (Elim d n) where
  setLoc loc (F x u _) = F x u loc
  setLoc loc (B i _) = B i loc
  setLoc loc (App fun arg _) = App fun arg loc
  setLoc loc (CasePair qty pair ret body _) =
    CasePair qty pair ret body loc
  setLoc loc (Fst pair _) = Fst pair loc
  setLoc loc (Snd pair _) = Fst pair loc
  setLoc loc (CaseEnum qty tag ret arms _) =
    CaseEnum qty tag ret arms loc
  setLoc loc (CaseNat qty qtyIH nat ret zero succ _) =
    CaseNat qty qtyIH nat ret zero succ loc
  setLoc loc (CaseBox qty box ret body _) =
    CaseBox qty box ret body loc
  setLoc loc (DApp fun arg _) =
    DApp fun arg loc
  setLoc loc (Ann tm ty _) =
    Ann tm ty loc
  setLoc loc (Coe ty p q val _) =
    Coe ty p q val loc
  setLoc loc (Comp ty p q val r zero one _) =
    Comp ty p q val r zero one loc
  setLoc loc (TypeCase ty ret arms def _) =
    TypeCase ty ret arms def loc
  setLoc loc (CloE (Sub term subst)) =
    CloE $ Sub (setLoc loc term) subst
  setLoc loc (DCloE (Sub term subst)) =
    DCloE $ Sub (setLoc loc term) subst

export
Relocatable (Term d n) where
  setLoc loc (TYPE l _) = TYPE l loc
  setLoc loc (IOState _) = IOState loc
  setLoc loc (Pi qty arg res _) = Pi qty arg res loc
  setLoc loc (Lam body _) = Lam body loc
  setLoc loc (Sig fst snd _) = Sig fst snd loc
  setLoc loc (Pair fst snd _) = Pair fst snd loc
  setLoc loc (Enum cases _) = Enum cases loc
  setLoc loc (Tag tag _) = Tag tag loc
  setLoc loc (Eq ty l r _) = Eq ty l r loc
  setLoc loc (DLam body _) = DLam body loc
  setLoc loc (NAT _) = NAT loc
  setLoc loc (Nat n _) = Nat n loc
  setLoc loc (Succ p _) = Succ p loc
  setLoc loc (STRING _) = STRING loc
  setLoc loc (Str s _) = Str s loc
  setLoc loc (BOX qty ty _) = BOX qty ty loc
  setLoc loc (Box val _) = Box val loc
  setLoc loc (Let qty rhs body _) = Let qty rhs body loc
  setLoc loc (E e) = E $ setLoc loc e
  setLoc loc (CloT (Sub term subst)) = CloT $ Sub (setLoc loc term) subst
  setLoc loc (DCloT (Sub term subst)) = DCloT $ Sub (setLoc loc term) subst

export
Relocatable1 f => Relocatable (ScopedBody s f n) where
  setLoc loc (Y body) = Y $ setLoc loc body
  setLoc loc (N body) = N $ setLoc loc body

export
Relocatable1 f => Relocatable (Scoped s f n) where
  setLoc loc (S names body) = S (setLoc loc <$> names) (setLoc loc body)