module Quox.Syntax.Term.Base

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


||| type-checkable terms, which consists of types and constructor forms.
|||
||| first argument `d` is dimension scope size; second `n` is term scope size
public export
data Term : (d, n : Nat) -> Type
%name Term s, t, r

||| inferrable terms, which consists of elimination forms like application and
||| `case` (as well as other terms with an annotation)
|||
||| first argument `d` is dimension scope size; second `n` is term scope size
public export
data Elim : (d, n : Nat) -> Type
%name Elim e, f


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

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

public export
ScopeTerm : TermLike
ScopeTerm = ScopeTermN 1

public export
DScopeTerm : TermLike
DScopeTerm = DScopeTermN 1


public export
TermT : TermLike
TermT = Thinned2 (\d, n => Term d n)

public export
ElimT : TermLike
ElimT = Thinned2 (\d, n => Elim d n)


public export
DimArg : TermLike
DimArg d n = Dim d


data Term where
  ||| type of types
  TYPE : (l : Universe) -> (loc : Loc) -> Term 0 0

  ||| function type
  Pi : Qty -> Subterms [Term, ScopeTerm] d n -> Loc -> Term d n
  ||| function value
  Lam : ScopeTerm d n -> Loc -> Term d n

  ||| pair type
  Sig : Subterms [Term, ScopeTerm] d n -> Loc -> Term d n
  ||| pair value
  Pair : Subterms [Term, Term] d n -> Loc -> Term d n

  ||| enum type
  Enum : List TagVal -> Loc -> Term 0 0
  ||| enum value
  Tag : TagVal -> Loc -> Term 0 0

  ||| equality type
  Eq : Subterms [DScopeTerm, Term, Term] d n -> Loc -> Term d n
  ||| equality value
  DLam : DScopeTerm d n -> Loc -> Term d n

  ||| natural numbers (temporary until 𝐖 gets added)
  Nat  : Loc -> Term 0 0
  Zero : Loc -> Term 0 0
  Succ : Term d n -> Loc -> Term 0 0

  ||| package a value with a quantity
  ||| e.g. a value of [ω. A], when unpacked, can be used ω times,
  ||| even if the box itself is linear
  BOX : Qty -> Term d n -> Loc -> Term d n
  Box : Term d n -> Loc -> Term d n

  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

||| first argument `d` is dimension scope size, second `n` is term scope size
public export
data Elim 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 : Name -> Universe -> Loc -> Elim 0 0
  ||| bound variable
  B : Loc -> Elim 0 1

  ||| term application
  App : Subterms [Elim, Term] d n -> Loc -> Elim d n

  ||| pair match
  ||| - the subterms are, in order: [head, return type, body]
  ||| - the quantity is that of the head, and since pairs only have one
  |||   constructor, can be 0
  CasePair : Qty -> Subterms [Elim, ScopeTerm, ScopeTermN 2] d n ->
             Loc -> Elim d n

  ||| enum match
  CaseEnum : Qty -> (arms : List TagVal) ->
             Subterms (Elim :: ScopeTerm :: (Term <$ arms)) d n ->
             Loc -> Elim d n

  ||| nat match
  CaseNat : Qty -> Qty ->
            Subterms [Elim, ScopeTerm, Term, ScopeTermN 2] d n ->
            Loc -> Elim d n

  ||| box match
  CaseBox : Qty -> Subterms [Elim, ScopeTerm, ScopeTerm] d n -> Loc -> Elim d n

  ||| dim application
  DApp : Subterms [Elim, DimArg] d n -> Loc -> Elim d n

  ||| type-annotated term
  Ann : Subterms [Term, Term] d n -> Loc -> Elim d n

  ||| coerce a value along a type equality, or show its coherence
  ||| [@xtt; §2.1.1]
  Coe : Subterms [DScopeTerm, DimArg, DimArg, Term] d n ->
        Loc -> Elim d n

  ||| "generalised composition" [@xtt; §2.1.2]
  Comp : Subterms [Term, DimArg, DimArg, Term,
                   DimArg, DScopeTerm, DScopeTerm] d n ->
         Loc -> Elim d n

  ||| match on types. needed for b.s. of coercions [@xtt; §2.2]
  TypeCase : Subterms [Elim, Term,   -- head, type
                       Term,         -- ★
                       ScopeTermN 2, -- pi
                       ScopeTermN 2, -- sig
                       Term,         -- enum
                       ScopeTermN 5, -- eq
                       Term,         -- nat
                       ScopeTerm     -- box
                      ] d n -> 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


-- this kills the idris ☹
-- export %hint
-- EqTerm : Eq (Term d n)

-- export %hint
-- EqElim : Eq (Elim d n)

-- EqTerm = deriveEq
-- EqElim = 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 0 0
FT x u loc = E $ F x u loc

||| abbreviation for a bound variable like `BV 4` instead of
||| `B (VS (VS (VS (VS VZ))))`
public export %inline
BV : (i : Fin n) -> (loc : Loc) -> ElimT d n
BV i loc = Th2 zero (one' i) $ B loc

||| same as `BV` but as a term
public export %inline
BVT : (i : Fin n) -> (loc : Loc) -> TermT d n
BVT i loc = Th2 zero (one' i) $ E $ B loc

public export
makeNat : Nat -> Loc -> Term 0 0
makeNat 0     loc = Zero loc
makeNat (S k) loc = Succ (makeNat k loc) loc


export
Located (Elim d n) where
  (F _ _ loc).loc          = loc
  (B loc).loc              = loc
  (App _ loc).loc          = loc
  (CasePair _ _ 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
  (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
  (Zero loc).loc        = loc
  (Succ _ loc).loc      = loc
  (BOX _ _ loc).loc     = loc
  (Box _ loc).loc       = loc
  (E e).loc             = e.loc
  (CloT (Sub t _)).loc  = t.loc
  (DCloT (Sub t _)).loc = t.loc


export
Relocatable (Elim d n) where
  setLoc loc (F x u _)                = F x u loc
  setLoc loc (B _)                    = B loc
  setLoc loc (App ts _)               = App ts loc
  setLoc loc (CasePair qty ts _)      = CasePair qty ts loc
  setLoc loc (CaseEnum qty arms ts _) = CaseEnum qty arms ts loc
  setLoc loc (CaseNat qty qtyIH ts _) = CaseNat qty qtyIH ts loc
  setLoc loc (CaseBox qty ts _)       = CaseBox qty ts loc
  setLoc loc (DApp ts _)              = DApp ts loc
  setLoc loc (Ann ts _)               = Ann ts loc
  setLoc loc (Coe ts _)               = Coe ts loc
  setLoc loc (Comp ts _)              = Comp ts loc
  setLoc loc (TypeCase ts _)          = TypeCase ts 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 (Pi qty ts _)            = Pi qty ts loc
  setLoc loc (Lam body _)             = Lam body loc
  setLoc loc (Sig ts _)               = Sig ts loc
  setLoc loc (Pair ts _)              = Pair ts loc
  setLoc loc (Enum cases _)           = Enum cases loc
  setLoc loc (Tag tag _)              = Tag tag loc
  setLoc loc (Eq ts _)                = Eq ts loc
  setLoc loc (DLam body _)            = DLam body loc
  setLoc loc (Nat _)                  = Nat loc
  setLoc loc (Zero _)                 = Zero loc
  setLoc loc (Succ p _)               = Succ p loc
  setLoc loc (BOX qty ty _)           = BOX qty ty loc
  setLoc loc (Box val _)              = Box val 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