module Core.Term.Substitution

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

import Syntax.PreorderReasoning

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

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

export
wknHomo :
  (sub : Subst j k) ->
  (thin1 : k `Thins` m) ->
  (thin2 : m `Thins` n) ->
  wkn (wkn sub thin1) thin2 = wkn sub (thin2 . thin1)
wknHomo (Wkn thin) thin1 thin2 = cong Wkn (compAssoc thin2 thin1 thin)
wknHomo (sub :< t) thin1 thin2 = cong2 (:<) (wknHomo sub thin1 thin2) (wknHomo t thin1 thin2)

-- Composition -----------------------------------------------------------------

export
(.) : Subst m n -> k `Thins` m -> Subst k n
comp : Subst m n -> Thinned n -> {0 thin : k `Thins` S m} -> View thin -> Subst k n

Wkn thin1 . thin = Wkn (thin1 . thin)
(sub :< t) . thin = comp sub t (view thin)

comp sub t (Id (S m)) = sub :< t
comp sub t (Drop thin) = sub . thin
comp sub t (Keep thin) = (sub . thin) :< t

export
compId : (sub : Subst m n) -> sub . id m = sub
compId (Wkn thin) = cong Wkn (compRightId thin)
compId (sub :< t) = viewId m (\v => comp sub t v = sub :< t) Refl

export
compKeep :
  (sub : Subst m n) ->
  (t : Thinned n) ->
  (thin : k `Thins` m) ->
  (sub :< t) . keep thin = (sub . thin) :< t
compKeep sub t thin =
  viewKeep thin (\v => comp sub t v = (sub . thin) :< t)
    (\Refl, Refl => sym $ cong (:< t) $ compId sub)
    Refl

export
wknComp :
  (sub : Subst k m) ->
  (thin1 : j `Thins` k) ->
  (thin2 : m `Thins` n) ->
  wkn (sub . thin1) thin2 = wkn sub thin2 . thin1
wknComp' :
  (sub : Subst k m) ->
  (t : Thinned m) ->
  {0 thin1 : j `Thins` S k} ->
  (v : View thin1) ->
  (thin2 : m `Thins` n) ->
  wkn (comp sub t v) thin2 = comp (wkn sub thin2) (wkn t thin2) v

wknComp (Wkn thin) thin1 thin2 = cong Wkn (compAssoc thin2 thin thin1)
wknComp (sub :< t) thin1 thin2 = wknComp' sub t (view thin1) thin2

wknComp' sub t (Id (S k)) thin2 = Refl
wknComp' sub t (Drop thin1) thin2 = wknComp sub thin1 thin2
wknComp' sub t (Keep thin1) thin2 = cong (:< wkn t thin2) (wknComp sub thin1 thin2)

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

export
indexHomo :
  (sub : Subst k m) ->
  (thin : m `Thins` n) ->
  (i : Fin k) ->
  index (wkn sub thin) i = wkn (index sub i) thin
