module Subst

import Control.Order
import Control.Relation
import Syntax.PreorderReasoning
import Syntax.PreorderReasoning.Generic
import Thinning

%prefix_record_projections off

-- Interfaces ------------------------------------------------------------------

public export
interface Point (0 a : Type) (0 t : a -> SnocList a -> Type) | t where
  point : {0 x : a} -> {0 sx : SnocList a} -> Elem x sx -> t x sx

public export
interface Weaken (0 a : Type) (0 t : a -> SnocList a -> Type) | t where
  wkn :
    {0 x : a} ->
    {0 sx, sy : SnocList a} ->
    t x sx ->
    sx `Thins` sy ->
    t x sy
  shift : {0 x, y : a} -> {0 sx : SnocList a} -> t y sx -> t y (sx :< x)
  shift v = wkn v (drop (id sx) x)

public export
interface Point a t => Weaken a t => PointWeaken a t | t where

public export
interface Weaken a t => FullWeaken a t | t where
  wknId : {0 x : a} -> {0 sx : SnocList a} -> (v : t x sx) -> wkn v (id sx) = v
  wknHomo :
    {0 x : a} ->
    {0 sx, sy, sz : SnocList a} ->
    (v : t x sx) ->
    (thin1 : sx `Thins` sy) ->
    (thin2 : sy `Thins` sz) ->
    wkn (wkn v thin1) thin2 = wkn v (thin2 . thin1)
  shiftIsWkn :
    {0 x, y : a} ->
    {0 sx : SnocList a} ->
    (v : t y sx) ->
    wkn v (drop (id sx) x) = shift v

public export
interface Point a t => FullWeaken a t => FullPointWeaken a t | t where
  wknPoint :
    {0 x : a} ->
    {0 sx, sy : SnocList a} ->
    (i : Elem x sx) ->
    (thin : sx `Thins` sy) ->
    wkn {t} (point i) thin = point (index thin i)

-- Substitutions ---------------------------------------------------------------

public export
data Subst : (t : a -> SnocList a -> Type) -> (sx, sy : SnocList a) -> Type where
  Base : sx `Thins` sy -> Subst t sx sy
  (:<) : Subst t sx sy -> t x sy -> Subst t (sx :< x) sy

%name Subst sub

export
index : Point a t => Subst t sx sy -> Elem x sx -> t x sy
index (Base thin) i = point (index thin i)
index (sub :< v) Here = v
index (sub :< v) (There i) = index sub i

namespace Subst
  export
  wkn : Weaken a t => Subst t sx sy -> sy `Thins` sz -> Subst t sx sz
  wkn (Base thin') thin = Base (thin . thin')
  wkn (sub :< v) thin = wkn sub thin :< wkn v thin

  export
  shift : Weaken a t => Subst t sx sy -> Subst t sx (sy :< y)
  shift (Base thin) = Base (drop thin y)
  shift (sub :< v) = shift sub :< shift v

  public export
  lift : Point a t => Weaken a t => Subst t sx sy -> Subst t (sx :< z) (sy :< z)
  lift sub = shift sub :< point Here

export
restrict : Subst t sy sz -> sx `Thins` sy -> Subst t sx sz
restrictView :
  Subst t sy sz ->
  t y sz ->
  {0 thin : sx `Thins` sy :< y} ->
  View thin ->
  Subst t sx sz

restrict (Base thin') thin = Base (thin' . thin)
restrict (sub :< v) thin = restrictView sub v (view thin)

restrictView sub v Id = sub :< v
restrictView sub v (Drop thin z) = restrict sub thin
restrictView sub v (Keep thin z) = restrict sub thin :< v

-- Equivalence -----------------------------------------------------------------

public export
record Equiv {auto 0 p : Point a t} (sub1, sub2 : Subst t sx sy) where
  constructor IsEquiv
  equiv : (forall x. (i : Elem x sx) -> index sub1 i = index sub2 i)

%name Subst.Equiv e

export
Point a t => Reflexive (Subst t sx sy) Equiv where
  reflexive = IsEquiv (\i => Refl)

