module Quox.Syntax.Term.Subst

import Quox.No
import Quox.Syntax.Term.Base
import Data.SnocVect

%default total

namespace CanDSubst
  public export
  interface CanDSubst (0 tm : TermLike) where
    (//) : tm dfrom n -> Lazy (DSubst dfrom dto) -> tm dto n

||| 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 dim-closure
||| - otherwise, wraps in a new closure
export
CanDSubst Term where
  s          // Shift SZ = s
  TYPE l     // _        = TYPE l
  DCloT s ph // th       = DCloT s $ ph . th
  s          // th       = DCloT s th

private
subDArgs : Elim dfrom n -> DSubst dfrom dto -> Elim dto n
subDArgs (f :% d) th = subDArgs f th :% (d // th)
subDArgs e        th = DCloE e th

||| does the minimal reasonable work:
||| - deletes the closure around a term variable
||| - deletes an identity substitution
||| - composes (lazily) with an existing top-level dim-closure
||| - immediately looks up bound variables in a
|||   top-level sequence of dimension applications
||| - otherwise, wraps in a new closure
export
CanDSubst Elim where
  e          // Shift SZ = e
  F x        // _        = F x
  B i        // _        = B i
  f :% d     // th       = subDArgs (f :% d) th
  DCloE e ph // th       = DCloE e $ ph . th
  e          // th       = DCloE e th

namespace DSubst.ScopeTermN
  export %inline
  (//) : ScopeTermN s dfrom n -> Lazy (DSubst dfrom dto) ->
         ScopeTermN s dto n
  S ns (Y body) // th = S ns $ Y $ body // th
  S ns (N body) // th = S ns $ N $ body // th

namespace DSubst.DScopeTermN
  export %inline
  (//) : {s : Nat} ->
         DScopeTermN s dfrom n -> Lazy (DSubst dfrom dto) ->
         DScopeTermN s dto n
  S ns (Y body) // th = S ns $ Y $ body // pushN s th
  S ns (N body) // th = S ns $ N $ body // th


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


||| does the minimal reasonable work:
||| - deletes the closure around a *free* name
||| - 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
CanSubstSelf (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

namespace CanTSubst
  public export
  interface CanTSubst (0 tm : TermLike) where
    (//) : tm d from -> Lazy (TSubst d from to) -> tm d to

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

namespace ScopeTermN
  export %inline
  (//) : {s : Nat} ->
         ScopeTermN s d from -> Lazy (TSubst d from to) ->
         ScopeTermN s d to
  S ns (Y body) // th = S ns $ Y $ body // pushN s th
  S ns (N body) // th = S ns $ N $ body // th

namespace DScopeTermN
  export %inline
  (//) : {s : Nat} ->
         DScopeTermN s d from -> Lazy (TSubst d from to) ->
         DScopeTermN s d to
  S ns (Y body) // th = S ns $ Y $ body // map (// shift s) th
  S ns (N body) // th = S ns $ N $ body // th

export %inline CanShift (Term d) where s // by = s // Shift by
export %inline CanShift (Elim d) where e // by = e // Shift by

export %inline
{s : Nat} -> CanShift (ScopeTermN s d) where
  b // by = b // Shift by


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


public export %inline
dweakT : (by : Nat) -> Term d n -> Term (by + d) n
dweakT by t = t // shift by

public export %inline
dweakE : (by : Nat) -> Elim d n -> Elim (by + d) n
dweakE by t = t // shift by


public export %inline
weakT : (by : Nat) -> Term d n -> Term d (by + n)
weakT by t = t // shift by

public export %inline
weakE : (by : Nat) -> Elim d n -> Elim d (by + n)
weakE by t = t // shift by


parameters {s : Nat}
  namespace ScopeTermBody
    export %inline
    (.term) : ScopedBody s (Term d) n -> Term d (s + n)
    (Y b).term = b
    (N b).term = weakT s b

  namespace ScopeTermN
    export %inline
    (.term) : ScopeTermN s d n -> Term d (s + n)
    t.term = t.body.term

  namespace DScopeTermBody
    export %inline
    (.term) : ScopedBody s (\d => Term d n) d -> Term (s + d) n
    (Y b).term = b
    (N b).term = dweakT s b

  namespace DScopeTermN
    export %inline
    (.term) : DScopeTermN s d n -> Term (s + d) n
    t.term = t.body.term


export %inline
subN : ScopeTermN s d n -> SnocVect s (Elim d n) -> Term d n
subN (S _ (Y body)) es = body // fromSnocVect es
subN (S _ (N body)) _  = body

export %inline
sub1 : ScopeTerm d n -> Elim d n -> Term d n
sub1 t e = subN t [< e]

export %inline
dsubN : DScopeTermN s d n -> SnocVect s (Dim d) -> Term d n
dsubN (S _ (Y body)) ps = body // fromSnocVect ps
dsubN (S _ (N body)) _  = body

export %inline
dsub1 : DScopeTerm d n -> Dim d -> Term d n
dsub1 t p = dsubN t [< p]


public export %inline
(.zero) : DScopeTerm d n -> Term d n
body.zero = dsub1 body $ K Zero

public export %inline
(.one) : DScopeTerm d n -> Term d n
body.one = dsub1 body $ K One


public export
0 CloTest : TermLike -> Type
CloTest tm = forall d, n. tm d n -> Bool

interface PushSubsts (0 tm : TermLike) (0 isClo : CloTest tm) | tm where
  pushSubstsWith : DSubst dfrom dto -> TSubst dto from to ->
                   tm dfrom from -> Subset (tm dto to) (No . isClo)

public export
0 NotClo : {isClo : CloTest tm} -> PushSubsts tm isClo => Pred (tm d n)
NotClo = No . isClo

public export
0 NonClo : (tm : TermLike) -> {isClo : CloTest tm} ->
           PushSubsts tm isClo => TermLike
NonClo tm d n = Subset (tm d n) NotClo

public export %inline
nclo : {isClo : CloTest tm} -> (0 _ : PushSubsts tm isClo) =>
       (t : tm d n) -> (0 nc : NotClo t) => NonClo tm d n
nclo t = Element t nc

parameters {0 isClo : CloTest tm} {auto _ : PushSubsts tm isClo}
  ||| if the input term has any top-level closures, push them under one layer of
  ||| syntax
  export %inline
  pushSubsts : tm d n -> NonClo tm d n
  pushSubsts s = pushSubstsWith id id s

  export %inline
  pushSubstsWith' : DSubst dfrom dto -> TSubst dto from to ->
                    tm dfrom from -> tm dto to
  pushSubstsWith' th ph x = fst $ pushSubstsWith th ph x

mutual
  public export
  isCloT : CloTest Term
  isCloT (CloT  {}) = True
  isCloT (DCloT {}) = True
  isCloT (E e)      = isCloE e
  isCloT _          = False

  public export
  isCloE : CloTest Elim
  isCloE (CloE  {}) = True
  isCloE (DCloE {}) = True
  isCloE _          = False

mutual
  export
  PushSubsts Term Subst.isCloT where
    pushSubstsWith th ph (TYPE l) =
      nclo $ TYPE l
    pushSubstsWith th ph (Pi qty a body) =
      nclo $ Pi qty (a // th // ph) (body // th // ph)
    pushSubstsWith th ph (Lam body) =
      nclo $ Lam (body // th // ph)
    pushSubstsWith th ph (Sig a b) =
      nclo $ Sig (a // th // ph) (b // th // ph)
    pushSubstsWith th ph (Pair s t) =
      nclo $ Pair (s // th // ph) (t // th // ph)
    pushSubstsWith th ph (Enum tags) =
      nclo $ Enum tags
    pushSubstsWith th ph (Tag tag) =
      nclo $ Tag tag
    pushSubstsWith th ph (Eq ty l r) =
      nclo $ Eq (ty // th // ph) (l // th // ph) (r // th // ph)
    pushSubstsWith th ph (DLam body) =
      nclo $ DLam (body // th // ph)
    pushSubstsWith _ _ Nat = nclo Nat
    pushSubstsWith _ _ Zero = nclo Zero
    pushSubstsWith th ph (Succ n) = nclo $ Succ $ n // th // ph
    pushSubstsWith th ph (BOX pi ty) = nclo $ BOX pi $ ty // th // ph
    pushSubstsWith th ph (Box val) = nclo $ Box $ val // th // ph
    pushSubstsWith th ph (E e) =
      let Element e nc = pushSubstsWith th ph e in nclo $ E e
    pushSubstsWith th ph (CloT s ps) =
      pushSubstsWith th (comp th ps ph) s
    pushSubstsWith th ph (DCloT s ps) =
      pushSubstsWith (ps . th) ph s

  export
  PushSubsts Elim Subst.isCloE where
    pushSubstsWith th ph (F x) =
      nclo $ F x
    pushSubstsWith th ph (B i) =
      let res = ph !! i in
      case nchoose $ isCloE res of
        Left  yes => assert_total pushSubsts res
        Right no  => Element res no
    pushSubstsWith th ph (f :@ s) =
      nclo $ (f // th // ph) :@ (s // th // ph)
    pushSubstsWith th ph (CasePair pi p r b) =
      nclo $ CasePair pi (p // th // ph) (r // th // ph) (b // th // ph)
    pushSubstsWith th ph (CaseEnum pi t r arms) =
      nclo $ CaseEnum pi (t // th // ph) (r // th // ph)
               (map (\b => b // th // ph) arms)
    pushSubstsWith th ph (CaseNat pi pi' n r z s) =
      nclo $ CaseNat pi pi' (n // th // ph) (r // th // ph)
               (z // th // ph) (s // th // ph)
    pushSubstsWith th ph (CaseBox pi x r b) =
      nclo $ CaseBox pi (x // th // ph) (r // th // ph) (b // th // ph)
    pushSubstsWith th ph (f :% d) =
      nclo $ (f // th // ph) :% (d // th)
    pushSubstsWith th ph (s :# a) =
      nclo $ (s // th // ph) :# (a // th // ph)
    pushSubstsWith th ph (Coe ty p q val) =
      nclo $ Coe (ty // th // ph) (p // th) (q // th) (val // th // ph)
    pushSubstsWith th ph (Comp ty p q val r zero one) =
      nclo $ Comp (ty // th // ph) (p // th) (q // th)
                  (val // th // ph) (r // th)
                  (zero // th // ph) (one // th // ph)
    pushSubstsWith th ph (TypeCase ty ret arms def) =
      nclo $ TypeCase (ty // th // ph) (ret  // th // ph)
        (map (\t => t // th // ph) arms) (def // th // ph)
    pushSubstsWith th ph (CloE e ps) =
      pushSubstsWith th (comp th ps ph) e
    pushSubstsWith th ph (DCloE e ps) =
      pushSubstsWith (ps . th) ph e


||| heterogeneous composition, using Comp and Coe (and subst)
|||
||| comp [i ⇒ A] @p @q s { (r=0) j ⇒ t₀; (r=1) j ⇒ t₁ }
||| ≔
||| comp [A‹q/i›] @p @q (coe [i ⇒ A] @p @q s) {
|||   (r=0) j ⇒ coe [i ⇒ A] @j @q t₀;
|||   (r=1) j ⇒ coe [i ⇒ A] @j @q t₁
||| }
public export %inline
CompH : (ty : DScopeTerm d n) -> (p, q : Dim d) -> (val : Term d n) ->
        (r : Dim d) -> (zero, one : DScopeTerm d n) -> Elim d n
CompH {ty, p, q, val, r, zero, one} =
  let ty' = SY ty.names $ ty.term // (B VZ ::: shift 2) in
  Comp {
    ty = dsub1 ty q, p, q,
    val = E $ Coe ty p q val, r,
    zero = SY zero.names $ E $ Coe ty' (B VZ) (weakD 1 q) zero.term,
    one  = SY one.names  $ E $ Coe ty' (B VZ) (weakD 1 q) one.term
  }