module bt 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 bst : Set where
  bt-empty : bst
  bt-node  : (node : tree (node col ℕ) bst bst ) → bst

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

-- 右回転
rotate_right : bst → bst
rotate_right bt-empty = bt-empty
rotate_right node@(bt-node (record { key = _; ltree = bt-empty; rtree = _ }) ) = node
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 }) )
  = (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 }) ) }) )

-- 左回転
rotate_left : bst → bst
rotate_left bt-empty = bt-empty
rotate_left node@(bt-node (record { key = _; ltree = _; rtree = bt-empty }) ) = node
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 }) ) }) )
  = (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 }) )

-- skew
skew : bst → bst
skew bt-empty = bt-empty
skew node@(bt-node (record { key = _; ltree = _ ; rtree = bt-empty }) ) = node
skew node@(bt-node (record { key = record { coler = y ; number = x }
  ; ltree = l
  ; rtree = (bt-node (record { key = record { coler = red ; number = lx }; ltree = ll; rtree = lr }) )}) )
  = rotate_left node
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 }) ) }) )
  = node

split : bst → bst
split bt-empty = bt-empty
split node@(bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = bt-empty }) ) = node
split node@(bt-node (record { key = record { coler = y ; number = x }; ltree = bt-empty; rtree = r }) ) = 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 }) )}) )
  = 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 }) )})

-- split
split_branch : bst → bst
split_branch bt-empty = bt-empty
split_branch node@(bt-node (record { key = _; ltree = bt-empty; rtree = _}) ) = 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 }) )
    = split (rotate_right 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 = black ; number = rrx }; ltree = lll; rtree = rrr }))
    ; rtree = lr  }) )
  ; rtree = r }) )
    = 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-empty
    ; rtree = rr }) )
  ; rtree = r}) )
    = node

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 }) )
  = node

insert_pipeline : bst → bst
insert_pipeline node = split_branch (skew node)

-- insert
bt-insert : bst → ℕ → bst
bt-insert bt-empty n = ( bt-node (record { key = record { coler = red ; number = n }; ltree = bt-empty; rtree = bt-empty }) )
bt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n  with <-cmp x n
bt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n | tri≈ ¬a b ¬c = (bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = r }) )
bt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n | tri> ¬a ¬b c = (insert_pipeline (bt-node (record { key = record { coler = y ; number = x }; ltree = (bt-insert l n); rtree = r }) ) )
bt-insert (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n | tri< a ¬b ¬c = (insert_pipeline (bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = (bt-insert r n) }) ) )

init_node_coler : bst → bst
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 : bst → ℕ → bst
insert node x = init_node_coler (bt-insert node x)

-- 赤黒木のinsertをする


rbt-insert2 : bst → ℕ → List bst → bst
rbt-insert2 bt-empty n stack = ( bt-node (record { key = record { coler = red ; number = n }; ltree = bt-empty; rtree = bt-empty }) )
rbt-insert2 (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack with <-cmp x n
rbt-insert2 (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-insert2 (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack | tri> ¬a ¬b c = (insert_pipeline (bt-node (record { key = record { coler = y ; number = x }; ltree = (rbt-insert2 l n stack); rtree = r }) ) )
rbt-insert2 (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) n stack | tri< a ¬b ¬c = (insert_pipeline (bt-node (record { key = record { coler = y ; number = x }; ltree = l; rtree = (rbt-insert2 r n stack) }) ) )

--treeをmerge
merge-tree : bst → List bst → bst
merge-tree node []  = init_node_coler (insert_pipeline node)
merge-tree x (bt-empty ∷ xs) =  bt-empty
merge-tree bt-empty s = bt-empty
merge-tree target@(bt-node (record { key = record { coler = ty; number = tx }; ltree = tl; rtree = tr }) ) ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) with <-cmp tx x
merge-tree target@(bt-node (record { key = record { coler = ty; number = tx }; ltree = tl; rtree = tr }) ) ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri≈ ¬a b ¬c
  = merge-tree ( insert_pipeline  (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r } ) ) ) xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r } ) ) xs

merge-tree target@(bt-node (record { key = record { coler = ty; number = tx }; ltree = tl; rtree = tr }) ) ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri> ¬a ¬b c
  = merge-tree ( insert_pipeline (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = target } ) ) ) xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = target } ) ) xs

merge-tree target@(bt-node (record { key = record { coler = ty; number = tx }; ltree = tl; rtree = tr }) ) ( (bt-node (record { key = record { coler = y; number = x }; ltree = l; rtree = r }) ) ∷ xs) | tri< a ¬b ¬c
  =  merge-tree ( insert_pipeline  (bt-node (record { key = record { coler = y; number = x }; ltree = target; rtree = r } ) ) ) xs
  -- merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = target; rtree = r } ) ) xs

--(bt-node (record { key = record { coler = red ; number = 1 }; ltree = bt-empty; rtree = bt-empty }))
-- merge-tree node ( (bt-node (record { key = record { coler = y; number = x }; ltree = bt-empty; rtree = bt-node r }) ) ∷ xs) = merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = node ; rtree = bt-node r }) ) xs
-- merge-tree node ( (bt-node (record { key = record { coler = y; number = x }; ltree = bt-node l; rtree = bt-empty }) ) ∷ xs) = merge-tree (bt-node (record { key = record { coler = y; number = x }; ltree = bt-node l ; rtree = node }) ) xs
--merge-tree node ( (bt-node (record { key = record { coler = y; number = x }; ltree = bt-node _; rtree = bt-node _ }) ) ∷ xs) = bt-empty


rbt-insert : bst → ℕ → List bst → bst
rbt-insert bt-empty n stack = merge-tree (bt-node (record { key = record { coler = red ; number = n }; ltree = bt-empty; rtree = bt-empty }) ) stack
-- ( bt-node (record { key = record { coler = red ; number = 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-insert1 : bst
test-insert1 = insert bt-empty 0

test-insert2 : bst
test-insert2 = insert test-insert1 1

test-insert3 : bst
test-insert3 = insert test-insert2 2

test-insert4 : bst
test-insert4 = insert test-insert3 3

test-insert5 : bst
test-insert5 = insert test-insert4 4

test-insert6 : bst
test-insert6 = insert test-insert5 5

test-insert7 : bst
test-insert7 = insert test-insert6 6