module RedBlackTree where

open import stack
open import Level

record TreeMethods {n m : Level } {a : Set n } {t : Set m } (treeImpl : Set n ) : Set (m Level.⊔ n) where
  field
    putImpl : treeImpl -> a -> (treeImpl -> t) -> t
    getImpl  : treeImpl -> (treeImpl -> Maybe a -> t) -> t
open TreeMethods

record Tree  {n m : Level } {a : Set n } {t : Set m } (treeImpl : Set n ) : Set (m Level.⊔ n) where
  field
    tree : treeImpl
    treeMethods : TreeMethods {a} {t}
  putTree : a -> (Tree -> t) -> t
  putTree d next = putImpl (treeMethods ) tree d (\t1 -> next (record {tree = t1 ; treeMethods = treeMethods} ))
  getTree : (Tree -> Maybe a -> t) -> t
  getTree next = getImpl (treeMethods ) tree (\t1 d -> next (record {tree = t1 ; treeMethods = treeMethods} ) d )

open Tree

data Color : Set where
  Red   : Color
  Black : Color

record Node (a : Set) : Set where
  field
    node  : Element a
    right : Maybe (Node a)
    left  : Maybe (Node a)
    color : Color

record RedBlackTree (a si : Set) : Set where
  field
    root : Maybe (Node a)
    stack : Stack si

open RedBlackTree

insertNode : ?
insertNode tree datum next = get2 (stack tree) (\ s d1 d2 -> insertCase1 ( record { root = root tree; stack = s }) datum d1 d2 next)

putRedBlackTree : {Data t : Set} -> RedBlackTree Data -> Data -> (Code : RedBlackTree Data -> t) -> t
putRedBlackTree tree datum next with (root tree)
...                                | Nothing = insertNode tree datum next
...                                | Just n  = findNode tree datum n (\ tree1 -> insertNode tree1 datum next)

findNode : {Data t : Set} -> RedBlackTree Data -> Data -> Node Data -> (Code : RedBlackTree Data (RedBlackTree Data -> t) -> t) -> t
findNode tree datum n next = pushStack (stack tree) n (\ s -> findNode1 (record tree {stack = s }) datum n next)

findNode1 : {Data t : Set} -> RedBlackTree Data -> Data -> Data -> (Code : RedBlackTree Data (RedBlackTree Data -> t) -> t) -> t
findNode1 tree datum n next with (compare datum n)
...                                | EQ = popStack (tree stack) (\s d -> findNode3 d (record tree { root = just (record n {node = datum}); stack = s }) next)
...                                | GT = findNode2 tree datum (right n) next
...                                | LT = findNode2 tree datum (left n) next
  where
    findNode2 tree datum nothing next = insertNode tree datum next
    findNode2 tree datum (just n) next = findNode (record tree {root = just n}) datum n next
    findNode3 nothing tree next = next tree
    findNode3 (just n) tree next = 
           popStack (tree stack) (\s d -> findNode3 d (record { root = record n {right = ? } }))


insertCase1 tree datum nothing grandparent next = next (record { root = ?; stack = createSingleLinkedStack })
insertCase1 tree datum (just parent) grandparent next = insertCase2 tree datum parent grandparent next

insertCase2 tree datum parent grandparent next with (color parent)
...                                | Red = insertCase3 tree datum parent grandparent next
...                                | Black = next (record { root = ?; stack = createSingleLinkedStack })

insertCase3 tree datum parent grandparent next 

getRedBlackTree : {a t : Set} -> RedBlackTree a -> (Code : RedBlackTree a -> (Maybe a) -> t) -> t
getRedBlackTree tree cs with (root tree)
...                                | Nothing = cs tree  Nothing
...                                | Just d  = cs stack1 (Just data1)
  where
    data1  = datum d
    stack1 = record { root = (next d) }