module Thinning

import public Data.SnocList.Elem

import Syntax.PreorderReasoning

%prefix_record_projections off

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

export
data Thins : SnocList a -> SnocList a -> Type
export
data NotId : sx `Thins` sy -> Type

data Thins where
  Id : sx `Thins` sx
  Drop : sx `Thins` sy -> sx `Thins` sy :< y
  Keep :
    (thin : sx `Thins` sy) ->
    {auto 0 ok : NotId thin} ->
    sx :< z `Thins` sy :< z

data NotId where
  DropNotId : NotId (Drop thin)
  KeepNotId : NotId thin -> NotId (Keep thin)

%name Thins thin

-- Smart Constructors ----------------------------------------------------------

export
id : (0 sx : SnocList a) -> sx `Thins` sx
id sx = Id

export
drop : sx `Thins` sy -> (0 z : a) -> sx `Thins` sy :< z
drop thin z = Drop thin

export
keep : sx `Thins` sy -> (0 z : a) -> sx :< z `Thins` sy :< z
keep Id z = Id
keep (Drop thin) z = Keep (Drop thin)
keep (Keep thin) z = Keep (Keep thin)

-- Views -----------------------------------------------------------------------

notIdUnique : (p, q : NotId thin) -> p = q
notIdUnique DropNotId DropNotId = Refl
notIdUnique (KeepNotId p) (KeepNotId p) = Refl

keepIsKeep :
  (thin : sx `Thins` sy) ->
  {auto 0 ok : NotId thin} ->
  (0 z : a) ->
  keep thin z = Keep {ok} thin
keepIsKeep (Drop thin) z =
  rewrite notIdUnique ok DropNotId in Refl
