module Core.Term.Substitution

import Core.Term
import Core.Term.Thinned
import Core.Thinning

infix 4 =~=

-- Definition ------------------------------------------------------------------

public export
data Subst : Nat -> Nat -> Type where
  Wkn : m `Thins` n -> Subst m n
  (:<) : Subst m n -> Thinned n -> Subst (S m) n

namespace Eq
  public export
  data (=~=) : Subst m n -> Subst m n -> Type where
    Refl : sub =~= sub
    (:<) : sub1 =~= sub2 -> t =~= u -> sub1 :< t =~= sub2 :< u

%name Subst sub
%name Eq.(=~=) prf

-- Indexing --------------------------------------------------------------------

public export
index : Subst m n -> Fin m -> Thinned n
index (Wkn thin) i = Var i `Over` thin
index (sub :< t) FZ = t
index (sub :< t) (FS i) = index sub i

export
indexCong :
  {0 sub1, sub2 : Subst m n} ->
  sub1 =~= sub2 ->
  (i : Fin m) ->
  index sub1 i =~= index sub2 i
indexCong Refl i = Refl
indexCong (prf :< prf') FZ = prf'
indexCong (prf :< prf') (FS i) = indexCong prf i

-- Weakening -------------------------------------------------------------------

public export
wkn : Subst k m -> m `Thins` n -> Subst k n
wkn (Wkn thin') thin = Wkn (thin . thin')
wkn (sub :< t) thin = wkn sub thin :< wkn t thin

export
wknCong :
  {0 sub1, sub2 : Subst k m} ->
  sub1 =~= sub2 ->
  (thin : m `Thins` n) ->
  wkn sub1 thin =~= wkn sub2 thin
wknCong Refl thin = Refl
wknCong (prf :< prf') thin = wknCong prf thin :< wknCong prf' thin

-- Constructors ----------------------------------------------------------------

public export
lift : Subst m n -> Subst (S m) (S n)
lift sub = wkn sub (drop $ id n) :< (Var 0 `Over` id (S n))

-- Substitution ----------------------------------------------------------------

public export
subst : Term m -> Subst m n -> Term n
subst (Var i) sub = expand $ index sub i
subst Set sub = Set
subst (Pi t u) sub = Pi (subst t sub) (subst u $ lift sub)
subst (Abs t) sub = Abs (subst t $ lift sub)
subst (App t u) sub = App (subst t sub) (subst u sub)

export
substCong :
  (t : Term m) ->
  {0 sub1, sub2 : Subst m n} ->
  sub1 =~= sub2 ->
  subst t sub1 = subst t sub2
substCong (Var i) prf = indexCong prf i
substCong Set prf = Refl
substCong (Pi t u) prf =
  cong2 Pi
    (substCong t prf)
    (substCong u $ wknCong prf (drop $ id n) :< Refl)
substCong (Abs t) prf = cong Abs (substCong t $ wknCong prf (drop $ id n) :< Refl)
substCong (App t u) prf = cong2 App (substCong t prf) (substCong u prf)