||| "order preserving embeddings", for recording a correspondence between
||| a smaller scope and part of a larger one.
module Quox.OPE

import Quox.NatExtra
import public Data.DPair
import public Data.SnocList
import public Data.SnocList.Elem

%default total

LTE_n = Nat.LTE
%hide Nat.LTE


public export
Scope : Type -> Type
Scope = SnocList

public export
data LTE : Scope a -> Scope a -> Type where
  End  : [<] `LTE` [<]
  Keep : xs `LTE` ys -> xs :< z `LTE` ys :< z
  Drop : xs `LTE` ys -> xs      `LTE` ys :< z
%name LTE p, q

-- [todo] bitmask representation???


export
dropLast : (xs :< x) `LTE` ys -> xs `LTE` ys
dropLast (Keep p) = Drop p
dropLast (Drop p) = Drop $ dropLast p

export
Uninhabited (xs :< x `LTE` [<]) where uninhabited _ impossible

export
Uninhabited (xs :< x `LTE` xs) where
  uninhabited (Keep p) = uninhabited p
  uninhabited (Drop p) = uninhabited $ dropLast p


export
0 lteLen : xs `LTE` ys -> length xs `LTE_n` length ys
lteLen End      = LTEZero
lteLen (Keep p) = LTESucc $ lteLen p
lteLen (Drop p) = lteSuccRight $ lteLen p

export
0 lteNilRight : xs `LTE` [<] -> xs = [<]
lteNilRight End = Refl

export
0 lteNilLeftDrop : (p : [<] `LTE` (xs :< x)) -> Exists (\q => p = Drop q)
lteNilLeftDrop (Drop q) = Evidence q Refl

export
0 lteNil2End : (p : [<] `LTE` [<]) -> p = End
lteNil2End End = Refl


public export
data Length : Scope a -> Type where
  Z : Length [<]
  S : (s : Length xs) -> Length (xs :< x)
%name Length s
%builtin Natural Length

namespace Length
  public export
  (.nat) : Length xs -> Nat
  (Z).nat   = Z
  (S s).nat = S s.nat
  %transform "Length.nat" Length.(.nat) xs = believe_me xs

  export
  0 lengthOk : (s : Length xs) -> s.nat = length xs
  lengthOk Z     = Refl
  lengthOk (S s) = cong S $ lengthOk s

export %hint
lengthLeft : xs `LTE` ys -> Length xs
lengthLeft End      = Z
lengthLeft (Keep p) = S (lengthLeft p)
lengthLeft (Drop p) = lengthLeft p

export %hint
lengthRight : xs `LTE` ys -> Length ys
lengthRight End      = Z
lengthRight (Keep p) = S (lengthRight p)
lengthRight (Drop p) = S (lengthRight p)


export
id : Length xs => xs `LTE` xs
id @{Z}   = End
id @{S s} = Keep id

export
zero : Length xs => [<] `LTE` xs
zero @{Z}   = End
zero @{S s} = Drop zero

export
single : Length xs => x `Elem` xs -> [< x] `LTE` xs
single @{S _} Here      = Keep zero
single @{S _} (There p) = Drop $ single p


export
(.) : ys `LTE` zs -> xs `LTE` ys -> xs `LTE` zs
End    . End    = End
Keep p . Keep q = Keep (p . q)
Keep p . Drop q = Drop (p . q)
Drop p . q      = Drop (p . q)

export
(++) : xs1 `LTE` ys1 -> xs2 `LTE` ys2 -> (xs1 ++ xs2) `LTE` (ys1 ++ ys2)
p ++ End    = p
p ++ Keep q = Keep (p ++ q)
p ++ Drop q = Drop (p ++ q)



public export
data LTEMaskView : (lte : xs `LTE` ys) -> (mask : Nat) -> Type where
  [search lte]
  END  : LTEMaskView End 0
  KEEP : (0 _ : LTEMaskView p n) -> LTEMaskView (Keep p) (S (2 * n))
  DROP : (0 _ : LTEMaskView p n) -> LTEMaskView (Drop p)    (2 * n)
%name LTEMaskView p, q

record LTEMask {a : Type} (xs, ys : Scope a) where
  constructor LTEM
  mask : Nat
  0 lte : xs `LTE` ys
  0 view0 : LTEMaskView lte mask
%name LTEMask m