keepIsKeep (Keep {ok = ok'} thin) z =
  rewrite notIdUnique ok (KeepNotId ok') in Refl

namespace View
  public export
  data View : sx `Thins` sy -> Type where
    Id : View (id sx)
    Drop : (thin : sx `Thins` sy) -> (0 z : a) -> View (drop thin z)
    Keep :
      (thin : sx `Thins` sy) ->
      {auto 0 ok : NotId thin} ->
      (0 z : a) ->
      View (keep thin z)

  %name View view

export
view : (thin : sx `Thins` sy) -> View thin
view Id = Id
view (Drop thin) = Drop thin _
view (Keep {ok} (Drop thin)) =
  rewrite notIdUnique ok DropNotId in
  Keep (Drop thin) _
view (Keep {ok} (Keep {ok = ok'} thin)) =
  rewrite notIdUnique ok (KeepNotId ok') in
  Keep (Keep thin) _

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

export
index : sx `Thins` sy -> Elem x sx -> Elem x sy
index Id i = i
index (Drop thin) i = There (index thin i)
index (Keep thin) Here = Here
index (Keep thin) (There i) = There (index thin i)

-- Properties (more at the end)

export
indexKeepHere :
  (thin : sx `Thins` sy) ->
  index (keep thin z) Here = Here
indexKeepHere Id = Refl
indexKeepHere (Drop thin) = Refl
indexKeepHere (Keep thin) = Refl

export
indexKeepThere :
  (thin : sx `Thins` sy) ->
  (i : Elem x sx) ->
  index (keep thin z) (There i) = There (index thin i)
indexKeepThere Id i = Refl
indexKeepThere (Drop thin) i = Refl
indexKeepThere (Keep thin) i = Refl

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

export
(.) : sy `Thins` sz -> sx `Thins` sy -> sx `Thins` sz
export
compNotId :
  {thin2 : sy `Thins` sz} ->
  {thin1 : sx `Thins` sy} ->
  NotId thin2 ->
  NotId thin1 ->
  NotId (thin2 . thin1)

Id . thin1 = thin1
Drop thin2 . Id = Drop thin2
Drop thin2 . thin1@(Drop _) = Drop (thin2 . thin1)
Drop thin2 . thin1@(Keep _) = Drop (thin2 . thin1)
Keep thin2 . Id = Keep thin2
Keep thin2 . Drop thin1 = Drop (thin2 . thin1)
Keep {ok} thin2 . Keep {ok = ok'} thin1 = Keep {ok = compNotId ok ok'} (thin2 . thin1)

compNotId DropNotId DropNotId = DropNotId
compNotId DropNotId (KeepNotId q) = DropNotId
compNotId (KeepNotId p) DropNotId = DropNotId
compNotId (KeepNotId p) (KeepNotId q) = KeepNotId (compNotId p q)

-- Properties (more at the end)

export
identityRight :
  (thin : sx `Thins` sy) ->
  thin . id sx = thin
identityRight Id = Refl
identityRight (Drop thin) = Refl
identityRight (Keep thin) = Refl

export
dropLeft :
  (thin2 : sy `Thins` sz) ->
  (thin1 : sx `Thins` sy) ->
  (0 z : a) ->
  drop thin2 z . thin1 = drop (thin2 . thin1) z
dropLeft thin2 Id z = sym $ cong Drop $ identityRight thin2
dropLeft thin2 thin1@(Drop _) z = Refl
dropLeft thin2 thin1@(Keep _) z = Refl

export
keepDrop :
  (thin2 : sy `Thins` sz) ->
  (thin1 : sx `Thins` sy) ->
  (0 z : a) ->
  keep thin2 z . drop thin1 z = drop (thin2 . thin1) z
keepDrop Id thin1 z = Refl
keepDrop (Drop thin2) thin1 z = Refl
keepDrop (Keep thin2) thin1 z = Refl

export
keepHomo :
  (thin2 : sy `Thins` sz) ->
  (thin1 : sx `Thins` sy) ->
  (0 z : a) ->
  keep thin2 z . keep thin1 z = keep (thin2 . thin1) z
keepHomo Id thin1 z = Refl
keepHomo (Drop thin2) Id z = Refl
keepHomo (Drop thin2) (Drop thin1) z = Refl
keepHomo (Drop thin2) (Keep thin1) z = Refl
keepHomo (Keep thin2) Id z = Refl
keepHomo (Keep thin2) (Drop thin1) z = Refl
keepHomo (Keep thin2) (Keep thin1) z = Refl

-- Coverings and Coproducts ----------------------------------------------------

export
Covering : {a : Type} -> {sx, sy, sz : SnocList a} -> sx `Thins` sz -> sy `Thins` sz -> Type
Covering thin1 thin2 =
  {x : a} ->
  (i : Elem x sz) ->
  Either (j ** index thin1 j = i) (k ** index thin2 k = i)

public export
record Coproduct {0 sx, sy : SnocList a} (thin1 : sx `Thins` sz) (thin2 : sy `Thins` sz) where
  constructor MkCoprod
  {0 sw : SnocList a}
  {thin1' : sx `Thins` sw}
  {thin2' : sy `Thins` sw}
  {thin : sw `Thins` sz}
  0 left : thin . thin1' = thin1
  0 right : thin . thin2' = thin2
  0 cover : Covering thin1' thin2'

%name Coproduct cp

coprodSym : Coproduct thin1 thin2 -> Coproduct thin2 thin1
coprodSym cp = MkCoprod
  cp.right
  cp.left
  (\i => case cp.cover i of Left x => Right x; Right x => Left x)

coprodId : (thin : sx `Thins` sy) -> Coproduct Id thin
coprodId thin = MkCoprod
  { thin1' = Id
  , thin2' = thin
  , thin = Id
  , left = Refl
  , right = Refl
  , cover = \i => Left (i ** Refl)
  }

coprod : (thin1 : sx `Thins` sz) -> (thin2 : sy `Thins` sz) -> Coproduct thin1 thin2
coprod Id thin2 = coprodId thin2
coprod thin1@(Drop _) Id = coprodSym $ coprodId thin1
coprod (Drop thin1) (Drop thin2) =
  let cp = coprod thin1 thin2 in
  MkCoprod
    { thin1' = cp.thin1'
    , thin2' = cp.thin2'
    , thin = Drop cp.thin
    , left = trans (dropLeft cp.thin cp.thin1' _) (cong Drop cp.left)
    , right = trans (dropLeft cp.thin cp.thin2' _) (cong Drop cp.right)
    , cover = cp.cover
    }
coprod (Drop thin1) (Keep {z} thin2) =
  let cp = coprod thin1 thin2 in
  MkCoprod
    { thin1' = Drop cp.thin1'
    , thin2' = keep cp.thin2' z
    , thin = keep cp.thin z
    , left = trans (keepDrop cp.thin cp.thin1' z) (cong Drop cp.left)
    , right = Calc $
      |~ keep cp.thin z . keep cp.thin2' z
      ~~ keep (cp.thin . cp.thin2') z      ...(keepHomo cp.thin cp.thin2' z)
      ~~ keep thin2 z                      ...(cong (\thin => keep thin z) cp.right)
      ~~ Keep thin2                        ...(keepIsKeep thin2 z)
    , cover = \i =>
      case i of
        Here => Right (Here ** indexKeepHere cp.thin2')
        There i => case cp.cover i of
          Left (j ** prf) => Left (j ** cong There prf)
          Right (k ** prf) => Right (There k ** trans (indexKeepThere cp.thin2' k) (cong There prf))
    }
coprod thin1@(Keep _) Id = coprodSym $ coprodId thin1
coprod (Keep {z} thin1) (Drop thin2) =
  let cp = coprod thin1 thin2 in
  MkCoprod
    { thin1' = keep cp.thin1' z
    , thin2' = Drop cp.thin2'
    , thin = keep cp.thin z
    , left = Calc $
      |~ keep cp.thin z . keep cp.thin1' z
      ~~ keep (cp.thin . cp.thin1') z      ...(keepHomo cp.thin cp.thin1' z)
      ~~ keep thin1 z                      ...(cong (\thin => keep thin z) cp.left)
      ~~ Keep thin1                        ...(keepIsKeep thin1 z)
    , right = trans (keepDrop cp.thin cp.thin2' z) (cong Drop cp.right)
    , cover = \i =>
      case i of
        Here => Left (Here ** indexKeepHere cp.thin1')
        There i => case cp.cover i of
          Left (j ** prf) => Left (There j ** trans (indexKeepThere cp.thin1' j) (cong There prf))
          Right (k ** prf) => Right (k ** cong There prf)
    }
coprod (Keep {z} thin1) (Keep thin2) =
  let cp = coprod thin1 thin2 in
  MkCoprod
    { thin1' = keep cp.thin1' z
    , thin2' = keep cp.thin2' z
    , thin = keep cp.thin z
    , left = Calc $
      |~ keep cp.thin z . keep cp.thin1' z
      ~~ keep (cp.thin . cp.thin1') z      ...(keepHomo cp.thin cp.thin1' z)
      ~~ keep thin1 z                      ...(cong (\thin => keep thin z) cp.left)
      ~~ Keep thin1                        ...(keepIsKeep thin1 z)
    , right = Calc $
      |~ keep cp.thin z . keep cp.thin2' z
      ~~ keep (cp.thin . cp.thin2') z      ...(keepHomo cp.thin cp.thin2' z)
      ~~ keep thin2 z                      ...(cong (\thin => keep thin z) cp.right)
      ~~ Keep thin2                        ...(keepIsKeep thin2 z)
    , cover = \i =>
      case i of
        Here => Left (Here ** indexKeepHere cp.thin1')
        There i => case cp.cover i of
          Left (j ** prf) => Left (There j ** trans (indexKeepThere cp.thin1' j) (cong There prf))
          Right (k ** prf) => Right (There k ** trans (indexKeepThere cp.thin2' k) (cong There prf))
    }

-- Thinned Things --------------------------------------------------------------

public export
record Thinned (T : SnocList a -> Type) (sx : SnocList a) where
  constructor Over
  {0 support : SnocList a}
  value : T support
  thin : support `Thins` sx

public export
record Pair (T, U : SnocList a -> Type) (sx : SnocList a) where
  constructor MakePair
  left : Thinned T sx
  right : Thinned U sx
  0 cover : Covering left.thin right.thin

public export
MkPair : Thinned t sx -> Thinned u sx -> Thinned (Pair t u) sx
MkPair (t `Over` thin1) (u `Over` thin2) =
  let cp = coprod thin1 thin2 in
  MakePair (t `Over` cp.thin1') (u `Over` cp.thin2') cp.cover `Over` cp.thin

public export
map : (forall ctx. t ctx -> u ctx) -> Thinned t ctx -> Thinned u ctx
map f (value `Over` thin) = f value `Over` thin

public export
shift : Thinned t sx -> Thinned t (sx :< x)
shift (value `Over` thin) = value `Over` drop thin x

public export
wkn : Thinned t sx -> sx `Thins` sy -> Thinned t sy
wkn (value `Over` thin) thin' = value `Over` thin' . thin

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

-- Construction

export
keepId : (0 sx : SnocList a) -> (0 x : a) -> keep (id sx) x = id (sx :< x)
keepId sx x = Refl

-- Views

export
dropNotId : (thin : sx `Thins` sy) -> (0 y : a) -> NotId (drop thin y)
dropNotId thin y = DropNotId

export
keepNotId : (thin : sx `Thins` sy) -> {auto ok : NotId thin} -> (0 y : a) -> NotId (keep thin y)
keepNotId (Drop thin) y = KeepNotId DropNotId
keepNotId (Keep thin) {ok = KeepNotId ok} y = KeepNotId (KeepNotId ok)

invertView : {0 thin : sx `Thins` sy} -> View thin -> sx `Thins` sy
invertView Id = id sx
invertView (Drop thin1 z) = drop thin1 z
invertView (Keep thin1 z) = keep thin1 z

invertViewCorrect : {0 thin : sx `Thins` sy} -> (view : View thin) -> invertView view = thin
invertViewCorrect Id = Refl
invertViewCorrect (Drop thin1 z) = Refl
invertViewCorrect (Keep thin1 z) = Refl

invertViewCorrect' :
  {0 thin : sx `Thins` sy} ->
  (v : View thin) ->
  view (invertView v) = (rewrite invertViewCorrect v in v)
invertViewCorrect' Id = Refl
invertViewCorrect' (Drop thin z) = Refl
invertViewCorrect' (Keep (Drop thin) {ok} z) = rewrite notIdUnique ok DropNotId in Refl
invertViewCorrect' (Keep (Keep {ok = ok'} thin) {ok} z) =
  rewrite notIdUnique ok (KeepNotId ok') in Refl

export
viewUnique : {0 thin : sx `Thins` sy} -> (v1, v2 : View thin) -> v1 = v2
viewUnique v1 v2 =
  rewrite sym $ invertViewCorrect' v1 in
  rewrite sym $ invertViewCorrect' v2 in
  rewrite invertViewCorrect v1 in
  rewrite invertViewCorrect v2 in
  Refl

-- Index

export
indexId : (i : Elem x sx) -> index (id sx) i = i
indexId i = Refl

export
indexDrop :
  (thin : sx `Thins` sy) ->
  (i : Elem x sx) ->
  index (drop thin z) i = There (index thin i)
indexDrop thin i = Refl

export
indexHomo :
  (thin1 : sx `Thins` sy) ->
  (thin2 : sy `Thins` sz) ->
  (i : Elem x sx) ->
  index thin2 (index thin1 i) = index (thin2 . thin1) i
indexHomo thin1 Id i = Refl
indexHomo Id (Drop thin2) i = Refl
indexHomo thin1@(Drop _) (Drop thin2) i = cong There $ indexHomo thin1 thin2 i
indexHomo thin1@(Keep _) (Drop thin2) i = cong There $ indexHomo thin1 thin2 i
indexHomo Id (Keep thin2) i = Refl
indexHomo (Drop thin1) (Keep thin2) i = cong There $ indexHomo thin1 thin2 i
indexHomo (Keep thin1) (Keep thin2) Here = Refl
indexHomo (Keep thin1) (Keep thin2) (There i) = cong There $ indexHomo thin1 thin2 i

-- Composition

export
identityLeft :
  (thin : sx `Thins` sy) ->
  id sy . thin = thin
identityLeft thin = Refl

export
assoc :
  (thin3 : sz `Thins` sw) ->
  (thin2 : sy `Thins` sz) ->
  (thin1 : sx `Thins` sy) ->
  thin3 . (thin2 . thin1) = (thin3 . thin2) . thin1
assoc Id thin2 thin1 = Refl
assoc (Drop thin3) Id thin1 = Refl
assoc (Drop thin3) thin2@(Drop _) Id = Refl
assoc (Drop thin3) thin2@(Drop _) thin1@(Drop _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Drop thin3) thin2@(Drop _) thin1@(Keep _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Drop thin3) thin2@(Keep _) Id = Refl
assoc (Drop thin3) thin2@(Keep _) thin1@(Drop _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Drop thin3) thin2@(Keep _) thin1@(Keep _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Keep thin3) Id thin1 = Refl
assoc (Keep thin3) (Drop thin2) Id = Refl
assoc (Keep thin3) (Drop thin2) thin1@(Drop _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Keep thin3) (Drop thin2) thin1@(Keep _) = cong Drop $ assoc thin3 thin2 thin1
assoc (Keep thin3) (Keep thin2) Id = Refl
assoc (Keep thin3) (Keep thin2) (Drop thin1) = cong Drop $ assoc thin3 thin2 thin1
assoc (Keep {ok = ok3} thin3) (Keep {ok = ok2} thin2) (Keep {ok = ok1} thin1) =
  rewrite assoc thin3 thin2 thin1 in
  rewrite
    notIdUnique
      (compNotId ok3 (compNotId ok2 ok1))
      (rewrite assoc thin3 thin2 thin1 in compNotId (compNotId ok3 ok2) ok1) in
  Refl