module GearsRBTree where

open import Data.Nat hiding (compare)
open import Data.Nat.Properties as NatProp -- <-cmp
open import Relation.Binary
open import Data.List

data col : Set where
  black : col
  red : col

record node (A B : Set) : Set where
  field
    coler : A
    number : B

record tree (A B C : Set) : Set where
  field
    key : A
    ltree : B
    rtree : C

data rbt : Set where
  bt-empty : rbt
  bt-node  : (node : tree (node col ℕ) rbt rbt ) → rbt

test : rbt
test = bt-node (record { key = record { coler = red ; number = zero }; ltree = bt-empty; rtree = bt-empty })

init_node_coler : rbt → rbt
init_node_coler bt-empty = bt-empty
init_node_coler (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) = (bt-node (record { key = record { coler = black; number = x }; ltree = l; rtree = r }) )

-- 赤黒木のinsertをする

-- skewの実行後にmergeに遷移するので型定義のみ上で行う
merge-tree : rbt → ℕ → List rbt → rbt

split : rbt → ℕ → List rbt → rbt
split bt-empty n stack = bt-empty
split node@(bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = bt-empty }) ) n stack = node
split node@(bt-node (record { key = record { coler = y ; number = x }; ltree = bt-empty; rtree = r }) ) n stack = node
split node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = ly ; number = lx }; ltree = ll; rtree = lr }) )
  ; rtree = (bt-node (record { key = record { coler = ry ; number = rx }; ltree = rl; rtree = rr }) )}) ) n stack
  = merge-tree (bt-node (record { key = record { coler = red ; number = x }
      ; ltree = (bt-node (record { key = record { coler = black ; number = lx }; ltree = ll; rtree = lr }) )
      ; rtree = (bt-node (record { key = record { coler = black ; number = rx }; ltree = rl; rtree = rr }) )})) n stack


-- 右回転
merge-rotate_right : rbt → ℕ → List rbt → rbt
merge-rotate_right bt-empty n stack = bt-empty
merge-rotate_right node@(bt-node (record { key = _; ltree = bt-empty; rtree = _ }) ) n stack = node
merge-rotate_right (bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = ly ; number = lx }; ltree = ll; rtree = lr }) )
  ; rtree = r }) ) n stack
    = split (bt-node (record { key = record { coler = y ; number = lx }
      ; ltree = ll
      ; rtree = (bt-node (record { key = record { coler = red ; number = x }; ltree = lr; rtree = r }) ) }) ) n stack


-- split
split_branch : rbt → ℕ → List rbt → rbt
split_branch bt-empty n stack = bt-empty
split_branch node@(bt-node (record { key = _; ltree = bt-empty; rtree = _}) ) n stack = node
split_branch node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = red ; number = rx }
    ; ltree = (bt-node (record { key = record { coler = red ; number = rrx }; ltree = rrl; rtree = rrr }) )
    ; rtree = lr }))
  ; rtree = r }) ) n stack
    = merge-rotate_right node n stack

split_branch node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = red ; number = rx }
    ; ltree = (bt-node (record { key = record { coler = black ; number = rrx }; ltree = lll; rtree = rrr }))
    ; rtree = lr  }) )
  ; rtree = r }) ) n stack
    = merge-tree node n stack

split_branch node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = red ; number = rx }
    ; ltree =  bt-empty
    ; rtree = rr }) )
  ; rtree = r}) ) n stack
    = merge-tree node n stack

split_branch node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = (bt-node (record { key = record { coler = black ; number = rx }; ltree = rl; rtree = rr }) )
  ; rtree = r }) ) n stack
    = merge-tree node n stack


merge-rotate_left : rbt → ℕ → List rbt → rbt
merge-rotate_left bt-empty n stack = bt-empty
merge-rotate_left node@(bt-node (record { key = _; ltree = _; rtree = bt-empty }) ) n stack = node
merge-rotate_left (bt-node (record { key = record { coler = y ; number = x }
  ; ltree = l
  ; rtree = (bt-node (record { key = record { coler = ry ; number = rx }; ltree = rl; rtree = rr }) ) }) ) n stack
  = split_branch (bt-node (record { key = record { coler = y ; number = rx }
    ; ltree = (bt-node (record { key = record { coler = red ; number = x }; ltree = l; rtree = rl }) )
    ; rtree = rr }) ) n stack

-- skew
skew : rbt → ℕ → List rbt → rbt
skew bt-empty n stack = bt-empty
skew node@(bt-node (record { key = _; ltree = _ ; rtree = bt-empty }) ) n stack = node
skew node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = l
  ; rtree = r@(bt-node (record { key = record { coler = red ; number = lx }; ltree = ll; rtree = lr }) )}) ) n stack
    = merge-rotate_left node n stack

skew node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = l
  ; rtree = (bt-node (record { key = record { coler = black ; number = lx }; ltree = ll; rtree = lr }) ) }) ) n stack
    = split_branch node n stack


--treeをmerge

merge-tree node n []  = init_node_coler node
merge-tree node n (bt-empty ∷ xs) =  bt-empty
merge-tree target n ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) with <-cmp n x
merge-tree target n ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri≈ ¬a b ¬c
  = skew (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r } ) ) x xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r } ) ) xs

merge-tree target n ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri> ¬a ¬b c
  = skew (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = target } ) ) x xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = target } ) ) xs

merge-tree target n ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri< a ¬b ¬c
  = skew (bt-node (record { key = record { coler = y; number = x }; ltree = target; rtree = r } ) ) x xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = target; rtree = r } ) ) xs


rbt-insert : rbt → ℕ → List rbt → rbt
rbt-insert bt-empty n stack = merge-tree (bt-node (record { key = record { coler = red ; number = n }; ltree = bt-empty; rtree = bt-empty }) ) n stack
-- ( bt-node (record { key = record { coler = red ; nu mber = n }; ltree = bt-empty; rtree = bt-empty }) ) -- mergeに遷移する
rbt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack with <-cmp x n
rbt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack | tri≈ ¬a b ¬c
  = (bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = r }) )
rbt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack | tri> ¬a ¬b c
  = rbt-insert l n ( (bt-node (record { key = record { coler = y ; number = x }; ltree = bt-empty; rtree = r }) ) ∷ stack)
rbt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack | tri< a ¬b ¬c
  = rbt-insert r n ( (bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = bt-empty }) ) ∷ stack )


-- 以下test部分

test-rbt1 = rbt-insert bt-empty 0 []
test-rbt2 = rbt-insert test-rbt1 1 []
test-rbt3 = rbt-insert test-rbt2 2 []
test-rbt4 = rbt-insert test-rbt3 3 []
test-rbt5 = rbt-insert test-rbt4 4 []
test-rbt6 = rbt-insert test-rbt5 5 []
test-rbt7 = rbt-insert test-rbt6 6 []