indexHomo (Wkn thin') thin i = Refl
indexHomo (sub :< t) thin FZ = Refl
indexHomo (sub :< t) thin (FS i) = indexHomo sub thin i

export
indexComp :
  (sub : Subst m n) ->
  (thin : k `Thins` m) ->
  (i : Fin k) ->
  index (sub . thin) i =~= index sub (wkn i thin)
indexComp' :
  (sub : Subst m n) ->
  (t : Thinned n) ->
  {0 thin : k `Thins` S m} ->
  (v : View thin) ->
  (i : Fin k) ->
  index (comp sub t v) i =~= index (sub :< t) (wkn i thin)

indexComp (Wkn thin1) thin i = sym $ cong Var $ wknHomo i thin thin1
indexComp (sub :< t) thin i = indexComp' sub t (view thin) i

indexComp' sub t (Id (S m)) i = rewrite wknId i in Refl
indexComp' sub t (Drop thin) i = rewrite wknDrop i thin in indexComp sub thin i
indexComp' sub t (Keep thin) FZ = rewrite wknKeepFZ thin in Refl
indexComp' sub t (Keep thin) (FS i) = rewrite wknKeepFS i thin in indexComp sub thin i

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

public export
Refl' : {0 sub1, sub2 : Subst m n} -> sub1 = sub2 -> sub1 =~= sub2
Refl' Refl = Refl

public export
lift : Subst m n -> Subst (S m) (S n)
lift sub = wkn sub (wkn1 n) :< pure (Var 0)

export
wknLift :
  (sub : Subst k m) ->
  (thin : m `Thins` n) ->
  wkn (lift sub) (keep thin) =~= lift (wkn sub thin)
wknLift sub thin =
  let
    eq1 : wkn (wkn sub (wkn1 m)) (keep thin) =~= wkn (wkn sub thin) (wkn1 n)
    eq1 = Refl' $ Calc $
      |~ wkn (wkn sub (wkn1 m)) (keep thin)
      ~~ wkn sub (keep thin . wkn1 m)       ...(wknHomo sub (wkn1 m) (keep thin))
      ~~ wkn sub (wkn1 n . thin)            ...(cong (wkn sub) $ wkn1Comm thin)
      ~~ wkn (wkn sub thin) (wkn1 n)        ...(sym $ wknHomo sub thin (wkn1 n))
    eq2 : (Var (wkn 0 (keep thin . id (S m))) = Var (wkn FZ (id $ S n)))
    eq2 = cong Var $ Calc $
      |~ wkn 0 (keep thin . id (S m))
      ~~ wkn 0 (keep thin)            ...(cong (wkn 0) $ compRightId (keep thin))
      ~~ 0                            ...(wknKeepFZ thin)
      ~~ wkn 0 (id $ S n)             ...(sym $ wknId 0)
  in
  eq1 :< eq2

export
compLift :
  (sub : Subst m n) ->
  (thin : k `Thins` m) ->
  lift sub . keep thin = lift (sub . thin)
compLift sub thin = Calc $
  |~ lift sub . keep thin
  ~~ (wkn sub (wkn1 n) . thin) :< pure (Var 0) ...(compKeep (wkn sub (wkn1 n)) (pure $ Var 0) thin)
  ~~ lift (sub . thin)                         ...(sym $ cong (:< pure (Var 0)) $ wknComp sub thin (wkn1 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)

public export
subst1 : Term (S n) -> Term n -> Term n
subst1 t u = subst t (Wkn (id n) :< pure u)

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 (wkn1 n) :< Refl)
substCong (Abs t) prf = cong Abs (substCong t $ wknCong prf (wkn1 n) :< Refl)
substCong (App t u) prf = cong2 App (substCong t prf) (substCong u prf)

-- Interaction with Weakening

export
wknSubst :
  (t : Term k) ->
  (sub : Subst k m) ->
  (thin : m `Thins` n) ->
  wkn (subst t sub) thin = subst t (wkn sub thin)
wknSubst (Var i) sub thin = sym $ Calc $
  |~ expand (index (wkn sub thin) i)
  ~~ expand (wkn (index sub i) thin) ...(cong expand $ indexHomo sub thin i)
  ~~ wkn (expand (index sub i)) thin ...(expandHomo (index sub i) thin)
wknSubst Set sub thin = Refl
wknSubst (Pi t u) sub thin = cong2 Pi (wknSubst t sub thin) $ Calc $
  |~ wkn (subst u $ lift sub) (keep thin)
  ~~ subst u (wkn (lift sub) (keep thin)) ...(wknSubst u (lift sub) (keep thin))
  ~~ subst u (lift $ wkn sub thin)        ...(substCong u $ wknLift sub thin)
wknSubst (Abs t) sub thin = cong Abs $ Calc $
  |~ wkn (subst t $ lift sub) (keep thin)
  ~~ subst t (wkn (lift sub) (keep thin)) ...(wknSubst t (lift sub) (keep thin))
  ~~ subst t (lift $ wkn sub thin)        ...(substCong t $ wknLift sub thin)
wknSubst (App t u) sub thin = cong2 App (wknSubst t sub thin) (wknSubst u sub thin)

export
substWkn :
  (t : Term k) ->
  (thin : k `Thins` m) ->
  (sub : Subst m n) ->
  subst (wkn t thin) sub = subst t (sub . thin)

substWkn (Var i) thin sub = sym (indexComp sub thin i)
substWkn Set thin sub = Refl
substWkn (Pi t u) thin sub = cong2 Pi (substWkn t thin sub) $ Calc $
  |~ subst (wkn u (keep thin)) (lift sub)
  ~~ subst u (lift sub . keep thin)       ...(substWkn u (keep thin) (lift sub))
  ~~ subst u (lift (sub . thin))          ...(cong (subst u) $ compLift sub thin)
substWkn (Abs t) thin sub = cong Abs $ Calc $
  |~ subst (wkn t (keep thin)) (lift sub)
  ~~ subst t (lift sub . keep thin)       ...(substWkn t (keep thin) (lift sub))
  ~~ subst t (lift (sub . thin))          ...(cong (subst t) $ compLift sub thin)
substWkn (App t u) thin sub = cong2 App (substWkn t thin sub) (substWkn u thin sub)

export
wknSubst1 :
  (t : Term (S m)) ->
  (u : Term m) ->
  (thin : m `Thins` n) ->
  wkn (subst1 t u) thin = subst1 (wkn t $ keep thin) (wkn u thin)
wknSubst1 t u thin = Calc $
  |~ wkn (subst t (Wkn (id m) :< pure u)) thin
  ~~ subst t (Wkn (thin . id m) :< (u `Over` thin . id m))        ...(wknSubst t (Wkn (id m) :< pure u) thin)
  ~~ subst t (Wkn thin :< (u `Over` thin))                        ...(cong (\thin => subst t (Wkn thin :< (u `Over` thin))) $ compRightId thin)
  ~~ subst t (Wkn (id n . thin) :< pure (wkn u thin))             ...(sym $ substCong t (Refl' (cong Wkn $ compLeftId thin) :< wknId (wkn u thin)))
  ~~ subst t ((Wkn (id n) :< pure (wkn u thin)) . keep thin)      ...(sym $ cong (subst t) $ compKeep (Wkn $ id n) (pure $ wkn u thin) thin)
  ~~ subst (wkn t $ keep thin) (Wkn (id n) :< pure (wkn u thin))  ...(sym $ substWkn t (keep thin) (Wkn (id n) :< pure (wkn u thin)))