module AATree where

open import Relation.Binary

postulate
  A   : Set
  _<_ : A → A → Set
  _≈A_ : A → A → Set
  trans : {a b c : A} → a < b → b < c → a < c
  compare : ∀ x y → Tri (x < y) (x ≈A y) (y < x)
 
open import Data.Unit
open import Data.Product
import Data.Nat as Nat
open Nat hiding (compare) renaming (_<_ to _<ℕ_; _≤_ to _≤ℕ_)
open import Data.Nat.Properties
open DecTotalOrder Nat.decTotalOrder hiding (_≈_; refl) renaming (trans to transℕ)
open StrictTotalOrder strictTotalOrder hiding (_<_; trans) renaming (compare to compareℕ)
import Relation.Binary.PropositionalEquality as RBP
open RBP hiding (trans)

{-
  Invariants:
  0. Binary search tree property.
  1. The level of a leaf node is one.
  2. The level of a left child is strictly less than that of its parent.
  3. The level of a right child is less than or equal to that of its parent.
  4. The level of a right grandchild is strictly less than that of its grandparent.
  5. Every node of level greater than one must have two children.
-} 

-- Implementation
------------------------------------------------------------------------
data Tree : Set where
  lf : Tree
  nd : A → ℕ → Tree → Tree → Tree 

level : Tree → ℕ
level lf           =  
level (nd _ n _ _) = n 

skew : Tree → Tree
skew (nd a n (nd la ln ll lr) r) with compareℕ n ln
... | tri≈ _ _ _ = nd la ln ll (nd a n lr r) 
... | _          = nd a n (nd la ln ll lr) r
skew t = t

split : Tree → Tree
split (nd a n l (nd ra rn rl rr)) with compareℕ n (level rr)  
... | tri≈ _ _ _ = nd ra (rn + ) (nd a n l rl) rr 
... | _          = nd a n l (nd ra rn rl rr)
split t = t

insert : A → Tree → Tree
insert a lf           = nd a  lf lf 
insert a {-t@-}(nd b n l r) with compare a b
... | tri< _ _ _ = split (skew (nd b n (insert a l) r))
... | tri≈ _ _ _ = {-t-} nd b n l r
... | tri> _ _ _ = split (skew (nd b n l (insert a r)))

-- Proof
------------------------------------------------------------------------

_>T_ : A → Tree → Set
_ >T lf           = ⊤
a >T (nd b _ l r) = (b < a) × (a >T l) × (a >T r)

_<T_ : A → Tree → Set
_ <T lf           = ⊤
a <T (nd b _ l r) = (a < b) × (a <T l) × (a <T r)

Sorted : Tree → Set 
Sorted lf           = ⊤
Sorted (nd a _ l r) = (a >T l × Sorted l) × (a <T r × Sorted r)

