record TreeMethods {n m : Level } {a : Set n } {t : Set m } (treeImpl : Set n ) : Set (m Level.@$\sqcup$@ n) where
  field
    putImpl : treeImpl @$\rightarrow$@ a @$\rightarrow$@ (treeImpl @$\rightarrow$@ t) @$\rightarrow$@ t
    getImpl  : treeImpl @$\rightarrow$@ (treeImpl @$\rightarrow$@ Maybe a @$\rightarrow$@ t) @$\rightarrow$@ t
open TreeMethods

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

open Tree

data Color {n : Level } : Set n where
  Red   : Color
  Black : Color

data CompareResult {n : Level } : Set n where
  LT : CompareResult
  GT : CompareResult
  EQ : CompareResult

record Node {n : Level } (a k : Set n) : Set n where
  inductive
    field
      key   : k
      value : a
      right : Maybe (Node a k)
      left  : Maybe (Node a k)
      color : Color {n}
open Node

record RedBlackTree {n m : Level } {t : Set m} (a k : Set n) : Set (m Level.@$\sqcup$@ n) where
  field
    root : Maybe (Node a k)
    nodeStack : SingleLinkedStack  (Node a k)
    compare : k @$\rightarrow$@ k @$\rightarrow$@ CompareResult {n}

open RedBlackTree


leafNode : {n : Level } {a k : Set n}  @$\rightarrow$@ k @$\rightarrow$@ a @$\rightarrow$@ Node a k
leafNode k1 value = record {
    key   = k1 ;
    value = value ;
    right = Nothing ;
    left  = Nothing ;
    color = Red
  }

putRedBlackTree : {n m : Level } {a k : Set n} {t : Set m} @$\rightarrow$@ RedBlackTree {n} {m} {t} a k @$\rightarrow$@ k @$\rightarrow$@ a @$\rightarrow$@ (RedBlackTree {n} {m} {t} a k @$\rightarrow$@ t) @$\rightarrow$@ t
putRedBlackTree {n} {m} {a} {k}  {t} tree k1 value next with (root tree)
...                                | Nothing = next (record tree {root = Just (leafNode k1 value) })
...                                | Just n2  = clearSingleLinkedStack (nodeStack tree) (\ s @$\rightarrow$@ findNode tree s (leafNode k1 value) n2 (\ tree1 s n1 @$\rightarrow$@ insertNode tree1 s n1 next))

getRedBlackTree : {n m : Level } {a k : Set n} {t : Set m} @$\rightarrow$@ RedBlackTree {n} {m} {t} a k @$\rightarrow$@ k @$\rightarrow$@ (RedBlackTree {n} {m} {t} a k @$\rightarrow$@ (Maybe (Node a k)) @$\rightarrow$@ t) @$\rightarrow$@ t
getRedBlackTree {_} {_} {a} {k} {t} tree k1 cs = checkNode (root tree)
  module GetRedBlackTree where 
    search : Node a k @$\rightarrow$@ t
    checkNode : Maybe (Node a k) @$\rightarrow$@ t
    checkNode Nothing = cs tree Nothing
    checkNode (Just n) = search n
    search n with compare tree k1 (key n) 
    search n | LT = checkNode (left n)
    search n | GT = checkNode (right n)
    search n | EQ = cs tree (Just n)