export
Point a t => Symmetric (Subst t sx sy) Equiv where
  symmetric e = IsEquiv (\i => sym (e.equiv i))

export
Point a t => Transitive (Subst t sx sy) Equiv where
  transitive e e' = IsEquiv (\i => trans (e.equiv i) (e'.equiv i))

export
Point a t => Preorder (Subst t sx sy) Equiv where

namespace Equiv
  export
  Base :
    (p : Point a t) =>
    {0 thin1, thin2 : sx `Thins` sy} ->
    thin1 = thin2 ->
    Equiv @{p} (Base thin1) (Base thin2)
  Base prf = IsEquiv (\i => cong (\thin => index (Base thin) i) prf)

  export
  (:<) :
    Point a t =>
    {0 sub1, sub2 : Subst t sx sy} ->
    Equiv sub1 sub2 ->
    u = v ->
    Equiv (sub1 :< u) (sub2 :< v)
  e :< prf = IsEquiv (\i => case i of Here => prf; There i => e.equiv i)

-- Properties ------------------------------------------------------------------

export
irrelevantEquiv :
  Point a t =>
  {0 sub1, sub2 : Subst t sx sy} ->
  (0 _ : Equiv sub1 sub2) ->
  Equiv sub1 sub2
irrelevantEquiv e = IsEquiv (\i => irrelevantEq $ e.equiv i)

-- Weakening