sorted-split : (t : Tree) → Sorted t → Sorted (split t)
sorted-split lf            _ = tt
sorted-split (nd a n l lf) p = p 
sorted-split (nd a n lf (nd ra rn rl rr)) 
  {-p@-}(p , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
   with compareℕ n (level rr)
... | tri≈ _ _ _ = ((a<ra , tt , ra>Trl) , (tt , tt) , a<Trl , srl) , ra<Trr , srr
... | tri< _ _ _ = {-p-}(p , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
... | tri> _ _ _ = {-p-}(p , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
sorted-split (nd a n (nd la ln ll lr) (nd ra rn rl rr)) 
  {-p@-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
  , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
   with compareℕ n (level rr)
... | tri≈ _ _ _ = ((a<ra , (trans la<a a<ra , lem ll a<ra a>Tll , lem lr a<ra a>Tlr) , ra>Trl)
  , ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , la<Tlr , slr) ,
      a<Trl , srl) , ra<Trr , srr
  where
    lem : {a b : A}(t : Tree) → b < a → b >T t → a >T t
    lem lf           _   _                   = tt
    lem (nd x n l r) b<a (x<b , b>Tl , b>Tr) = trans x<b b<a , lem l b<a b>Tl , lem r b<a b>Tr
... | tri< _ _ _ = {-p-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
                   , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
... | tri> _ _ _ = {-p-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
                   , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))

sorted-skew : (t : Tree) → Sorted t → Sorted (skew t)
sorted-skew lf                           _ = tt
sorted-skew (nd _ _ lf _)                p = p
sorted-skew (nd a n (nd la ln ll lr) lf)
  {-p@-}(((la<a , a>Tll , a>Tlr) , (la<Trl , sll) , (la<Tlr , slr)) , _) with compareℕ n ln
... | tri≈ _ _ _ = (la<Trl , sll) , (la<a , la<Tlr , tt) , (a>Tlr , slr) , tt , tt
... | tri< _ _ _ = {-p-}(((la<a , a>Tll , a>Tlr) , (la<Trl , sll) , (la<Tlr , slr)) , (tt , tt))
... | tri> _ _ _ = {-p-}(((la<a , a>Tll , a>Tlr) , (la<Trl , sll) , (la<Tlr , slr)) , (tt , tt))
sorted-skew (nd a n (nd la ln ll lr) (nd ra rn rl rr))
  {-p@-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
  , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
  with compareℕ n ln
... | tri≈ _ _ _ = (la>Tll , sll) 
                   , ((la<a , la<Tlr , trans la<a a<ra , lem rl la<a a<Trl , lem rr la<a a<Trr)
                   , (a>Tlr , slr) , (a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , ra<Trr , srr)
  where
    lem : {a b : A}(t : Tree) → a < b → b <T t → a <T t
    lem lf           _   _                   = tt
    lem (nd x n l r) a<b (b<x , b<Tl , b<Tr) = trans a<b b<x , lem l a<b b<Tl , lem r a<b b<Tr
... | tri< _ _ _ = {-p-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
                   , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))
... | tri> _ _ _ = {-p-}( ((la<a , a>Tll , a>Tlr) , (la>Tll , sll) , (la<Tlr , slr))
                   , ((a<ra , a<Trl , a<Trr) , (ra>Trl , srl) , (ra<Trr , srr)))

sorted-split-skew : (t : Tree) → Sorted t → Sorted (split (skew t))
sorted-split-skew t p = sorted-split (skew t) (sorted-skew t p) 

-- g-skew and l-skew are symmetric 
g-split : (b : A)(t : Tree) → b >T t → b >T split t
g-split _ lf            _ = tt
g-split _ (nd _ _ _ lf) p = p
g-split b (nd c n l (nd ra rn rl rr)) {-p@-}(c<b , pl , (ra<b , b>Trl , b>Trr))
  with compareℕ n (level rr)
... | tri≈ _ _ _ = ra<b , (c<b , pl , b>Trl) , b>Trr
... | tri< _ _ _ = {-p-}c<b , pl , ra<b , b>Trl , b>Trr
... | tri> _ _ _ = {-p-}c<b , pl , ra<b , b>Trl , b>Trr

l-split : (b : A)(t : Tree) → b <T t → b <T split t
l-split _ lf            _ = tt
l-split _ (nd _ _ _ lf) p = p
l-split b (nd c n l (nd ra rn rl rr)) {-p@-}(c<b , pl , (ra<b , b>Trl , b>Trr))
  with compareℕ n (level rr)
... | tri≈ _ _ _ = ra<b , (c<b , pl , b>Trl) , b>Trr
... | tri< _ _ _ = {-p-}c<b , pl , ra<b , b>Trl , b>Trr
... | tri> _ _ _ = {-p-}c<b , pl , ra<b , b>Trl , b>Trr

-- g-skew and l-skew are symmetric 
g-skew : (b : A)(t : Tree) → b >T t → b >T skew t
g-skew _ lf _ = tt
g-skew _ (nd _ _ lf r) p = p
g-skew b (nd a n (nd la ln ll lr) r) {-p@-}(a<b , (la<b , b>Tll , b>Tlr) , pr) with compareℕ n ln
... | tri≈ _ _ _ = la<b , b>Tll , a<b , b>Tlr , pr
... | tri< _ _ _ = {-p-}(a<b , (la<b , b>Tll , b>Tlr) , pr)
... | tri> _ _ _ = {-p-}(a<b , (la<b , b>Tll , b>Tlr) , pr)

l-skew : (b : A)(t : Tree) → b <T t → b <T skew t
l-skew _ lf _ = tt
l-skew _ (nd _ _ lf r) p = p
l-skew b (nd a n (nd la ln ll lr) r) {-p@-}(a<b , (la<b , b>Tll , b>Tlr) , pr) with compareℕ n ln
... | tri≈ _ _ _ = la<b , b>Tll , a<b , b>Tlr , pr 
... | tri< _ _ _ = {-p-}(a<b , (la<b , b>Tll , b>Tlr) , pr)
... | tri> _ _ _ = {-p-}(a<b , (la<b , b>Tll , b>Tlr) , pr)

g-split-skew : (b : A)(t : Tree) → b >T t → b >T split (skew t)
g-split-skew b t p = g-split b (skew t) (g-skew b t p)

l-split-skew : (b : A)(t : Tree) → b <T t → b <T split (skew t)
l-split-skew b t p = l-split b (skew t) (l-skew b t p)

-- lemma¹ and lemmaʳ are symmmetric.
lemmaˡ : (a b : A) → (t : Tree) → a < b → b >T t → b >T (insert a t)
lemmaˡ _ _ lf a<b _ = (a<b , tt , tt)
lemmaˡ a b (nd c n l r) a<b (c<b , b>Tl , b>Tr) with compare a c
... | tri< a<c _ _ = g-split-skew b (nd c n (insert a l) r) (c<b , lemmaˡ a b l a<b b>Tl , b>Tr) 
... | tri≈ _ _ _ = c<b , b>Tl , b>Tr
... | tri> _ _ c<a = g-split-skew b (nd c n l (insert a r)) (c<b , b>Tl , lemmaˡ a b r a<b b>Tr)

lemmaʳ : (a b : A) → (t : Tree) → b < a → b <T t → b <T (insert a t)
lemmaʳ _ _ lf a<b _ = (a<b , tt , tt)
lemmaʳ a b (nd c n l r) a<b (c<b , b>Tl , b>Tr) with compare a c
... | tri< a<c _ _ = l-split-skew b (nd c n (insert a l) r) (c<b , lemmaʳ a b l a<b b>Tl , b>Tr) 
... | tri≈ _ _ _ = c<b , b>Tl , b>Tr
... | tri> _ _ c<a = l-split-skew b (nd c n l (insert a r)) (c<b , b>Tl , lemmaʳ a b r a<b b>Tr)

sorted : (a : A)(t : Tree) → Sorted t → Sorted (insert a t)
sorted _ lf           _ = (tt , tt) , (tt , tt)
sorted a (nd b n l r) {-p@-}((b>Tl , sl) , (b<Tr , sr))  with compare a b
... | tri< a<b _ _ = sorted-split-skew (nd b n (insert a l) r)
  ((lemmaˡ a b l a<b b>Tl , sorted a l sl) , b<Tr , sr)
... | tri≈ _ _ _ = {-p-}((b>Tl , sl) , (b<Tr , sr))
... | tri> _ _ b<a = sorted-split-skew (nd b n l (insert a r)) 
  ((b>Tl , sl) , lemmaʳ a b r b<a b<Tr , sorted a r sr)