||| thinnings, covers, partitions, etc,
||| for co-de Bruijn representation [@egtbs]
module Quox.Thin

import Quox.NatExtra
import Data.Nat.Views
import Data.Singleton
import Data.DPair
import Syntax.PreorderReasoning

%default total


||| "order preserving embeddings", for recording a correspondence between a
||| smaller scope and part of a larger one. the third argument is a bitmask
||| representing this OPE, unique for a given `n`.
public export
data OPE : (m, n, mask : Nat) -> Type where
  [search m n]
  Stop : OPE 0 0 0
  Drop : OPE m n mask -> mask' = mask + mask       -> OPE m     (S n) mask'
  Keep : OPE m n mask -> mask' = (S (mask + mask)) -> OPE (S m) (S n) mask'
%name OPE ope


||| everything selected
public export
id : {m : Nat} -> Subset Nat (OPE m m)
id {m = 0}   = Element _ Stop
id {m = S m} = Element _ $ Keep id.snd Refl

||| nothing selected
public export
zero : {m : Nat} -> OPE 0 m 0
zero {m = 0}   = Stop
zero {m = S m} = Drop zero Refl


infix 6 `Eqv`

public export
data Eqv : OPE m1 n1 mask1 -> OPE m2 n2 mask2 -> Type where
  EqvStop : Eqv Stop Stop
  EqvDrop : {0 p : OPE m1 n1 mask1} ->
            {0 q : OPE m2 n2 mask2} ->
            Eqv p q -> Eqv (Drop p eq1) (Drop q eq2)
  EqvKeep : {0 p : OPE m1 n1 mask1} ->
            {0 q : OPE m2 n2 mask2} ->
            Eqv p q -> Eqv (Keep p eq1) (Keep q eq2)
%name Eqv eqv

export
Reflexive (OPE m n mask) Eqv where
  reflexive {x = Stop}    = EqvStop
  reflexive {x = Drop {}} = EqvDrop reflexive
  reflexive {x = Keep {}} = EqvKeep reflexive

export
symmetric : p `Eqv` q -> q `Eqv` p
symmetric EqvStop       = EqvStop
symmetric (EqvDrop eqv) = EqvDrop $ symmetric eqv
symmetric (EqvKeep eqv) = EqvKeep $ symmetric eqv

export
transitive : p `Eqv` q -> q `Eqv` r -> p `Eqv` r
transitive EqvStop     EqvStop     = EqvStop
transitive (EqvDrop x) (EqvDrop y) = EqvDrop (transitive x y)
transitive (EqvKeep x) (EqvKeep y) = EqvKeep (transitive x y)

export
eqvIndices : {0 p : OPE m1 n1 mask1} -> {0 q : OPE m2 n2 mask2} ->
             p `Eqv` q -> (m1 = m2, n1 = n2, mask1 = mask2)
eqvIndices EqvStop = (Refl, Refl, Refl)
eqvIndices (EqvDrop eqv {eq1 = Refl, eq2 = Refl}) =
  let (Refl, Refl, Refl) = eqvIndices eqv in (Refl, Refl, Refl)
eqvIndices (EqvKeep eqv {eq1 = Refl, eq2 = Refl}) =
  let (Refl, Refl, Refl) = eqvIndices eqv in (Refl, Refl, Refl)

export
0 eqvMask : (p : OPE m1 n mask1) -> (q : OPE m2 n mask2) ->
            mask1 = mask2 -> p `Eqv` q
eqvMask Stop Stop _ = EqvStop
eqvMask (Drop ope1 Refl) (Drop {mask = mm2} ope2 eq2) Refl =
  EqvDrop $ eqvMask ope1 ope2 (doubleInj _ _ eq2)
eqvMask (Drop ope1 Refl) (Keep ope2 eq2) Refl =
  void $ notEvenOdd _ _ eq2
eqvMask (Keep ope1 eq1) (Keep ope2 eq2) Refl =
  EqvKeep $ eqvMask ope1 ope2 (doubleInj _ _ $ inj S $ trans (sym eq1) eq2)
eqvMask (Keep ope1 eq1) (Drop ope2 eq2) Refl =
  void $ notEvenOdd _ _ $ trans (sym eq2) eq1

uip : (p, q : a = b) -> p = q
uip Refl Refl = Refl

