module Quox.Syntax.Term

import public Quox.Syntax.Var
import public Quox.Syntax.Shift
import public Quox.Syntax.Subst
import public Quox.Syntax.Universe
import public Quox.Syntax.Qty
import public Quox.Syntax.Dim
import public Quox.Name

import Quox.Pretty

import public Data.DPair
import Data.List
import Data.Maybe
import Data.Nat
import public Data.So
import Data.String

%default total


infixl 8 :#
infixl 9 :@
mutual
  public export
  TSubst : Nat -> Nat -> Nat -> Type
  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, qtm : Qty) -> (x : Name) ->
          (arg : Term d n) -> (res : Term d (S n)) -> Term d n
    ||| function term
    Lam : (x : Name) -> (body : Term d (S 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 : (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

    ||| type-annotated term
    (:#) : (tm, ty : 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

%name Term s, t, r
%name Elim e, f


||| 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) -> {auto 0 p : LT i n} -> Elim d n
BV i = B $ V i

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


infixl 9 :@@
||| apply multiple arguments at once
public export %inline
(:@@) : Elim d n -> List (Term d n) -> Elim d n
f :@@ ss = foldl (:@) f ss

private
getArgs' : Elim d n -> List (Term d n) -> (Elim d n, List (Term d n))
getArgs' (f :@ s) args = getArgs' f (s :: args)
getArgs' f        args = (f, args)

||| splits an application into its head and arguments. if it's not an
||| application then the list is just empty
export %inline
getArgs : Elim d n -> (Elim d n, List (Term d n))
getArgs e = getArgs' e []


parameters {auto _ : Pretty.HasEnv m}
  private %inline arrowD : m (Doc HL)
  arrowD = hlF Syntax $ ifUnicode "→" "->"

  private %inline lamD : m (Doc HL)
  lamD = hlF Syntax $ ifUnicode "λ" "fun"

  private %inline annD : m (Doc HL)
  annD = hlF Syntax $ ifUnicode "⦂" "::"

private %inline typeD : Doc HL
typeD = hl Syntax "Type"

private %inline colonD : Doc HL
colonD = hl Syntax ":"

mutual
  export covering
  PrettyHL (Term d n) where
    prettyM (TYPE l) =
      parensIfM App $ typeD <//> !(withPrec Arg $ prettyM l)
    prettyM (Pi qty qtm x s t) =
      parensIfM Outer $ hang 2 $
        !(prettyBinder [qty, qtm] x s) <++> !arrowD
          <//> !(under T x $ prettyM t)
    prettyM (Lam x t) =
      parensIfM Outer $
        sep [!lamD, hl TVar !(prettyM x), !arrowD]
          <//> !(under T x $ prettyM t)
    prettyM (E e) =
      prettyM e
    prettyM (CloT s th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM s) </> prettyTSubst th|]
    prettyM (DCloT s th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM s) </> prettyDSubst th|]

  export covering
  PrettyHL (Elim d n) where
    prettyM (F x) =
      hl' Free <$> prettyM x
    prettyM (B i) =
      prettyVar TVar TVarErr (!ask).tnames i
    prettyM (e :@ s) =
      let (f, args) = getArgs' e [s] in
      parensIfM App =<< withPrec Arg
        [|prettyM f <//> (align . sep <$> traverse prettyM args)|]
    prettyM (s :# a) =
      parensIfM Ann $ hang 2 $
        !(withPrec AnnL $ prettyM s) <++> !annD
          <//> !(withPrec Ann $ prettyM a)
    prettyM (CloE e th)  =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM e) </> prettyTSubst th|]
    prettyM (DCloE e th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM e) </> prettyDSubst th|]

  export covering
  prettyTSubst : Pretty.HasEnv m => TSubst d from to -> m (Doc HL)
  prettyTSubst s = prettySubstM prettyM (!ask).tnames TVar "[" "]" s

  export covering
  prettyBinder : Pretty.HasEnv m => List Qty -> Name -> Term d n -> m (Doc HL)
  prettyBinder pis x a =
    pure $ parens $ hang 2 $
      !(prettyQtyBinds pis) <//>
      hsep [hl TVar !(prettyM x), colonD, !(withPrec Outer $ prettyM a)]


export FromVar (Elim d) where fromVar = B
export FromVar (Term d) where fromVar = E . fromVar

||| does the minimal reasonable work:
||| - deletes the closure around a free name since it doesn't do anything
||| - deletes an identity substitution
||| - composes (lazily) with an existing top-level closure
||| - immediately looks up a bound variable
||| - otherwise, wraps in a new closure
export
CanSubst (Elim d) (Elim d) where
  F x       // _  = F x
  B i       // th = th !! i
  CloE e ph // th = assert_total CloE e $ ph . th
  e         // th = case force th of
                         Shift SZ => e
                         th       => CloE e th

||| does the minimal reasonable work:
||| - deletes the closure around an atomic constant like `TYPE`
||| - deletes an identity substitution
||| - composes (lazily) with an existing top-level closure
||| - goes inside `E` in case it is a simple variable or something
||| - otherwise, wraps in a new closure
export
CanSubst (Elim d) (Term d) where
  TYPE l    // _  = TYPE l
  E e       // th = E $ e // th
  CloT s ph // th = CloT s $ ph . th
  s         // th = case force th of
                         Shift SZ => s
                         th       => CloT s th

