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 %inline
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
trans' : ys `Sub'` zs -> xs `Sub'` ys -> xs `Sub'` zs
trans' End      End      = End
trans' (Keep p) (Keep q) = Keep $ trans' p q
trans' (Keep p) (Drop q) = Drop $ trans' p q
trans' (Drop p) q        = Drop $ trans' p q

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


||| 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 SubView : (lte : xs `Sub'` ys) -> (mask : Nat) -> Type where
  [search lte]
  END  : SubView End 0
  KEEP : {n : Nat} -> {0 p : xs `Sub'` ys} ->
         (0 v : SubView p n) -> SubView (Keep {z} p) (S (2 * n))
  DROP : {n : Nat} -> {0 p : xs `Sub'` ys} ->
         (0 v : SubView p n) -> SubView (Drop {z} p)    (2 * n)
%name SubView v

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


export
getMask : SubView lte mask -> Subset Nat (Equal mask)

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

export
ltemNilLeftZero : (0 _ : SubView {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' : SubView p n -> p = End -> n = 0
ltemEnd' END Refl = Refl

private
0 ltemEven' : {p : xs `Sub'` (ys :< y)} ->
              n = 2 * n' -> SubView 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 _ : SubView 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' ->
              SubView (Drop lte) n -> SubView 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') -> SubView 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 _ : SubView 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') -> SubView 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 _ : SubView p (S (2 * n))) -> x = y
ltemOddHead q = irrelevantEq $ ltemOddHead' Refl q

private
0 fromKEEP' : {lte : xs `Sub'` ys} -> n = S (2 * n') ->
              SubView (Keep lte) n -> SubView 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 : Sub xs ys) -> SubView 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 : Sub xs ys) -> SubView m.lte m.mask
(.view) = view


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

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


public export %inline
end : [<] `Sub` [<]
end = SubM 0 End END

public export %inline
keep : xs `Sub` ys -> xs :< z `Sub` ys :< z
keep (SubM mask lte view0) = SubM (S (2 * mask)) (Keep lte) (KEEP view0)

public export %inline
drop : xs `Sub` ys -> xs `Sub` ys :< z
drop (SubM mask lte view0) = SubM (2 * mask) (Drop lte) (DROP view0)


export %inline
dropLast : Length xs => Length ys =>
           (xs :< x) `Sub` ys -> xs `Sub` ys
dropLast @{sx} @{sy} sub@(SubM mask lte _) with (sub.view)
  dropLast sub@(SubM (S (2 * n)) (Keep p) _) | (KEEP v) =
    SubM (2 * n) (Drop p) (DROP v)
  dropLast @{_} @{S sy} sub@(SubM (2 * n) (Drop p) _) | DROP v =
    drop $ dropLast $ SubM n p v


export
Uninhabited (xs :< x `Sub` [<]) where
  uninhabited sub = void $ uninhabited sub.lte

export
Length xs => Uninhabited (xs :< x `Sub` xs) where
  uninhabited @{sx} sub@(SubM mask lte view0) with (sub.view)
    uninhabited @{S sx} sub@(SubM (S (2 * n)) (Keep p) _) | KEEP v =
      uninhabited $ SubM n p v
    uninhabited @{S sx} sub@(SubM (2 * n) (Drop p) _) | DROP v =
      uninhabited $ dropLast $ SubM n p v


export
refl : Length xs => xs `Sub` xs
refl @{Z}   = end
refl @{S s} = keep refl

export Reflexive (Scope a) Sub where reflexive = refl

mutual
  private
  antisym_ : Length xs => Length ys =>
             {0 p : xs `Sub'` ys} -> {0 q : ys `Sub'` xs} ->
             SubView p m1 -> SubView q m2 -> xs = ys
  antisym_ END       END = Refl
  antisym_ (KEEP v1) (KEEP v2 {z}) @{S sx} @{S sy} =
    cong (:< z) $ antisym (SubM _ _ v1) (SubM _ _ v2)
  antisym_ (KEEP v1) (DROP v2) {p = Keep p} {q = Drop q} =
    void $ succNotLTEpred $ lteLen q `transitive` lteLen p
  antisym_ (DROP v1) (KEEP v2) {p = Drop p} {q = Keep q} =
    void $ succNotLTEpred $ lteLen p `transitive` lteLen q
  antisym_ (DROP v1) (DROP v2) {p = Drop p} {q = Drop q} =
    void $ succNotLTEpred $ lteLen p `transitive` lteSuccLeft (lteLen q)

  export
  antisym : Length xs => Length ys => xs `Sub` ys -> ys `Sub` xs -> xs = ys
  antisym p q = antisym_ p.view q.view

