module RedBlackTree where

open import stack
open import Level hiding (zero)
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

open SingleLinkedStack

--
-- put new node at parent node, and rebuild tree to the top
--
{-!$\#$! TERMINATING !$\#$!-}   -- https://agda.readthedocs.io/en/v2.5.3/language/termination-checking.html
replaceNode : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$!  Node a k !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
replaceNode {n} {m} {t} {a} {k} tree s n0 next = popSingleLinkedStack s (
      \s parent !$\rightarrow$! replaceNode1 s parent)
        where
          replaceNode1 : SingleLinkedStack (Node a k) !$\rightarrow$! Maybe ( Node a k ) !$\rightarrow$! t 
          replaceNode1 s Nothing = next ( record tree { root = Just (record n0 { color = Black}) } )
          replaceNode1 s (Just n1) with compare tree (key n1) (key n0)
          ... | EQ =  replaceNode tree s ( record n1 { value = value n0 ; left = left n0 ; right = right n0 } ) next
          ... | GT =  replaceNode tree s ( record n1 { left = Just n0 } ) next
          ... | LT =  replaceNode tree s ( record n1 { right = Just n0 } ) next


rotateRight : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$!
  (RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! t) !$\rightarrow$! t
rotateRight {n} {m} {t} {a} {k}  tree s n0 parent rotateNext = getSingleLinkedStack s (\ s n0 !$\rightarrow$! rotateRight1 tree s n0 parent rotateNext)
  where
        rotateRight1 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k)  !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! 
          (RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k)  !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! t) !$\rightarrow$! t
        rotateRight1 {n} {m} {t} {a} {k}  tree s n0 parent rotateNext with n0
        ... | Nothing  = rotateNext tree s Nothing n0 
        ... | Just n1 with parent
        ...           | Nothing = rotateNext tree s (Just n1 ) n0
        ...           | Just parent1 with left parent1
        ...                | Nothing = rotateNext tree s (Just n1) Nothing 
        ...                | Just leftParent with compare tree (key n1) (key leftParent)
        ...                                    | EQ = rotateNext tree s (Just n1) parent 
        ...                                    | _ = rotateNext tree s (Just n1) parent 


rotateLeft : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$!
  (RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$!  t) !$\rightarrow$! t
rotateLeft {n} {m} {t} {a} {k}  tree s n0 parent rotateNext = getSingleLinkedStack s (\ s n0 !$\rightarrow$! rotateLeft1 tree s n0 parent rotateNext)
  where
        rotateLeft1 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! 
          (RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node  a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! t) !$\rightarrow$! t
        rotateLeft1 {n} {m} {t} {a} {k}  tree s n0 parent rotateNext with n0
        ... | Nothing  = rotateNext tree s Nothing n0 
        ... | Just n1 with parent
        ...           | Nothing = rotateNext tree s (Just n1) Nothing 
        ...           | Just parent1 with right parent1
        ...                | Nothing = rotateNext tree s (Just n1) Nothing 
        ...                | Just rightParent with compare tree (key n1) (key rightParent)
        ...                                    | EQ = rotateNext tree s (Just n1) parent 
        ...                                    | _ = rotateNext tree s (Just n1) parent 

