module RedBlackTreeHoare where


open import Level hiding (zero)

open import Data.Nat hiding (compare)
open import Data.Nat.Properties as NatProp
open import Data.Maybe
open import Data.Bool
open import Data.Empty

open import Relation.Binary
open import Relation.Binary.PropositionalEquality

open import stack

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 {n} {m} {a} {t} treeImpl
  putTree : a → (Tree treeImpl → t) → t
  putTree d next = putImpl (treeMethods ) tree d (\t1 → next (record {tree = t1 ; treeMethods = treeMethods} ))
  getTree : (Tree treeImpl → Maybe a → t) → t
  getTree next = getImpl (treeMethods ) tree (\t1 d → next (record {tree = t1 ; treeMethods = treeMethods} ) d )

open Tree

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


record Node {n : Level } (a : Set n) (k : ℕ) : Set n where
  inductive
  field
    key   : ℕ
    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 : Set n) (k : ℕ) : Set (m Level.⊔ n) where
  field
    root : Maybe (Node a k)
    nodeStack : SingleLinkedStack  (Node a k)
    -- compare : k → k → Tri A B C
    -- <-cmp 

open RedBlackTree

open SingleLinkedStack

compTri : ( x y : ℕ ) ->  Tri ( x < y )  ( x ≡ y )  ( x > y )
compTri = IsStrictTotalOrder.compare (Relation.Binary.StrictTotalOrder.isStrictTotalOrder <-strictTotalOrder)
  where open import  Relation.Binary

-- 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 : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} a k → SingleLinkedStack (Node a k) →  Node a k → (RedBlackTree {n} {m} {t} a k → t) → t
replaceNode {n} {m} {t} {a} {k} tree s n0 next = popSingleLinkedStack s (
      \s parent → replaceNode1 s parent)
       module ReplaceNode where
          replaceNode1 : SingleLinkedStack (Node a k) → Maybe ( Node a k ) → t
          replaceNode1 s nothing = next ( record tree { root = just (record n0 { color = Black}) } )
          replaceNode1 s (just n1) with compTri  (key n1) (key n0)
          replaceNode1 s (just n1) | tri< lt ¬eq ¬gt = replaceNode {n} {m} {t} {a} {k} tree s ( record n1 { value = value n0 ; left = left n0 ; right = right n0 } ) next
          replaceNode1 s (just n1) | tri≈ ¬lt eq ¬gt = replaceNode {n} {m} {t} {a} {k} tree s ( record n1 { left = just n0 } ) next
          replaceNode1 s (just n1) | tri> ¬lt ¬eq gt = replaceNode {n} {m} {t} {a} {k} tree s ( record n1 { right = just n0 } ) next


rotateRight : {n m : Level } {t : Set m } {a : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) →
  (RedBlackTree {n} {m} {t} a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) → t) → t
rotateRight {n} {m} {t} {a} {k} tree s n0 parent rotateNext = getSingleLinkedStack s (\ s n0 → rotateRight1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext)
  where
        rotateRight1 : {n m : Level } {t : Set m } {a : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} a k → SingleLinkedStack (Node  a k)  → Maybe (Node a k) → Maybe (Node a k) →
          (RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node  a k)  → Maybe (Node a k) → Maybe (Node a k) → t) → 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 compTri (key n1) (key leftParent)
        rotateRight1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just leftParent | tri< a₁ ¬b ¬c = rotateNext tree s (just n1) parent
        rotateRight1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just leftParent | tri≈ ¬a b ¬c = rotateNext tree s (just n1) parent
        rotateRight1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just leftParent | tri> ¬a ¬b c = rotateNext tree s (just n1) parent


rotateLeft : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) →
  (RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) →  t) → t