export
Antisymmetric (Scope a) Sub where
  antisymmetric p q = antisym p q


mutual
  private
  trans_ : Length xs => Length ys => Length zs =>
           {0 p : xs `Sub'` ys} -> {0 q : ys `Sub'` zs} ->
           SubView p m1 -> SubView q m2 -> xs `Sub` zs
  trans_ END END = end
  trans_ (KEEP v1) (KEEP v2) @{S sx} @{S sy} @{S sz} =
    keep $ SubM _ _ v1 `trans` SubM _ _ v2
  trans_ (DROP v1) (KEEP v2) @{sx} @{S sy} @{S sz} =
    drop $ SubM _ _ v1 `trans` SubM _ _ v2
  trans_ v1 (DROP v2) @{sx} @{sy} @{S sz} =
    let Element m1' eq = getMask v1 in
    drop $ SubM m1' _ (rewrite sym eq in v1) `trans` SubM _ _ v2

  export
  trans : Length xs => Length ys => Length zs =>
          xs `Sub` ys -> ys `Sub` zs -> xs `Sub` zs
  trans p q = trans_ p.view q.view

  export
  (.) : Length xs => Length ys => Length zs =>
        xs `Sub` ys -> ys `Sub` zs -> xs `Sub` zs
  (.) = trans

export
Transitive (Scope a) Sub where
  transitive p q = trans p q


export
zero : Length xs => [<] `Sub` xs
zero @{Z}   = end
zero @{S s} = drop zero

export
single : Length xs => x `Elem` xs -> [< x] `Sub` xs
single @{S sx} Here      = keep zero
single @{S sx} (There p) = drop $ single p


export
(++) : Length xs2 => Length ys2 =>
       xs1 `Sub` ys1 -> xs2 `Sub` ys2 -> (xs1 ++ xs2) `Sub` (ys1 ++ ys2)
(++) sub1 sub2@(SubM {}) @{sx2} @{sy2} with (sub2.view)
  (++) sub1 sub2@(SubM {}) | END = sub1
  (++) sub1 sub2@(SubM {}) @{S sx2} @{S sy2} | KEEP v = keep $ sub1 ++ SubM _ _ v
  (++) sub1 sub2@(SubM {}) @{sx2} @{S sy2} | DROP v = drop $ sub1 ++ SubM _ _ v


||| if `p` holds for all elements of the main list,
||| it holds for all elements of the sublist
export
subAll : Length xs => Length ys =>
         xs `Sub` ys -> All prop ys -> All prop xs
subAll sub@(SubM {}) ps @{sx} @{sy} with (sub.view)
  subAll sub@(SubM {}) [<] | END = [<]
  subAll sub@(SubM {}) (ps :< p) @{S sx} @{S sy} | KEEP v =
    subAll (SubM _ _ v) ps :< p
  subAll sub@(SubM {}) (ps :< p) @{sx} @{S sy} | DROP v =
    subAll (SubM _ _ v) ps


||| if `p` holds for one element of the sublist,
||| it holds for one element of the main list
export
subAny : Length xs => Length ys =>
         xs `Sub` ys -> Any prop xs -> Any prop ys
subAny sub@(SubM {}) p @{sx} @{sy} with (sub.view)
  subAny sub@(SubM {}) p | END impossible
  subAny sub@(SubM {}) (Here p) | KEEP v = Here p
  subAny sub@(SubM {}) (There p) @{S sx} @{S sy} | KEEP v =
    There $ subAny (SubM _ _ v) p
  subAny sub@(SubM {}) p @{sx} @{S sy} | DROP v =
    There $ subAny (SubM _ _ v) p