export
0 eqvEq' : (p, q : OPE m n mask) -> p `Eqv` q -> p === q
eqvEq' Stop Stop EqvStop = Refl
eqvEq' (Drop p eq1) (Drop q eq2) (EqvDrop eqv)
    with (doubleInj _ _ $ trans (sym eq1) eq2)
  _ | Refl = cong2 Drop (eqvEq' p q eqv) (uip eq1 eq2)
eqvEq' (Keep p eq1) (Keep q eq2) (EqvKeep eqv)
    with (doubleInj _ _ $ inj S $ trans (sym eq1) eq2)
  _ | Refl = cong2 Keep (eqvEq' p q eqv) (uip eq1 eq2)


export
0 eqvEq : (p : OPE m1 n1 mask1) -> (q : OPE m2 n2 mask2) ->
          p `Eqv` q -> p ~=~ q
eqvEq p q eqv = let (Refl, Refl, Refl) = eqvIndices eqv in eqvEq' p q eqv


public export
data View : OPE m n mask -> Type where
  StopV : View Stop
  DropV : (mask : Nat) -> (0 ope : OPE m n mask) -> View (Drop ope Refl)
  KeepV : (mask : Nat) -> (0 ope : OPE m n mask) -> View (Keep ope Refl)
%name Thin.View v

private
0 stopEqs : OPE m 0 mask -> (m = 0, mask = 0)
stopEqs Stop = (Refl, Refl)

private
0 fromStop : (ope : OPE 0 0 0) -> ope = Stop
fromStop Stop = Refl

private
0 fromDrop : (ope : OPE m (S n) (k + k)) ->
             (inner : OPE m n k ** ope === Drop inner Refl)
fromDrop (Drop ope eq) with (doubleInj _ _ eq)
  fromDrop (Drop ope Refl) | Refl = (ope ** Refl)
fromDrop (Keep ope eq) = void $ notEvenOdd _ _ eq

private
0 fromKeep : (ope : OPE (S m) (S n) (S (k + k))) ->
             (inner : OPE m n k ** ope === Keep inner Refl)
fromKeep (Drop ope eq) = void $ notEvenOdd _ _ $ sym eq
fromKeep (Keep ope eq) with (doubleInj _ _ $ inj S eq)
  fromKeep (Keep ope Refl) | Refl = (ope ** Refl)

private
0 keepIsSucc : (ope : OPE m n (S (k + k))) -> IsSucc m
keepIsSucc (Drop ope eq) = void $ notEvenOdd _ _ $ sym eq
keepIsSucc (Keep ope _)  = ItIsSucc

export
view : {0 m : Nat} -> {n, mask : Nat} -> (0 ope : OPE m n mask) -> View ope
view {n = 0} ope with 0 (fst $ stopEqs ope) | 0 (snd $ stopEqs ope)
  _ | Refl | Refl = rewrite fromStop ope in StopV
view {n = S n} ope with (half mask)
  _ | HalfOdd mask' with 0 (keepIsSucc ope)
    _ | ItIsSucc with 0 (fromKeep ope)
      _ | (ope' ** eq) = rewrite eq in KeepV mask' ope'
  _ | HalfEven mask' with 0 (fromDrop ope)
    _ | (ope' ** eq) = rewrite eq in DropV mask' ope'

export
appOpe : {0 m : Nat} -> (n : Nat) -> {mask : Nat} ->
         (0 ope : OPE m n mask) -> Singleton m
appOpe n ope with (view ope)
  appOpe 0     Stop          | StopV        = Val 0
  appOpe (S n) (Drop ope' _) | DropV _ ope' = appOpe n ope'
  appOpe (S n) (Keep ope' _) | KeepV _ ope' = [|S $ appOpe n ope'|]


||| inductive definition of OPE composition
public export
data Comp : (l : OPE n p mask1) -> (r : OPE m n mask2) ->
            (res : OPE m p mask3) -> Type where
  [search l r]
  StopZ : Comp Stop Stop Stop
  DropZ : Comp a b c -> Comp (Drop a Refl) b (Drop c Refl)
  KeepZ : Comp a b c -> Comp (Keep a Refl) (Keep b Refl) (Keep c Refl)
  KDZ   : Comp a b c -> Comp (Keep a Refl) (Drop b Refl) (Drop c Refl)

public export
record CompResult (ope1 : OPE n p mask1) (ope2 : OPE m n mask2) where
  constructor MkComp
  {resultMask : Nat}
  {0 result : OPE m p resultMask}
  0 comp : Comp ope1 ope2 result
%name CompResult comp

||| compose two OPEs, if the middle scope size is already known at runtime
export
comp' : {n, p, mask1, mask2 : Nat} ->
        (0 ope1 : OPE n p mask1) -> (0 ope2 : OPE m n mask2) ->
        CompResult ope1 ope2
comp' ope1 ope2 with (view ope1) | (view ope2)
  comp' Stop Stop | StopV | StopV =
    MkComp StopZ
  comp' (Drop ope1 Refl) ope2 | DropV _ ope1 | _ =
    MkComp $ DropZ (comp' ope1 ope2).comp
  comp' (Keep ope1 Refl) (Drop ope2 Refl) | KeepV _ ope1 | DropV _ ope2 =
    MkComp $ KDZ (comp' ope1 ope2).comp
  comp' (Keep ope1 Refl) (Keep ope2 Refl) | KeepV _ ope1 | KeepV _ ope2 =
    MkComp $ KeepZ (comp' ope1 ope2).comp

||| compose two OPEs, after recomputing the middle scope size using `appOpe`
export
comp : {p, mask1, mask2 : Nat} ->
       (0 ope1 : OPE n p mask1) -> (0 ope2 : OPE m n mask2) ->
       CompResult ope1 ope2
comp ope1 ope2 = let Val n = appOpe p ope1 in comp' ope1 ope2

-- [todo] is there a quick way to compute the mask of comp?


export
app' : OPE m1 n1 mask1 -> OPE m2 n2 mask2 -> Exists (OPE (m1 + m2) (n1 + n2))
app' Stop             ope2 = Evidence _ ope2
app' (Drop ope1 Refl) ope2 = Evidence _ $ Drop (app' ope1 ope2).snd Refl
app' (Keep ope1 Refl) ope2 = Evidence _ $ Keep (app' ope1 ope2).snd Refl

export
(++) : {n1, n2, mask1, mask2 : Nat} ->
       (0 ope1 : OPE m1 n1 mask1) -> (0 ope2 : OPE m2 n2 mask2) ->
       Subset Nat (OPE (m1 + m2) (n1 + n2))
ope1 ++ ope2 with (view ope1)
  Stop           ++ ope2 | StopV = Element _ ope2
  Drop ope1 Refl ++ ope2 | DropV mask ope1 =
    Element _ $ Drop (ope1 ++ ope2).snd Refl
  Keep ope1 Refl ++ ope2 | KeepV mask ope1 =
    Element _ $ Keep (ope1 ++ ope2).snd Refl

-- [todo] this mask is just (mask1 << n2) | mask2
-- prove it and add %transform


||| the tail of a non-empty OPE is its behaviour on all but the innermost slot
public export
data Tail : (full : OPE m1 (S n) mask1) -> (tail : OPE m2 n mask2) -> Type where
  [search full]
  DropT : Tail (Drop ope eq) ope
  KeepT : Tail (Keep ope eq) ope
%name Tail tail

public export
record TailRes (0 ope : OPE m (S n) mask) where
  constructor MkTail
  {0 scope  : Nat}
  {tailMask : Nat}
  0 tail : OPE scope n tailMask
  0 isTail : Tail ope tail

export
tail : {n, mask : Nat} -> (0 ope : OPE m (S n) mask) -> TailRes ope
tail ope with (view ope)
  tail (Drop ope _) | DropV _ ope = MkTail ope DropT
  tail (Keep ope _) | KeepV _ ope = MkTail ope KeepT


namespace OPEList
  ||| a list of OPEs of a given outer scope size
  public export
  data OPEList : Nat -> Type where
    Nil  : OPEList n
    (::) : OPE m n mask -> OPEList n -> OPEList n
  %name OPEList opes

namespace Tails -- 🦊
  ||| `Tails opes tails` if each i'th element of `tails` is the tail of
  ||| the i'th element of `opes`
  public export
  data Tails : (full : OPEList (S n)) -> (tails : OPEList n) -> Type where
    [search full]
    Nil  : Tails [] []
    (::) : Tail ope tail -> Tails opes tails ->
           Tails (ope :: opes) (tail :: tails)

namespace Cover
  ||| an OPE list is a cover if at least one of the OPEs has `Keep` as the head,
  ||| and the tails are also a cover
  |||
  ||| in @egtbs it is a binary relation which is fine for ×ᵣ but i don't want to
  ||| write my AST in universe-of-syntaxes style. sorry
  public export data Cover : OPEList n -> Type

  ||| the "`Keep` in the head" condition of a cover
  public export data Cover1 : OPEList n -> Type

  data Cover where
    Nil  : Cover opes {n = 0}
    (::) : Cover1 opes -> Tails opes tails => Cover tails -> Cover opes
  %name Cover cov

  data Cover1 where
    Here  : Cover1 (Keep p eq :: opes)
    There : Cover1 opes -> Cover1 (ope :: opes)
  %name Cover1 cov1


public export
record Coprod (ope1 : OPE m1 n mask1) (ope2 : OPE m2 n mask2) where
  constructor MkCoprod
  {0 size    : Nat}
  {sizeMask  : Nat}
  {leftMask  : Nat}
  {rightMask : Nat}
  {0 sub     : OPE size n sizeMask}
  {0 left    : OPE m1 size leftMask}
  {0 right   : OPE m2 size rightMask}
  0 leftComp  : Comp sub left  ope1
  0 rightComp : Comp sub right ope2
  {auto 0 isCover : Cover [left, right]}
%name Coprod cop

export
coprod : {n, mask1, mask2 : Nat} ->
         (0 ope1 : OPE m1 n mask1) -> (0 ope2 : OPE m2 n mask2) ->
         Coprod ope1 ope2
coprod ope1 ope2 with (view ope1) | (view ope2)
  coprod Stop Stop | StopV | StopV = MkCoprod StopZ StopZ
  coprod (Drop ope1 Refl) (Drop ope2 Refl) | DropV _ ope1 | DropV _ ope2 =
    let MkCoprod l r = coprod ope1 ope2 in MkCoprod (DropZ l) (DropZ r)
  coprod (Drop ope1 Refl) (Keep ope2 Refl) | DropV _ ope1 | KeepV _ ope2 =
    let MkCoprod l r = coprod ope1 ope2 in MkCoprod (KDZ l) (KeepZ r)
  coprod (Keep ope1 Refl) (Drop ope2 Refl) | KeepV _ ope1 | DropV _ ope2 =
    let MkCoprod l r = coprod ope1 ope2 in MkCoprod (KeepZ l) (KDZ r)
  coprod (Keep ope1 Refl) (Keep ope2 Refl) | KeepV _ ope1 | KeepV _ ope2 =
    let MkCoprod l r = coprod ope1 ope2 in MkCoprod (KeepZ l) (KeepZ r)

-- [todo] n-ary coprod


public export
record Chunks m n where
  constructor MkChunks
  {leftMask  : Nat}
  {rightMask : Nat}
  0 left  : OPE m (m + n) leftMask
  0 right : OPE n (m + n) rightMask
  {auto 0 isCover : Cover [left, right]}
%name Chunks chunks

export
chunks : (m, n : Nat) -> Chunks m n
chunks 0 0 = MkChunks Stop Stop
chunks 0 (S n) =
  let MkChunks l r = chunks 0 n in
  MkChunks (Drop l Refl) (Keep r Refl)
chunks (S m) n =
  let MkChunks l r = chunks m n in
  MkChunks (Keep l Refl) (Drop r Refl)

-- [todo] the masks here are just ((2 << m) - 1) << n and (2 << n) - 1


public export
record SplitAt m n1 n2 (ope : OPE m (n1 + n2) mask) where
  constructor MkSplitAt
  {leftMask, rightMask : Nat}
  {0 leftScope, rightScope : Nat}
  0 left     : OPE leftScope n1 leftMask
  0 right    : OPE rightScope n2 rightMask
  0 scopePrf : m = leftScope + rightScope
  0 opePrf   : ope `Eqv` (left `app'` right).snd
%name SplitAt split

export
splitAt : (n1 : Nat) -> {n2, mask : Nat} -> (0 ope : OPE m (n1 + n2) mask) ->
          SplitAt m n1 n2 ope
splitAt 0     ope = MkSplitAt zero ope Refl reflexive
splitAt (S n1) ope with (view ope)
  splitAt (S n1) (Drop ope Refl) | DropV _ ope with (splitAt n1 ope)
    _ | MkSplitAt left right scopePrf opePrf =
      MkSplitAt (Drop left Refl) right scopePrf (EqvDrop opePrf)
  splitAt (S n1) (Keep ope Refl) | KeepV _ ope with (splitAt n1 ope)
    _ | MkSplitAt left right scopePrf opePrf =
      MkSplitAt (Keep left Refl) right (cong S scopePrf) (EqvKeep opePrf)


public export
record Thinned f n where
  constructor Th
  {0 scope   : Nat}
  {scopeMask : Nat}
  0 ope : OPE scope n scopeMask
  term  : f scope
%name Thinned s, t, u

export
pure : {n : Nat} -> f n -> Thinned f n
pure term = Th id.snd term

export
join : {n : Nat} -> Thinned (Thinned f) n -> Thinned f n
join (Th ope1 (Th ope2 term)) = Th (comp ope1 ope2).result term