module Quox.OPE.Sub

import Quox.OPE.Basics
import Quox.OPE.Length
import Quox.NatExtra
import Data.DPair
import Data.SnocList.Elem
import Data.SnocList.Quantifiers

%default total

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

export
keepInjective : Keep p = Keep q -> p = q
keepInjective Refl = Refl

export
dropInjective : Drop p = Drop q -> p = q
dropInjective Refl = Refl


-- these need to be `public export` so that
-- `id`, `zero`, and maybe others can reduce
public export %hint
lengthLeft : xs `Sub` ys -> Length xs
lengthLeft End      = Z
lengthLeft (Keep p) = S (lengthLeft p)
lengthLeft (Drop p) = lengthLeft p

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


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

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

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


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

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



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

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

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


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

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

export
0 appZeroRight : (p : xs `Sub` ys) -> p ++ zero @{len} {xs = [<]} = p
appZeroRight {len = Z} p = Refl


||| if `p` holds for all elements of the main list,
||| it holds for all elements of the sublist
public export
subAll : xs `Sub` 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

||| if `p` holds for one element of the sublist,
||| it holds for one element of the main list
public export
subAny : xs `Sub` ys -> Any p xs -> Any p ys
subAny (Keep q) (Here  x) = Here x
subAny (Keep q) (There x) = There (subAny q x)
subAny (Drop q) x         = There (subAny q x)


public export
data SubMaskView : (lte : xs `Sub` ys) -> (mask : Nat) -> Type where
  [search lte]
  END  : SubMaskView End 0
  KEEP : {n : Nat} -> {0 p : xs `Sub` ys} ->
         (0 v : SubMaskView p n) -> SubMaskView (Keep {z} p) (S (2 * n))
  DROP : {n : Nat} -> {0 p : xs `Sub` ys} ->
         (0 v : SubMaskView p n) -> SubMaskView (Drop {z} p)    (2 * n)
%name SubMaskView v

public export
record SubMask {a : Type} (xs, ys : Scope a) where
  constructor SubM
  mask : Nat
  0 lte : xs `Sub` ys
  0 view0 : SubMaskView lte mask
%name SubMask m

private
0 ltemNilLeftZero' : SubMaskView {xs = [<]} lte mask -> mask = 0
ltemNilLeftZero' END      = Refl
ltemNilLeftZero' (DROP v) = cong (2 *) $ ltemNilLeftZero' v

export
ltemNilLeftZero : (0 _ : SubMaskView {xs = [<]} lte mask) -> mask = 0
ltemNilLeftZero v = irrelevantEq $ ltemNilLeftZero' v


private
0 lteNilLeftDrop0 : (p : [<] `Sub` (xs :< x)) -> (q ** p = Drop q)
lteNilLeftDrop0 (Drop q) = (q ** Refl)

private
lteNilLeftDrop : (0 p : [<] `Sub` (xs :< x)) -> Exists (\q => p = Drop q)
lteNilLeftDrop q =
  let 0 res = lteNilLeftDrop0 q in
  Evidence res.fst (irrelevantEq res.snd)

private
0 lteNil2End : (p : [<] `Sub` [<]) -> p = End
lteNil2End End = Refl

private
0 ltemEnd' : SubMaskView p n -> p = End -> n = 0
ltemEnd' END Refl = Refl

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

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

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

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

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

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

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

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

export
view : Length xs => Length ys =>
       (m : SubMask xs ys) -> SubMaskView m.lte m.mask
view @{Z} @{Z} (SubM {lte, view0, _}) =
  rewrite lteNil2End lte in
  rewrite ltemEnd' view0 (lteNil2End lte) in
  END
view @{S _} @{Z} (SubM {lte, _}) = void $ absurd lte
view @{Z} @{S sy} (SubM mask lte view0) with (ltemNilLeftZero view0)
  view @{Z} @{S sy} (SubM 0 lte view0)
      | Refl with (lteNilLeftDrop lte)
    view @{Z} @{S sy} (SubM 0 (Drop lte) view0)
        | Refl | Evidence lte Refl =
            DROP {n = 0} $ let DROP {n = 0} p = view0 in p
view @{S sx} @{S sy} (SubM mask lte view0) with (viewLsb mask)
  view @{S sx} @{S sy} (SubM (2 * n) lte view0)
      | Evidence Even (Lsb0 n) with (ltemEven view0)
    view @{S sx} @{S sy} (SubM (2 * m) (Drop lte) view0)
        | Evidence Even (Lsb0 m) | Evidence lte Refl =
            DROP $ fromDROP' Refl view0
  view @{S sx} @{S sy} (SubM (S (2 * n)) lte view0)
      | Evidence Odd (Lsb1 n) with (ltemOddHead view0)
    view @{S sx} @{S sy} (SubM (S (2 * n)) lte view0)
        | Evidence Odd (Lsb1 n) | Refl with (ltemOdd view0)
      view @{S sx} @{S sy} (SubM (S (2 * n)) (Keep lte) view0)
          | Evidence Odd (Lsb1 n) | Refl | Evidence lte Refl =
              KEEP $ fromKEEP' Refl view0

export
(.view) : Length xs => Length ys =>
          (m : SubMask xs ys) -> SubMaskView m.lte m.mask
(.view) = view


export
ltemLen : Length xs => Length ys =>
          xs `SubMask` ys -> length xs `LTE` length ys
ltemLen @{sx} @{sy} lte@(SubM m l _) with (lte.view)
  ltemLen @{sx} @{sy} lte@(SubM 0 End _) | END = LTEZero
  ltemLen @{S sx} @{S sy} lte@(SubM (S (2 * n)) (Keep p) _) | (KEEP q) =
    LTESucc $ ltemLen $ SubM n p q
  ltemLen @{sx} @{S sy} lte@(SubM (2 * n) (Drop p) _) | (DROP q) =
    lteSuccRight $ ltemLen $ SubM n p q

export
ltemNilRight : xs `SubMask` [<] -> xs = [<]
ltemNilRight m = irrelevantEq $ lteNilRight m.lte