{-!$\#$! TERMINATING !$\#$!-}
insertCase5 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
insertCase5 {n} {m} {t} {a} {k}  tree s n0 parent grandParent next = pop2SingleLinkedStack s (\ s parent grandParent !$\rightarrow$! insertCase51 tree s n0 parent grandParent next)
  where
    insertCase51 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
    insertCase51 {n} {m} {t} {a} {k}  tree s n0 parent grandParent next with n0
    ...     | Nothing = next tree
    ...     | Just n1  with  parent | grandParent
    ...                 | Nothing | _  = next tree
    ...                 | _ | Nothing  = next tree
    ...                 | Just parent1 | Just grandParent1 with left parent1 | left grandParent1
    ...                                                     | Nothing | _  = next tree
    ...                                                     | _ | Nothing  = next tree
    ...                                                     | Just leftParent1 | Just leftGrandParent1
      with compare tree (key n1) (key leftParent1) | compare tree (key leftParent1) (key leftGrandParent1)
    ...     | EQ | EQ = rotateRight tree s n0 parent 
                 (\ tree s n0 parent !$\rightarrow$! insertCase5 tree s n0 parent1 grandParent1 next)
    ...     | _ | _ = rotateLeft tree s n0 parent 
                 (\ tree s n0 parent !$\rightarrow$! insertCase5 tree s n0 parent1 grandParent1 next)

insertCase4 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
insertCase4 {n} {m} {t} {a} {k}  tree s n0 parent grandParent next
       with  (right parent) | (left grandParent)
...    | Nothing | _ = insertCase5 tree s (Just n0) parent grandParent next
...    | _ | Nothing = insertCase5 tree s (Just n0) parent grandParent next       
...    | Just rightParent | Just leftGrandParent with compare tree (key n0) (key rightParent) | compare tree (key parent) (key leftGrandParent)
...                                              | EQ | EQ = popSingleLinkedStack s (\ s n1 !$\rightarrow$! rotateLeft tree s (left n0) (Just grandParent)
   (\ tree s n0 parent !$\rightarrow$! insertCase5 tree s n0 rightParent grandParent next))
...                                              | _ | _  = insertCase41 tree s n0 parent grandParent next
  where
    insertCase41 : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
    insertCase41 {n} {m} {t} {a} {k}  tree s n0 parent grandParent next
                 with  (left parent) | (right grandParent)       
    ...    | Nothing | _ = insertCase5 tree s (Just n0) parent grandParent next
    ...    | _ | Nothing = insertCase5 tree s (Just n0) parent grandParent next
    ...    | Just leftParent | Just rightGrandParent with compare tree (key n0) (key leftParent) | compare tree (key parent) (key rightGrandParent)
    ...                                              | EQ | EQ = popSingleLinkedStack s (\ s n1 !$\rightarrow$! rotateRight tree s (right n0) (Just grandParent)
       (\ tree s n0 parent !$\rightarrow$! insertCase5 tree s n0 leftParent grandParent next))
    ...                                              | _ | _  = insertCase5 tree s (Just n0) parent grandParent next

colorNode : {n : Level } {a k : Set n}  !$\rightarrow$! Node a k !$\rightarrow$! Color  !$\rightarrow$! Node a k
colorNode old c = record old { color = c }

{-!$\#$! TERMINATING !$\#$!-}
insertNode : {n m : Level } {t : Set m } {a k : Set n} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! t) !$\rightarrow$! t
insertNode {n} {m} {t} {a} {k}  tree s n0 next = get2SingleLinkedStack s (insertCase1 n0)
   where
    insertCase1 : Node a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! Maybe (Node a k) !$\rightarrow$! t    -- placed here to allow mutual recursion
          -- http://agda.readthedocs.io/en/v2.5.2/language/mutual-recursion.html
    insertCase3 : SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! t
    insertCase3 s n0 parent grandParent with left grandParent | right grandParent
    ... | Nothing | Nothing = insertCase4 tree s n0 parent grandParent next
    ... | Nothing | Just uncle  = insertCase4 tree s n0 parent grandParent next
    ... | Just uncle | _  with compare tree ( key uncle ) ( key parent )
    ...                   | EQ =  insertCase4 tree s n0 parent grandParent next
    ...                   | _ with color uncle
    ...                           | Red = pop2SingleLinkedStack s ( \s p0 p1 !$\rightarrow$! insertCase1  (
           record grandParent { color = Red ; left = Just ( record parent { color = Black } )  ; right = Just ( record uncle { color = Black } ) }) s p0 p1 )
    ...                           | Black = insertCase4 tree s n0 parent grandParent next
    insertCase2 : SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! Node a k !$\rightarrow$! t
    insertCase2 s n0 parent grandParent with color parent
    ... | Black = replaceNode tree s n0 next
    ... | Red =   insertCase3 s n0 parent grandParent
    insertCase1 n0 s Nothing Nothing = next tree
    insertCase1 n0 s Nothing (Just grandParent) = next tree
    insertCase1 n0 s (Just parent) Nothing = replaceNode tree s (colorNode n0 Black) next
    insertCase1 n0 s (Just parent) (Just grandParent) = insertCase2 s n0 parent grandParent

----
-- find node potition to insert or to delete, the path will be in the stack
-- 
findNode : {n m : Level } {a k : Set n} {t : Set m} !$\rightarrow$! RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! (Node a k) !$\rightarrow$! (Node a k) !$\rightarrow$! (RedBlackTree {n} {m} {t} a k !$\rightarrow$! SingleLinkedStack (Node a k) !$\rightarrow$! Node a k !$\rightarrow$! t) !$\rightarrow$! t
findNode {n} {m} {a} {k}  {t} tree s n0 n1 next = pushSingleLinkedStack s n1 (\ s !$\rightarrow$! findNode1 s n1)
  where
    findNode2 : SingleLinkedStack (Node a k) !$\rightarrow$! (Maybe (Node a k)) !$\rightarrow$! t
    findNode2 s Nothing = next tree s n0
    findNode2 s (Just n) = findNode tree s n0 n next
    findNode1 : SingleLinkedStack (Node a k) !$\rightarrow$! (Node a k)  !$\rightarrow$! t
    findNode1 s n1 with (compare tree (key n0) (key n1))
    ...                                | EQ = popSingleLinkedStack s ( \s _ !$\rightarrow$! next tree s (record n1 { key = key n1 ; value = value n0 } ) )
    ...                                | GT = findNode2 s (right n1)
    ...                                | LT = findNode2 s (left n1)


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                     -- http://agda.readthedocs.io/en/v2.5.2/language/let-and-where.html
    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)

open import Data.Nat hiding (compare)

compare!$\mathbb{N}$! :  !$\mathbb{N}$! !$\rightarrow$! !$\mathbb{N}$! !$\rightarrow$! CompareResult {Level.zero}
compare!$\mathbb{N}$! x y with Data.Nat.compare x y
... | less _ _ = LT
... | equal _ = EQ
... | greater _ _ = GT

compare2 : (x y : !$\mathbb{N}$! ) !$\rightarrow$! CompareResult {Level.zero}
compare2 zero zero = EQ
compare2 (suc _) zero = GT
compare2  zero (suc _) = LT
compare2  (suc x) (suc y) = compare2 x y


createEmptyRedBlackTree!$\mathbb{N}$! : { m : Level } (a : Set Level.zero) {t : Set m} !$\rightarrow$! RedBlackTree {Level.zero} {m} {t} a !$\mathbb{N}$! 
createEmptyRedBlackTree!$\mathbb{N}$!  {m} a {t} = record {
        root = Nothing
     ;  nodeStack = emptySingleLinkedStack
     ;  compare = compare2
   }