namespace Subst
  export
  indexWkn :
    FullPointWeaken a t =>
    (sub : Subst t sx sy) ->
    (thin : sy `Thins` sz) ->
    (i : Elem x sx) ->
    wkn (index sub i) thin = index (wkn sub thin) i
  indexWkn (Base thin') thin i = Calc $
    |~ wkn (point (index thin' i)) thin
    ~~ point {t} (index thin (index thin' i)) ...(wknPoint (index thin' i) thin)
    ~~ point (index (thin . thin') i)         ...(cong point $ indexHomo thin' thin i)
  indexWkn (sub :< v) thin Here = Refl
  indexWkn (sub :< v) thin (There i) = indexWkn sub thin i

  export
  wknCong :
    FullPointWeaken a t =>
    {0 sub1, sub2 : Subst t sx sy} ->
    {0 thin1, thin2 : sy `Thins` sz} ->
    Equiv sub1 sub2 ->
    thin1 = thin2 ->
    Equiv (wkn sub1 thin1) (wkn sub2 thin2)
  wknCong e prf = irrelevantEquiv $ IsEquiv (\i => Calc $
    |~ index (wkn sub1 thin1) i
    ~~ wkn (index sub1 i) thin1 ..<(indexWkn sub1 thin1 i)
    ~~ wkn (index sub2 i) thin2 ...(cong2 wkn (e.equiv i) prf)
    ~~ index (wkn sub2 thin2) i ...(indexWkn sub2 thin2 i))

  export
  wknId :
    FullPointWeaken a t =>
    (sub : Subst t sx sy) ->
    Equiv (wkn sub (id sy)) sub
  wknId (Base thin) = Base (identityLeft thin)
  wknId (sub :< v) = wknId sub :< wknId v

  export
  wknHomo :
    FullPointWeaken a t =>
    (sub : Subst t sx sy) ->
    (thin1 : sy `Thins` sz) ->
    (thin2 : sz `Thins` sw) ->
    Equiv (wkn (wkn sub thin1) thin2) (wkn sub (thin2 . thin1))
  wknHomo (Base thin) thin1 thin2 = Base (assoc thin2 thin1 thin)
  wknHomo (sub :< v) thin1 thin2 = wknHomo sub thin1 thin2 :< wknHomo v thin1 thin2

-- Restrict

export
restrictBase :
  (p : Point a t) =>
  (thin' : sy `Thins` sz) ->
  (thin : sx `Thins` sy) ->
  Equiv @{p} (Base (thin' . thin)) (restrict (Base thin') thin)
restrictBase thin' thin = reflexive

export
indexRestrict :
  Point a t =>
  (sub : Subst t sy sz) ->
  (thin : sx `Thins` sy) ->
  (i : Elem x sx) ->
  index sub (index thin i) = index (restrict sub thin) i
indexRestrictView :
  Point a t =>
  (sub : Subst t sy sz) ->
  (v : t y sz) ->
  {0 thin : sx `Thins` sy :< y} ->
  (view : View thin) ->
  (i : Elem x sx) ->
  index (sub :< v) (index thin i) = index (restrictView sub v view) i

indexRestrict (Base thin') thin i = cong point (indexHomo thin thin' i)
indexRestrict (sub :< v) thin i = indexRestrictView sub v (view thin) i

indexRestrictView sub v Id i = cong (index (sub :< v)) (indexId i)
indexRestrictView sub v (Drop thin z) i = Calc $
  |~ index (sub :< v) (index (drop thin z) i)
  ~~ index (sub :< v) (There (index thin i))  ...(cong (index (sub :< v)) $ indexDrop thin i)
  ~~ index sub (index thin i)                 ...(Refl)
  ~~ index (restrict sub thin) i              ...(indexRestrict sub thin i)
indexRestrictView sub v (Keep thin x) Here = cong (index (sub :< v)) (indexKeepHere thin)
indexRestrictView sub v (Keep thin z) (There i) = Calc $
  |~ index (sub :< v) (index (keep thin z) (There i))
  ~~ index (sub :< v) (There (index thin i))          ...(cong (index (sub :< v)) $ indexKeepThere thin i)
  ~~ index sub (index thin i)                         ...(Refl)
  ~~ index (restrict sub thin) i                      ...(indexRestrict sub thin i)

export
restrictCong :
  Point a t =>
  {0 sub1, sub2 : Subst t sy sz} ->
  {0 thin1, thin2 : sx `Thins` sy} ->
  Equiv sub1 sub2 ->
  thin1 = thin2 ->
  Equiv (restrict sub1 thin1) (restrict sub2 thin2)
restrictCong e prf = irrelevantEquiv $ IsEquiv (\i => Calc $
  |~ index (restrict sub1 thin1) i
  ~~ index sub1 (index thin1 i)    ..<(indexRestrict sub1 thin1 i)
  ~~ index sub2 (index thin1 i)    ...(e.equiv (index thin1 i))
  ~~ index sub2 (index thin2 i)    ...(cong (\thin => index sub2 (index thin i)) prf)
  ~~ index (restrict sub2 thin2) i ...(indexRestrict sub2 thin2 i))

export
restrictId :
  Point a t =>
  (sub : Subst t sx sy) ->
  Equiv (restrict sub (id sx)) sub
restrictId (Base thin) = Base (identityRight thin)
restrictId {sx = .(sx :< x)} (sub :< v) =
  rewrite viewUnique (view (id (sx :< x))) Id in
  reflexive

export
restrictHomo :
  Point a t =>
  (sub : Subst t sz sw) ->
  (thin1 : sy `Thins` sz) ->
  (thin2 : sx `Thins` sy) ->
  Equiv (restrict (restrict sub thin1) thin2) (restrict sub (thin1 . thin2))
restrictViewHomo1 :
  Point a t =>
  (sub : Subst t sz sw) ->
  (v : t z sw) ->
  {0 thin1 : sy `Thins` sz :< z} ->
  (view1 : View thin1) ->
  (thin2 : sx `Thins` sy) ->
  Equiv (restrict (restrictView sub v view1) thin2) (restrictView sub v (view (thin1 . thin2)))
restrictViewHomo2 :
  Point a t =>
  (sub : Subst t sz sw) ->
  (v : t y sw) ->
  (thin1 : sy `Thins` sz) ->
  {auto 0 ok : NotId thin1} ->
  {0 thin2 : sx `Thins` sy :< y} ->
  (view2 : View thin2) ->
  Equiv
    (restrictView (restrict sub thin1) v view2)
    (restrictView sub v (view (keep thin1 y . thin2)))

restrictHomo (Base thin) thin1 thin2 = symmetric $ Base (assoc thin thin1 thin2)
restrictHomo (sub :< v) thin1 thin2 = restrictViewHomo1 sub v (view thin1) thin2

restrictViewHomo1 sub v Id thin2 =
  rewrite identityLeft thin2 in
  reflexive
restrictViewHomo1 sub v (Drop thin1 z) thin2 =
  rewrite dropLeft thin1 thin2 z in
  rewrite viewUnique (view (drop (thin1 . thin2) z)) (Drop (thin1 . thin2) z) in
  restrictHomo sub thin1 thin2
restrictViewHomo1 sub v (Keep thin1 z) thin2 = restrictViewHomo2 sub v thin1 (view thin2)

restrictViewHomo2 sub v thin1 Id =
  rewrite identityRight (keep thin1 y) in
  rewrite viewUnique (view (keep thin1 y)) (Keep thin1 y) in
  reflexive
restrictViewHomo2 sub v thin1 (Drop thin2 y) =
  rewrite keepDrop thin1 thin2 y in
  rewrite viewUnique (view (drop (thin1 . thin2) y)) (Drop (thin1 . thin2) y) in
  restrictHomo sub thin1 thin2
restrictViewHomo2 sub v thin1 (Keep thin2 {ok = ok'} y) =
  rewrite keepHomo thin1 thin2 y in
  rewrite
    viewUnique
      (view (keep (thin1 . thin2) y))
      (Keep (thin1 . thin2) {ok = compNotId ok ok'} y)
  in
  restrictHomo sub thin1 thin2 :< Refl

namespace Subst
  export
  wknRestrict :
    FullPointWeaken a t =>
    (sub : Subst t sy sz) ->
    (thin1 : sx `Thins` sy) ->
    (thin2 : sz `Thins` sw) ->
    Equiv (wkn (restrict sub thin1) thin2) (restrict (wkn sub thin2) thin1)
  wknRestrictView :
    FullPointWeaken a t =>
    (sub : Subst t sy sz) ->
    (v : t y sz) ->
    {0 thin1 : sx `Thins` sy :< y} ->
    (view1 : View thin1) ->
    (thin2 : sz `Thins` sw) ->
    Equiv (wkn (restrictView sub v view1) thin2) (restrictView (wkn sub thin2) (wkn v thin2) view1)

  wknRestrict (Base thin) thin1 thin2 = Base (assoc thin2 thin thin1)
  wknRestrict (sub :< v) thin1 thin2 = wknRestrictView sub v (view thin1) thin2

  wknRestrictView sub v Id thin2 = reflexive
  wknRestrictView sub v (Drop thin1 y) thin2 = wknRestrict sub thin1 thin2
  wknRestrictView sub v (Keep thin1 y) thin2 = wknRestrict sub thin1 thin2 :< Refl

-- Shift

  export
  shiftIsWkn :
    Point a t =>
    FullWeaken a t =>
    (sub : Subst t sx sy) ->
    Equiv (wkn sub (drop (id sy) y)) (shift sub)
  shiftIsWkn (Base thin) =
    Base $ trans
      (dropLeft (id sy) thin y)
      (cong (\thin => drop thin y) $ identityLeft thin)
  shiftIsWkn (sub :< v) = shiftIsWkn sub :< shiftIsWkn v

export
shiftCong :
  FullPointWeaken a t =>
  {0 sub1, sub2 : Subst t sx sy} ->
  Equiv sub1 sub2 ->
  Equiv (shift {y} sub1) (shift sub2)
shiftCong e = irrelevantEquiv $ CalcWith $
  |~ shift sub1
  <~ wkn sub1 (drop (id sy) y) ..<(shiftIsWkn sub1)
  <~ wkn sub2 (drop (id sy) y) ...(wknCong e reflexive)
  <~ shift sub2                ...(shiftIsWkn sub2)

export
wknShiftKeep :
  FullPointWeaken a t =>
  (sub : Subst t sx sy) ->
  (thin : sy `Thins` sz) ->
  Equiv (wkn (shift sub) (keep thin y)) (shift (wkn sub thin))
wknShiftKeep sub thin = CalcWith $
  |~ wkn (shift sub) (keep thin y)
  <~ wkn (wkn sub (drop (id sy) y)) (keep thin y) ..<(wknCong (shiftIsWkn sub) reflexive)
  <~ wkn sub (keep thin y . drop (id sy) y)       ...(wknHomo sub (drop (id sy) y) (keep thin y))
  ~~ wkn sub (drop (thin . (id sy)) y)            ...(cong (wkn sub) (keepDrop thin (id sy) y))
  ~~ wkn sub (drop thin y)                        ...(cong (\thin => wkn sub (drop thin y)) (identityRight thin))
  ~~ wkn sub (drop (id sz . thin) y)              ..<(cong (\thin => wkn sub (drop thin y)) (identityLeft thin))
  <~ wkn sub (drop (id sz) y . thin)              ..<(wknCong reflexive (dropLeft (id sz) thin y))
  <~ wkn (wkn sub thin) (drop (id sz) y)          ..<(wknHomo sub thin (drop (id sz) y))
  <~ shift (wkn sub thin)                         ...(shiftIsWkn (wkn sub thin))


export
restrictShift :
  FullPointWeaken a t =>
  (sub : Subst t sy sz) ->
  (thin : sx `Thins` sy) ->
  Equiv (restrict (shift {y = z} sub) thin) (shift (restrict sub thin))
restrictShift sub thin = CalcWith $
  |~ restrict (shift sub) thin
  <~ restrict (wkn sub (drop (id sz) z)) thin ..<(restrictCong (shiftIsWkn sub) reflexive)
  <~ wkn (restrict sub thin) (drop (id sz) z) ..<(wknRestrict sub thin (drop (id sz) z))
  <~ shift (restrict sub thin)                ...(shiftIsWkn (restrict sub thin))

-- Lift

namespace Subst
  export
  liftCong :
    FullPointWeaken a t =>
    {0 sub1, sub2 : Subst t sx sy} ->
    Equiv sub1 sub2 ->
    Equiv (lift sub1) (lift sub2)
  liftCong e = shiftCong e :< Refl

  export
  wknLift :
    FullPointWeaken a t =>
    (sub : Subst t sx sy) ->
    (thin : sy `Thins` sz) ->
    Equiv (wkn (lift sub) (keep thin z)) (lift (wkn sub thin))
  wknLift sub thin =
    wknShiftKeep sub thin :<
    trans
      (wknPoint Here (keep thin z))
      (cong point $ indexKeepHere thin)

  restrictLiftView :
    FullPointWeaken a t =>
    (sub : Subst t sy sz) ->
    {0 thin : sx `Thins` sy} ->
    (view : View thin) ->
    Equiv (restrict (lift sub) (keep thin y)) (lift (restrict sub thin))
  restrictLiftView sub Id =
    rewrite keepId sy y in
    rewrite viewUnique (view (id (sy :< y))) Id in
    symmetric (shiftCong (restrictId sub)) :< Refl
  restrictLiftView sub (Drop thin z) =
    rewrite
      viewUnique
        (view (keep (drop thin z) y))
        (Keep (drop thin z) {ok = dropNotId thin z} y)
    in
    restrictShift sub (drop thin z) :< Refl
  restrictLiftView sub (Keep thin z) =
    rewrite
      viewUnique
        (view (keep (keep thin z) y))
        (Keep (keep thin z) {ok = keepNotId thin z} y)
    in
    restrictShift sub (keep thin z) :< Refl

  export
  restrictLift :
    FullPointWeaken a t =>
    (sub : Subst t sy sz) ->
    (thin : sx `Thins` sy) ->
    Equiv (restrict (lift sub) (keep thin y)) (lift (restrict sub thin))
  restrictLift sub thin = restrictLiftView sub (view thin)