module Core.Term.Substitution

import Core.Context
import Core.Var
import Core.Term
import Core.Thinning

import Syntax.PreorderReasoning

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

public export
data Subst : Context -> Context -> Type where
  Base : sx `Thins` sy -> Subst sx sy
  (:<) : Subst sx sy -> Thinned sy -> Subst (sx :< x) sy

%name Subst sub

-- Equality --------------------------------------------------------------------

namespace Eq
  public export
  data SubstEq : Subst sx sy -> Subst sx sy -> Type where
    Base : SubstEq env env
    (:<) : SubstEq sub1 sub2 -> expand t = expand u -> SubstEq (sub1 :< t) (sub2 :< u)

  %name SubstEq prf

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

export
index : Subst sx sy -> Var sx -> Term sy
doIndex : Subst sx sy -> Thinned sy -> {0 i : Var (sx :< n)} -> View i -> Term sy

index (Base thin) i = Var $ wkn i thin
index (sub :< t) i = doIndex sub t (view i)

doIndex sub t Here = expand t
doIndex sub t (There i) = index sub i

-- Preserves Equality

export
indexCong : SubstEq sub1 sub2 -> (i : Var sx) -> index sub1 i = index sub2 i
doIndexCong :
  SubstEq sub1 sub2 ->
  expand t = expand u ->
  {0 i : Var (sx :< n)} ->
  (v : View i) ->
  doIndex sub1 t v = doIndex sub2 u v

indexCong Base i = Refl
indexCong (prf :< eq) i = doIndexCong prf eq (view i)

doIndexCong prf eq Here = eq
doIndexCong prf eq (There i) = indexCong prf i

-- Restriction -----------------------------------------------------------------

export
restrict : sx `Thins` sy -> Subst sy sz -> Subst sx sz
doRestrict :
  {0 thin : sx `Thins` sy :< n} ->
  View thin ->
  Subst sy sz ->
  Thinned sz ->
  Subst sx sz

restrict thin (Base thin') = Base (thin' . thin)
restrict thin (sub :< t) = doRestrict (view thin) sub t

doRestrict (Id _) sub t = sub :< t
doRestrict (Drop thin n) sub t = restrict thin sub
doRestrict (Keep thin n) sub t = restrict thin sub :< t

-- Basic Simplification

export
restrictBase :
  (thin1 : sx `Thins` sy) ->
  (thin2 : sy `Thins` sz) ->
  restrict thin1 (Base thin2) = Base (thin2 . thin1)
restrictBase thin1 thin2 = Refl

export
restrictId :
  (sub : Subst sx sy) ->
  restrict (id sx) sub = sub
restrictId (Base thin) = cong Base $ compRightIdentity thin
restrictId {sx = sx :< n} (sub :< t) =
  rewrite viewUnique (view $ id _) (Id $ sx :< n) in
  Refl

doRestrictKeep :
  {0 thin : sx `Thins` sy} ->
  (v : View thin) ->
  (0 n : Name) ->
  (sub : Subst sy sz) ->
  (t : Thinned sz) ->
  restrict (keep thin n) (sub :< t) = restrict thin sub :< t
doRestrictKeep (Id sy) n sub t =
  rewrite keepIdIsId sy n in
  rewrite viewUnique (view $ id (sy :< n)) (Id _) in
  sym $ cong (:< t) $ restrictId sub
doRestrictKeep (Drop thin m) n sub t =
  rewrite viewUnique (view $ keep (drop thin m) n) (Keep {prf = dropIsNotId thin m} _ n) in
  Refl
doRestrictKeep (Keep {prf} thin m) n sub t =
  rewrite viewUnique (view $ keep (keep thin m) n) (Keep {prf = keepNotIdIsNotId prf m} _ n) in
  Refl