export CanSubst Var (Term d) where s // th = s // map (B {d}) th
export CanSubst Var (Elim d) where e // th = e // map (B {d}) th

export CanShift (Term d) where i // by = i // Shift by {env=(Elim d)}
export CanShift (Elim d) where i // by = i // Shift by {env=(Elim d)}


infixl 8 ///
mutual
  namespace Term
    ||| applies a dimension substitution with the same behaviour as `(//)`
    ||| above
    export
    (///) : Term dfrom n -> DSubst dfrom dto -> Term dto n
    TYPE l     /// _        = TYPE l
    E e        /// th       = E $ e /// th
    DCloT s ph /// th       = DCloT s $ ph . th
    s          /// Shift SZ = s
    s          /// th       = DCloT s th

    ||| applies a term and dimension substitution
    public export %inline
    subs : Term dfrom from -> DSubst dfrom dto -> TSubst dto from to ->
           Term dto to
    subs s th ph = s /// th // ph

  namespace Elim
    ||| applies a dimension substitution with the same behaviour as `(//)`
    ||| above
    export
    (///) : Elim dfrom n -> DSubst dfrom dto -> Elim dto n
    F x        /// _        = F x
    B i        /// _        = B i
    DCloE e ph /// th       = DCloE e $ ph . th
    e          /// Shift SZ = e
    e          /// th       = DCloE e th

    ||| applies a term and dimension substitution
    public export %inline
    subs : Elim dfrom from -> DSubst dfrom dto -> TSubst dto from to ->
           Elim dto to
    subs e th ph = e /// th // ph


private %inline
comp' : DSubst dfrom dto -> TSubst dfrom from mid -> TSubst dto mid to ->
        TSubst dto from to
comp' th ps ph = map (/// th) ps . ph


||| true if an elimination has a closure or dimension closure at the top level
public export %inline
topCloE : Elim d n -> Bool
topCloE (CloE  _ _) = True
topCloE (DCloE _ _) = True
topCloE _           = False

||| true if a term has a closure or dimension closure at the top level,
||| or is `E` applied to such an elimination
public export %inline
topCloT : Term d n -> Bool
topCloT (CloT  _ _) = True
topCloT (DCloT _ _) = True
topCloT (E e)       = topCloE e
topCloT _           = False

||| an elimination which is not a top level closure
public export NotCloElim : Nat -> Nat -> Type
NotCloElim d n = Subset (Elim d n) $ So . not . topCloE

||| a term which is not a top level closure
public export NotCloTerm : Nat -> Nat -> Type
NotCloTerm d n = Subset (Term d n) $ So . not . topCloT


mutual
  ||| if the input term has any top-level closures, push them under one layer of
  ||| syntax
  export
  pushSubstsT : Term d n -> NotCloTerm d n
  pushSubstsT s = pushSubstsT' id id s

  ||| if the input elimination has any top-level closures, push them under one
  ||| layer of syntax
  export
  pushSubstsE : Elim d n -> NotCloElim d n
  pushSubstsE e = pushSubstsE' id id e

  private
  pushSubstsT' : DSubst dfrom dto -> TSubst dto from to ->
                 Term dfrom from -> NotCloTerm dto to
  pushSubstsT' th ph (TYPE l) =
    Element (TYPE l) Oh
  pushSubstsT' th ph (Pi qty qtm x a b) =
    Element (Pi qty qtm x (subs a th ph) (subs b th (push ph))) Oh
  pushSubstsT' th ph (Lam x t) =
    Element (Lam x $ subs t th $ push ph) Oh
  pushSubstsT' th ph (E e) =
    let Element e prf = pushSubstsE' th ph e in Element (E e) prf
  pushSubstsT' th ph (CloT  s ps) =
    pushSubstsT' th (comp' th ps ph) s
  pushSubstsT' th ph (DCloT s ps) =
    pushSubstsT' (ps . th) ph s

  private
  pushSubstsE' : DSubst dfrom dto -> TSubst dto from to ->
                 Elim dfrom from -> NotCloElim dto to
  pushSubstsE' th ph (F x) =
    Element (F x) Oh
  pushSubstsE' th ph (B i) =
    assert_total pushSubstsE $ ph !! i
  pushSubstsE' th ph (f :@ s) =
    Element (subs f th ph :@ subs s th ph) Oh
  pushSubstsE' th ph (s :# a) =
    Element (subs s th ph :# subs a th ph) Oh
  pushSubstsE' th ph (CloE  e ps) =
    pushSubstsE' th (comp' th ps ph) e
  pushSubstsE' th ph (DCloE e ps) =
    pushSubstsE' (ps . th) ph e


parameters {auto _ : Alternative f}
  ||| `(λx. t ⦂ (x: A) → B) s >>> (t ⦂ B)[x ≔ (s ⦂ A)]`
  export
  betaLam1 : Elim d n -> f (Elim d n)
  betaLam1 ((Lam {body, _} :# Pi {arg, res, _}) :@ s) =
    pure $ (body :# res) // (s :# arg ::: id)
  betaLam1 _ = empty

  ||| `(e ⦂ A) >>> e`    [if `e` is an elim]
  export
  upsilon1 : Elim d n -> f (Elim d n)
  upsilon1 (E e :# _) = pure e
  upsilon1 _          = empty

  export
  step : Elim d n -> f (Elim d n)
  step e = betaLam1 e <|> upsilon1 e

export
step' : Elim d n -> Elim d n
step' e = fromMaybe e $ step e