namespace View
  private
  0 lteMaskEnd' : LTEMaskView p n -> p = End -> n = 0
  lteMaskEnd' END Refl = Refl

  private
  0 lteMaskDrop' : LTEMaskView p n -> p = Drop q -> (n' ** n = 2 * n')
  lteMaskDrop' (DROP p {n = n'}) Refl = (n' ** Refl)

  private
  0 lteMaskEven' : {p : xs `LTE` (ys :< y)} ->
                   n = 2 * n' -> LTEMaskView p n -> (q ** p = Drop q)
  lteMaskEven' eq (KEEP q) = absurd $ lsbMutex' eq Refl
  lteMaskEven' eq (DROP q) = (_ ** Refl)

  private
  lteMaskEven : {0 p : xs `LTE` (ys :< y)} ->
                (0 _ : LTEMaskView p (2 * n)) -> Exists (\q => p = Drop q)
  lteMaskEven q =
    let 0 res = lteMaskEven' Refl q in
    Evidence res.fst (irrelevantEq res.snd)

  private
  0 fromDROP' : {lte : xs `LTE` ys} -> n = 2 * n' ->
                LTEMaskView (Drop lte) n -> LTEMaskView lte n'
  fromDROP' eq (DROP {n} p) =
    let eq = doubleInj eq {m = n, n = n'} in
    rewrite sym eq in p

  private
  0 fromDROP : LTEMaskView (Drop lte) (2 * n) -> LTEMaskView lte n
  fromDROP = fromDROP' Refl

  private
  0 lteMaskOdd' : {p : (xs :< x) `LTE` (ys :< x)} -> {n' : Nat} ->
                  n = S (2 * n') -> LTEMaskView p n -> (q ** p = Keep q)
  lteMaskOdd' eq (KEEP q) = (_ ** Refl)
  lteMaskOdd' eq (DROP q) = absurd $ lsbMutex' Refl eq
  lteMaskOdd' _  END      impossible

  private
  lteMaskOdd : (0 _ : LTEMaskView p (S (2 * n))) -> Exists (\q => p = Keep q)
  lteMaskOdd q =
    let 0 res = lteMaskOdd' Refl q in
    Evidence res.fst (irrelevantEq res.snd)

  private
  0 lteMaskOddHead' : {p : (xs :< x) `LTE` (ys :< y)} -> {n' : Nat} ->
                      n = S (2 * n') -> LTEMaskView p n -> x = y
  lteMaskOddHead' eq (KEEP q) = Refl
  lteMaskOddHead' eq (DROP q) = absurd $ lsbMutex' Refl eq
  lteMaskOddHead' eq END      impossible

  private
  lteMaskOddHead : {0 p : (xs :< x) `LTE` (ys :< y)} ->
                   (0 _ : LTEMaskView p (S (2 * n))) -> x = y
  lteMaskOddHead q = irrelevantEq $ lteMaskOddHead' Refl q

  private
  0 fromKEEP' : {lte : xs `LTE` ys} -> n = S (2 * n') ->
                LTEMaskView (Keep lte) n -> LTEMaskView lte n'
  fromKEEP' eq (KEEP {n} p) =
    let eq = doubleInj (injective eq) {m = n, n = n'} in
    rewrite sym eq in p

  private
  0 fromKEEP : LTEMaskView (Keep lte) (S (2 * n)) -> LTEMaskView lte n
  fromKEEP = fromKEEP' Refl

  export
  view : (sx : Length xs) => (sy : Length ys) =>
         (m : LTEMask xs ys) -> LTEMaskView m.lte m.mask
  view @{Z} @{Z} (LTEM {lte, view0, _}) =
    rewrite lteNil2End lte in
    rewrite lteMaskEnd' view0 (lteNil2End lte) in
    END
  view @{S _} @{Z} (LTEM {lte, _}) = void $ absurd lte
  view @{Z} @{S sy} (LTEM mask lte view0) =
    rewrite (lteNilLeftDrop lte).snd in
    rewrite (lteMaskDrop' view0 (lteNilLeftDrop lte).snd).snd in
    DROP $ let DROP p = view0 in p
  view @{S sx} @{S sy} (LTEM mask lte view0) with (viewLsb mask)
    view @{S sx} @{S sy} (LTEM (2 * n) lte view0)
        | Evidence Even (Lsb0 n) with (lteMaskEven view0)
      view @{S sx} @{S sy} (LTEM (2 * m) (Drop lte) view0)
          | Evidence Even (Lsb0 m) | Evidence lte Refl =
              DROP $ fromDROP view0
    view @{S sx} @{S sy} (LTEM (S (2 * n)) lte view0)
        | Evidence Odd (Lsb1 n) with (lteMaskOddHead view0)
      view @{S sx} @{S sy} (LTEM (S (2 * n)) lte view0)
          | Evidence Odd (Lsb1 n) | Refl with (lteMaskOdd view0)
        view @{S sx} @{S sy} (LTEM (S (2 * n)) (Keep lte) view0)
            | Evidence Odd (Lsb1 n) | Refl | Evidence lte Refl =
                KEEP $ fromKEEP view0



public export
record Split {a : Type} (xs, ys, zs : Scope a) (p : xs `LTE` ys ++ zs) where
  constructor MkSplit
  {0 leftSub, rightSub : Scope a}
  leftThin : leftSub `LTE` ys
  rightThin : rightSub `LTE` zs
  0 eqScope : xs = leftSub ++ rightSub
  0 eqThin : p ~=~ leftThin ++ rightThin

export
split : (zs : Scope a) -> (p : xs `LTE` ys ++ zs) -> Split xs ys zs p
split [<] p = MkSplit p zero Refl Refl
split (zs :< z) (Keep p) with (split zs p)
  split (zs :< z) (Keep (l ++ r)) | MkSplit l r Refl Refl =
    MkSplit l (Keep r) Refl Refl
split (zs :< z) (Drop p) {xs} with (split zs p)
  split (zs :< z) (Drop (l ++ r)) {xs = _} | MkSplit l r Refl Refl =
    MkSplit l (Drop r) Refl Refl


public export
data Comp : ys `LTE` zs -> xs `LTE` ys -> xs `LTE` zs -> Type where
  CEE : Comp End End End
  CKK : Comp p q pq -> Comp (Keep p) (Keep q) (Keep pq)
  CKD : Comp p q pq -> Comp (Keep p) (Drop q) (Drop pq)
  CD0 : Comp p q pq -> Comp (Drop p) q        (Drop pq)

export
comp : (p : ys `LTE` zs) -> (q : xs `LTE` ys) -> Comp p q (p . q)
comp End      End      = CEE
comp (Keep p) (Keep q) = CKK (comp p q)
comp (Keep p) (Drop q) = CKD (comp p q)
comp (Drop p) q        = CD0 (comp p q)

export
0 compOk : Comp p q r -> r = (p . q)
compOk CEE     = Refl
compOk (CKK z) = cong Keep $ compOk z
compOk (CKD z) = cong Drop $ compOk z
compOk (CD0 z) = cong Drop $ compOk z

export
compZero : (sx : Length xs) => (sy : Length ys) =>
           (p : xs `LTE` ys) -> Comp p (OPE.zero @{sx}) (OPE.zero @{sy})
compZero {sx = Z,   sy = Z}   End      = CEE
compZero {sx = S _, sy = S _} (Keep p) = CKD (compZero p)
compZero           {sy = S _} (Drop p) = CD0 (compZero p)

export
compIdLeft : (sy : Length ys) =>
             (p : xs `LTE` ys) -> Comp (OPE.id @{sy}) p p
compIdLeft {sy = Z}   End      = CEE
compIdLeft {sy = S _} (Keep p) = CKK (compIdLeft p)
compIdLeft {sy = S _} (Drop p) = CKD (compIdLeft p)

export
compIdRight : (sx : Length xs) =>
              (p : xs `LTE` ys) -> Comp p (OPE.id @{sx}) p
compIdRight {sx = Z}   End      = CEE
compIdRight {sx = S _} (Keep p) = CKK (compIdRight p)
compIdRight            (Drop p) = CD0 (compIdRight p)


export
0 compAssoc : (p : ys `LTE` zs) -> (q : xs `LTE` ys) -> (r : ws `LTE` xs) ->
              p . (q . r) = (p . q) . r
compAssoc End      End      End      = Refl
compAssoc (Keep p) (Keep q) (Keep r) = cong Keep $ compAssoc p q r
compAssoc (Keep p) (Keep q) (Drop r) = cong Drop $ compAssoc p q r
compAssoc (Keep p) (Drop q) r        = cong Drop $ compAssoc p q r
compAssoc (Drop p) q        r        = cong Drop $ compAssoc p q r
compAssoc End (Drop _) _ impossible


public export
Scoped : Type -> Type
Scoped a = Scope a -> Type

public export
Subscope : Scope a -> Type
Subscope ys = Exists (`LTE` ys)

public export
record SubMap {a : Type} {xs, ys, zs : Scope a}
              (p : xs `LTE` zs) (q : ys `LTE` zs) where
  constructor SM
  thin : xs `LTE` ys
  0 comp : Comp q thin p

parameters (p : xs `LTE` ys)
  export
  subId : SubMap p p
  subId = SM id (compIdRight p)

  export
  subZero : SubMap OPE.zero p
  subZero = SM zero (compZero p)


public export
data All : (a -> Type) -> Scoped a where
  Lin  : All p [<]
  (:<) : All p xs -> p x -> All p (xs :< x)
%name OPE.All ps, qs

export
mapAll : (forall x. p x -> q x) -> All p xs -> All q xs
mapAll f [<]      = [<]
mapAll f (x :< y) = mapAll f x :< f y

export
subAll : xs `LTE` ys -> All p ys -> All p xs
subAll End      [<]       = [<]
subAll (Keep q) (ps :< x) = subAll q ps :< x
subAll (Drop q) (ps :< x) = subAll q ps


public export
data Cover_ : (overlap : Bool) -> xs `LTE` zs -> ys `LTE` zs -> Type where
  CE : Cover_ ov End End
  CL : Cover_ ov p q -> Cover_ ov   (Keep p) (Drop q)
  CR : Cover_ ov p q -> Cover_ ov   (Drop p) (Keep q)
  C2 : Cover_ ov p q -> Cover_ True (Keep p) (Keep q)

public export
Cover : xs `LTE` zs -> ys `LTE` zs -> Type
Cover = Cover_ True

public export
Partition : xs `LTE` zs -> ys `LTE` zs -> Type
Partition = Cover_ False