export
restrictKeep :
  (thin : sx `Thins` sy) ->
  (0 n : Name) ->
  (sub : Subst sy sz) ->
  (t : Thinned sz) ->
  restrict (keep thin n) (sub :< t) = restrict thin sub :< t
restrictKeep thin n sub t = doRestrictKeep (view thin) n sub t

-- Interaction with Indexing

export
indexRestrict :
  (thin : sx `Thins` sy) ->
  (sub : Subst sy sz) ->
  (i : Var sx) ->
  index (restrict thin sub) i = index sub (wkn i thin)
doIndexRestrict :
  {0 thin : sx `Thins` sy :< n} ->
  (view : View thin) ->
  (sub : Subst sy sz) ->
  (t : Thinned sz) ->
  (i : Var sx) ->
  index (doRestrict view sub t) i = index (sub :< t) (wkn i thin)
doDoIndexRestrict :
  (thin : sx `Thins` sy) ->
  (0 n : Name) ->
  (sub : Subst sy sz) ->
  (t : Thinned sz) ->
  {0 i : Var (sx :< n)} ->
  (v : View i) ->
  doIndex (restrict thin sub) t v = doIndex sub t (view $ wkn i $ keep thin n)

indexRestrict thin (Base thin') i = sym $ cong Var $ wknHomo i thin thin'
indexRestrict thin (sub :< t) i = doIndexRestrict (view thin) sub t i

doIndexRestrict (Id _) sub t i = rewrite wknId i in Refl
doIndexRestrict (Drop thin n) sub t i =
  rewrite wknDropIsThere i thin n in
  rewrite viewUnique (view $ there {n} $ wkn i thin) (There $ wkn i thin) in
  indexRestrict thin sub i
doIndexRestrict (Keep thin n) sub t i = doDoIndexRestrict thin n sub t (view i)

doDoIndexRestrict thin n sub t Here =
  rewrite wknKeepHereIsHere thin n in
  sym $ cong (doIndex {n} sub t) $ viewUnique (view here) Here
doDoIndexRestrict thin n sub t (There i) =
  rewrite wknKeepThereIsThere i thin n in
  rewrite viewUnique (view $ there {n} $ wkn i thin) (There $ wkn i thin) in
  indexRestrict thin sub i

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

public export
wkn : Subst sx sy -> sy `Thins` sz -> Subst sx sz
wkn (Base thin') thin = Base (thin . thin')
wkn (sub :< t) thin = wkn sub thin :< shift t thin

-- Preserves Equality

export
wknCong : SubstEq sub1 sub2 -> (thin : sx `Thins` sy) -> SubstEq (wkn sub1 thin) (wkn sub2 thin)
wknCong Base thin = Base
wknCong ((:<) {t, u} prf eq) thin =
  (:<) (wknCong prf thin) $
  Calc $
  |~ expand (shift t thin)
  ~~ wkn (expand t) thin   ...(expandShift t thin)
  ~~ wkn (expand u) thin   ...(cong (\t => wkn t thin) eq)
  ~~ expand (shift u thin) ...(sym $ expandShift u thin)

-- Is Homomorphic

export
wknHomo :
  (sub : Subst sx sy) ->
  (thin1 : sy `Thins` sz) ->
  (thin2 : sz `Thins` sw) ->
  wkn (wkn sub thin1) thin2 = wkn sub (thin2 . thin1)
wknHomo (Base thin) thin1 thin2 = cong Base $ compAssoc thin2 thin1 thin
wknHomo (sub :< t) thin1 thin2 = cong2 (:<) (wknHomo sub thin1 thin2) (shiftHomo t thin1 thin2)

-- Interaction with Indexing

export
wknIndex :
  (sub : Subst sx sy) ->
  (thin : sy `Thins` sz) ->
  (i : Var sx) ->
  wkn (index sub i) thin = index (wkn sub thin) i
wknDoIndex :
  (sub : Subst sx sy) ->
  (t : Thinned sy) ->
  (thin : sy `Thins` sz) ->
  {0 i : Var (sx :< n)} ->
  (v : View i) ->
  wkn (doIndex sub t v) thin = doIndex (wkn sub thin) (shift t thin) v

wknIndex (Base thin') thin i = cong Var $ wknHomo i thin' thin
wknIndex (sub :< t) thin i = wknDoIndex sub t thin (view i)

wknDoIndex sub t thin Here = sym $ expandShift t thin
wknDoIndex sub t thin (There i) = wknIndex sub thin i

-- Interaction with Restriction

export
wknRestrictCommute :
  (thin1 : sx `Thins` sy) ->
  (sub : Subst sy sz) ->
  (thin2 : sz `Thins` sw) ->
  wkn (restrict thin1 sub) thin2 = restrict thin1 (wkn sub thin2)
wknDoRestrictCommute :
  {0 thin1 : sx `Thins` sy :< n} ->
  (v : View thin1) ->
  (sub : Subst sy sz) ->
  (t : Thinned sz) ->
  (thin2 : sz `Thins` sw) ->
  wkn (doRestrict v sub t) thin2 = doRestrict v (wkn sub thin2) (shift t thin2)

wknRestrictCommute thin1 (Base thin) thin2 = cong Base $ compAssoc thin2 thin thin1
wknRestrictCommute thin1 (sub :< t) thin2 = wknDoRestrictCommute (view thin1) sub t thin2

wknDoRestrictCommute (Id _) sub t thin2 = Refl
wknDoRestrictCommute (Drop thin1 n) sub t thin2 = wknRestrictCommute thin1 sub thin2
wknDoRestrictCommute (Keep thin1 n) sub t thin2 =
  cong (:< shift t thin2) $ wknRestrictCommute thin1 sub thin2

-- Lifting ---------------------------------------------------------------------

public export
lift : Subst sx sy -> (0 n : Name) -> Subst (sx :< n) (sy :< n)
lift sub n = wkn sub (wkn1 sy n) :< (Var here `Over` id (sy :< n))

-- Preserves Equality

export
liftCong : SubstEq sub1 sub2 -> (0 n : Name) -> SubstEq (lift sub1 n) (lift sub2 n)
liftCong prf n = wknCong prf (wkn1 _ n) :< Refl

-- Interaction with Restriction

export
restrictLift :
  (thin : sx `Thins` sy) ->
  (sub : Subst sy sz) ->
  (0 n : Name) ->
  restrict (keep thin n) (lift sub n) = lift (restrict thin sub) n
restrictLift thin sub n = Calc $
  |~ restrict (keep thin n) (wkn sub (wkn1 sz n) :< (Var here `Over` id _))
  ~~ restrict thin (wkn sub (wkn1 sz n)) :< (Var here `Over` id _)          ...(restrictKeep thin n (wkn sub (wkn1 sz n)) (Var here `Over` id _))
  ~~ wkn (restrict thin sub) (wkn1 sz n) :< (Var here `Over` id _)          ...(sym $ cong (:< Over (Var here) (id (sz :< n))) $ wknRestrictCommute thin sub (wkn1 sz n))

-- Interaction with Weakening

export
wknLift :
  (sub : Subst sx sy) ->
  (thin : sy `Thins` sz) ->
  (0 n : Name) ->
  SubstEq (wkn (lift sub n) (keep thin n)) (lift (wkn sub thin) n)
wknLift sub thin n =
  let
    eq1 = Calc $
      |~ wkn (wkn sub (wkn1 sy n)) (keep thin n)
      ~~ wkn sub (keep thin n . wkn1 sy n)       ...(wknHomo sub (wkn1 sy n) (keep thin n))
      ~~ wkn sub (drop thin n)                   ...(cong (wkn sub) $ compRightWkn1 thin n)
      ~~ wkn sub (wkn1 sz n . thin)              ...(sym $ cong (wkn sub) $ compLeftWkn1 thin n)
      ~~ wkn (wkn sub thin) (wkn1 sz n)          ...(sym $ wknHomo sub thin (wkn1 sz n))
  in
  let
    eq2 = Calc $
      |~ wkn here (keep thin n . id (sy :< n))
      ~~ wkn here (keep thin n)                ...(cong (wkn here) $ compRightIdentity (keep thin n))
      ~~ here                                  ...(wknKeepHereIsHere thin n)
      ~~ wkn here (id (sz :< n))               ...(sym $ wknId here)
  in
  (rewrite eq1 in Base) :< cong Var eq2

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

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

-- Substitutions are extensional

export
substCong :
  (t : Term sx) ->
  {0 sub1, sub2 : Subst sx sy} ->
  SubstEq sub1 sub2 ->
  subst t sub1 = subst t sub2
substCong (Var i) prf = indexCong prf i
substCong Set prf = Refl
substCong (Pi n t u) prf = cong2 (Pi n) (substCong t prf) (substCong u $ liftCong prf n)
substCong (Abs n t) prf = cong (Abs n) (substCong t $ liftCong prf n)
substCong (App t u) prf = cong2 App (substCong t prf) (substCong u prf)

-- Substitution Commutes with Weakening

export
wknSubst :
  (t : Term sx) ->
  (sub : Subst sx sy) ->
  (thin : sy `Thins` sz) ->
  wkn (subst t sub) thin = subst t (wkn sub thin)
wknSubst (Var i) sub thin = wknIndex sub thin i
wknSubst Set sub thin = Refl
wknSubst (Pi n t u) sub thin =
  cong2 (Pi n) (wknSubst t sub thin) $
  Calc $
  |~ wkn (subst u $ lift sub n) (keep thin n)
  ~~ subst u (wkn (lift sub n) (keep thin n)) ...(wknSubst u (lift sub n) (keep thin n))
  ~~ subst u (lift (wkn sub thin) n)          ...(substCong u $ wknLift sub thin n)
wknSubst (Abs n t) sub thin =
  cong (Abs n) $
  Calc $
  |~ wkn (subst t $ lift sub n) (keep thin n)
  ~~ subst t (wkn (lift sub n) (keep thin n)) ...(wknSubst t (lift sub n) (keep thin n))
  ~~ subst t (lift (wkn sub thin) n)          ...(substCong t $ wknLift sub thin n)
wknSubst (App t u) sub thin = cong2 App (wknSubst t sub thin) (wknSubst u sub thin)

export
substWkn :
  (t : Term sx) ->
  (thin : sx `Thins` sy) ->
  (sub : Subst sy sz) ->
  subst (wkn t thin) sub = subst t (restrict thin sub)
substWkn (Var i) thin sub = sym $ indexRestrict thin sub i
substWkn Set thin sub = Refl
substWkn (Pi n t u) thin sub =
  cong2 (Pi n) (substWkn t thin sub) $
  Calc $
  |~ subst (wkn u (keep thin n)) (lift sub n)
  ~~ subst u (restrict (keep thin n) $ lift sub n) ...(substWkn u (keep thin n) (lift sub n))
  ~~ subst u (lift (restrict thin sub) n)          ...(cong (subst u) $ restrictLift thin sub n)
substWkn (Abs n t) thin sub =
  cong (Abs n) $
  Calc $
  |~ subst (wkn t (keep thin n)) (lift sub n)
  ~~ subst t (restrict (keep thin n) $ lift sub n) ...(substWkn t (keep thin n) (lift sub n))
  ~~ subst t (lift (restrict thin sub) n)          ...(cong (subst t) $ restrictLift thin sub n)
substWkn (App t u) thin sub = cong2 App (substWkn t thin sub) (substWkn u thin sub)

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

public export
sub1 : Term sx -> Subst (sx :< n) sx
sub1 t = Base (id sx) :< (t `Over` id sx)