rotateLeft {n} {m} {t} {a} {k}  tree s n0 parent rotateNext = getSingleLinkedStack s (\ s n0 → rotateLeft1 tree s n0 parent rotateNext)
  where
        rotateLeft1 : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) →
          (RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node  a k) → Maybe (Node a k) → Maybe (Node a k) → t) → 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 compTri (key n1) (key rightParent)
        rotateLeft1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just rightParent | tri< a₁ ¬b ¬c = rotateNext tree s (just n1) parent
        rotateLeft1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just rightParent | tri≈ ¬a b ¬c = rotateNext tree s (just n1) parent
        rotateLeft1 {n} {m} {t} {a} {k} tree s n0 parent rotateNext | just n1 | just parent1 | just rightParent | tri> ¬a ¬b c = rotateNext tree s (just n1) parent
        -- ...                                    | EQ = rotateNext tree s (just n1) parent
        -- ...                                    | _ = rotateNext tree s (just n1) parent

{-# TERMINATING #-}
insertCase5 : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Maybe (Node a k) → Node a k → Node a k → (RedBlackTree {n} {m} {t}  a k → t) → t
insertCase5 {n} {m} {t} {a} {k}  tree s n0 parent grandParent next = pop2SingleLinkedStack s (\ s parent grandParent → insertCase51 tree s n0 parent grandParent next)
  where
    insertCase51 : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Maybe (Node a k) → Maybe (Node a k) → Maybe (Node a k) → (RedBlackTree {n} {m} {t}  a k → t) → 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 compTri (key n1) (key leftParent1) | compTri (key leftParent1) (key leftGrandParent1)
    ...    | tri≈ ¬a b ¬c | tri≈ ¬a1 b1 ¬c1  = rotateRight tree s n0 parent (\ tree s n0 parent → insertCase5 tree s n0 parent1 grandParent1 next)
    ...    | _            | _                = rotateLeft tree s n0 parent (\ tree s n0 parent → insertCase5 tree s n0 parent1 grandParent1 next)
    -- ...     | EQ | EQ = rotateRight tree s n0 parent (\ tree s n0 parent → insertCase5 tree s n0 parent1 grandParent1 next)
    -- ...     | _ | _ = rotateLeft tree s n0 parent (\ tree s n0 parent → insertCase5 tree s n0 parent1 grandParent1 next)

insertCase4 : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Node a k → Node a k → Node a k → (RedBlackTree {n} {m} {t}  a k → t) → 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 compTri (key n0) (key rightParent) | compTri (key parent) (key leftGrandParent) -- (key n0) (key rightParent) | (key parent) (key leftGrandParent)
-- ...                                              | EQ | EQ = popSingleLinkedStack s (\ s n1 → rotateLeft tree s (left n0) (just grandParent)
--    (\ tree s n0 parent → insertCase5 tree s n0 rightParent grandParent next))
-- ...                                              | _ | _  = insertCase41 tree s n0 parent grandParent next
...                                                 | tri≈ ¬a b ¬c | tri≈ ¬a1 b1 ¬c1 = popSingleLinkedStack s (\ s n1 → rotateLeft tree s (left n0) (just grandParent) (\ tree s n0 parent → insertCase5 tree s n0 rightParent grandParent next))
... | _            | _               = insertCase41 tree s n0 parent grandParent next
  where
    insertCase41 : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Node a k → Node a k → Node a k → (RedBlackTree {n} {m} {t}  a k → t) → 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 compTri (key n0) (key leftParent) | compTri (key parent) (key rightGrandParent)
    ... | tri≈ ¬a b ¬c | tri≈ ¬a1 b1 ¬c1 =  popSingleLinkedStack s (\ s n1 → rotateRight tree s (right n0) (just grandParent) (\ tree s n0 parent → insertCase5 tree s n0 leftParent grandParent next))
    ... | _ | _ = insertCase5 tree s (just n0) parent grandParent next
    -- ...                                              | EQ | EQ = popSingleLinkedStack s (\ s n1 → rotateRight tree s (right n0) (just grandParent)
    --    (\ tree s n0 parent → insertCase5 tree s n0 leftParent grandParent next))
    -- ...                                              | _ | _  = insertCase5 tree s (just n0) parent grandParent next

colorNode : {n : Level } {a : Set n} {k : ℕ} → Node a k → Color  → Node a k
colorNode old c = record old { color = c }

{-# TERMINATING #-}
insertNode : {n m  : Level } {t : Set m } {a : Set n} {k : ℕ}  → RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Node a k → (RedBlackTree {n} {m} {t}  a k → t) → t
insertNode {n} {m} {t} {a} {k}  tree s n0 next = get2SingleLinkedStack s (insertCase1 n0)
   where
    insertCase1 : Node a k → SingleLinkedStack (Node a k) → Maybe (Node a k) → Maybe (Node a k) → t    -- placed here to allow mutual recursion
          -- http://agda.readthedocs.io/en/v2.5.2/language/mutual-recursion.html
    insertCase3 : SingleLinkedStack (Node a k) → Node a k → Node a k → Node a k → 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 compTri ( key uncle ) ( key parent )
    insertCase3 s n0 parent grandParent | just uncle | _ | tri≈ ¬a b ¬c = insertCase4 tree s n0 parent grandParent next
    insertCase3 s n0 parent grandParent | just uncle | _ | tri< a ¬b ¬c with color uncle
    insertCase3 s n0 parent grandParent | just uncle | _ | tri< a ¬b ¬c | Red = pop2SingleLinkedStack s ( \s p0 p1 → insertCase1  (
           record grandParent { color = Red ; left = just ( record parent { color = Black } )  ; right = just ( record uncle { color = Black } ) }) s p0 p1 )
    insertCase3 s n0 parent grandParent | just uncle | _ | tri< a ¬b ¬c | Black = insertCase4 tree s n0 parent grandParent next
    insertCase3 s n0 parent grandParent | just uncle | _ | tri> ¬a ¬b c with color uncle
    insertCase3 s n0 parent grandParent | just uncle | _ | tri> ¬a ¬b c | Red = pop2SingleLinkedStack s ( \s p0 p1 → insertCase1  ( record grandParent { color = Red ; left = just ( record parent { color = Black } )  ; right = just ( record uncle { color = Black } ) }) s p0 p1 )
    insertCase3 s n0 parent grandParent | just uncle | _ | tri> ¬a ¬b c | Black = insertCase4 tree s n0 parent grandParent next
    -- ...                   | EQ =  insertCase4 tree s n0 parent grandParent next
    -- ...                   | _ with color uncle
    -- ...                           | Red = pop2SingleLinkedStack s ( \s p0 p1 → 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) → Node a k → Node a k → Node a k → 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 : Set n} {k : ℕ} {t : Set m}  → RedBlackTree {n} {m} {t}   a k → SingleLinkedStack (Node a k) → (Node a k) → (Node a k) → (RedBlackTree {n} {m} {t}  a k → SingleLinkedStack (Node a k) → Node a k → t) → t
findNode {n} {m} {a} {k} {t} tree s n0 n1 next = pushSingleLinkedStack s n1 (\ s → findNode1 s n1)
  module FindNode where
    findNode2 : SingleLinkedStack (Node a k) → (Maybe (Node a k)) → t
    findNode2 s nothing = next tree s n0
    findNode2 s (just n) = findNode tree s n0 n next
    findNode1 : SingleLinkedStack (Node a k) → (Node a k)  → t
    findNode1 s n1 with (compTri (key n0) (key n1))
    findNode1 s n1 | tri< a ¬b ¬c = popSingleLinkedStack s ( \s _ → next tree s (record n1 { key = key n1 ; value = value n0 } ) )
    findNode1 s n1 | tri≈ ¬a b ¬c = findNode2 s (right n1)
    findNode1 s n1 | tri> ¬a ¬b c = findNode2 s (left n1)
    -- ...                                | EQ = popSingleLinkedStack s ( \s _ → 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 : Set n } → a → (k : ℕ) → (Node a k)
leafNode v k1 = record { key = k1 ; value = v ; right = nothing ; left = nothing ; color = Red }

putRedBlackTree : {n m : Level} {t : Set m} {a : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} a k → {!!} → {!!} → (RedBlackTree {n} {m} {t} a k → t) → t
putRedBlackTree {n} {m} {t} {a} {k} tree val k1 next with (root tree)
putRedBlackTree {n} {m} {t} {a} {k} tree val k1 next | nothing = next (record tree {root = just (leafNode {!!} {!!}) })
putRedBlackTree {n} {m} {t} {a} {k} tree val k1 next | just n2 = clearSingleLinkedStack (nodeStack tree) (λ s → findNode tree s (leafNode {!!} {!!}) n2 (λ tree1 s n1 → insertNode tree1 s n1 next))
-- putRedBlackTree {n} {m} {t} {a} {k} tree value k1 next with (root tree)
-- ...                                | nothing = next (record tree {root = just (leafNode k1 value) })
-- ...                                | just n2  = clearSingleLinkedStack (nodeStack tree) (\ s → findNode tree s (leafNode k1 value) n2 (\ tree1 s n1 → insertNode tree1 s n1 next))


-- getRedBlackTree : {n m  : Level } {t : Set m}  {a : Set n} {k : ℕ} → RedBlackTree {n} {m} {t} {A}  a k → k → (RedBlackTree {n} {m} {t} {A}  a k → (Maybe (Node a k)) → t) → t
-- getRedBlackTree {_} {_} {t}  {a} {k} 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 → t
--     checkNode : Maybe (Node a k) → t
--     checkNode nothing = cs tree nothing
--     checkNode (just n) = search n
--     search n with compTri k1 (key n)
--     search n | tri< a ¬b ¬c = checkNode (left n)
--     search n | tri≈ ¬a b ¬c = cs tree (just n)
--     search n | tri> ¬a ¬b c = checkNode (right n)



-- compareT :  {A B C : Set } → ℕ → ℕ → Tri A B C
-- compareT x y with IsStrictTotalOrder.compare (Relation.Binary.StrictTotalOrder.isStrictTotalOrder <-strictTotalOrder) x y
-- compareT x y | tri< a ¬b ¬c = tri< {!!} {!!} {!!}
-- compareT x y | tri≈ ¬a b ¬c = {!!}
-- compareT x y | tri> ¬a ¬b c = {!!}
-- -- ... | tri≈ a b c = {!!}
-- -- ... | tri< a b c = {!!}
-- -- ... | tri> a b c = {!!}

-- compare2 : (x y : ℕ ) → CompareResult {Level.zero}
-- compare2 zero zero = EQ
-- compare2 (suc _) zero = GT
-- compare2  zero (suc _) = LT
-- compare2  (suc x) (suc y) = compare2 x y

-- -- putUnblanceTree : {n m : Level } {a : Set n} {k : ℕ} {t : Set m} → RedBlackTree {n} {m} {t} {A}  a k → k → a → (RedBlackTree {n} {m} {t} {A}  a k → t) → t
-- -- putUnblanceTree {n} {m} {A} {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 → findNode tree s (leafNode k1 value) n2 (λ  tree1 s n1 → replaceNode tree1 s n1 next))

-- -- checkT : {m : Level } (n : Maybe (Node  ℕ ℕ)) → ℕ → Bool
-- -- checkT nothing _ = false
-- -- checkT (just n) x with compTri (value n)  x
-- -- ...  | tri≈ _ _ _ = true
-- -- ...  | _ = false

-- -- checkEQ :  {m : Level }  ( x :  ℕ ) -> ( n : Node  ℕ ℕ ) -> (value n )  ≡ x  -> checkT {m} (just n) x ≡ true
-- -- checkEQ x n refl with compTri (value n)  x
-- -- ... |  tri≈ _ refl _ = refl
-- -- ... |  tri> _ neq gt =  ⊥-elim (neq refl)
-- -- ... |  tri< lt neq _ =  ⊥-elim (neq refl)


createEmptyRedBlackTreeℕ : {n m  : Level} {t : Set m} (a : Set n) (b : ℕ)
     → RedBlackTree {n} {m} {t} a b
createEmptyRedBlackTreeℕ a b = record {
        root = nothing
     ;  nodeStack = emptySingleLinkedStack
     -- ;  nodeComp = λ x x₁ → {!!}

   }

-- ( x y : ℕ ) ->  Tri  ( x < y )  ( x ≡ y )  ( x > y )

-- test = (λ x → (createEmptyRedBlackTreeℕ x x) 

-- ts = createEmptyRedBlackTreeℕ {ℕ} {?} {!!} 0

-- tes = putRedBlackTree {_} {_} {_} (createEmptyRedBlackTreeℕ {_} {_} {_} 3 3) 2 2 (λ t → t)