module rbt_t where

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

-- @$\rightarrow$@ t を適用するために必要だった
open import Level renaming ( suc to succ ; zero to Zero )
open import Level


data col : Set where
  black : col
  red : col

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

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

data rbt : Set where
  bt-empty : rbt
  bt-node  : (node : tree (node col @$\mathbb{N}$@) rbt rbt ) @$\rightarrow$@ rbt

record Env : Set (succ Zero) where
  field
    vart : rbt
    varn : @$\mathbb{N}$@
    varl : List rbt
open Env


whileTest-next : {l : Level} {t : Set l}  @$\rightarrow$@ (c10 : rbt) @$\rightarrow$@ (n : @$\mathbb{N}$@) @$\rightarrow$@ (list : List rbt) @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
whileTest-next c10 n list next exit = next (record {vart = c10 ; varn = n ; varl = list} )

merge-tree :  {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t

-- 全てmerge-treeへ
split : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
split node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
split node@record { vart = (bt-node record { key = key ; ltree = bt-empty ; rtree = rtree }) ; varn = varn ; varl = varl } next exit = exit node
split node@record { vart = (bt-node record { key = key ; ltree = (bt-node node@$\_{1}$@) ; rtree = bt-empty }) ; varn = varn ; varl = varl } next exit = exit node
split record { vart = (bt-node record { key = record { coler = coler@$\_{1}$@ ; number = number@$\_{1}$@ }
  ; ltree = (bt-node record { key = record { coler = ly ; number = ln } ; ltree = ll ; rtree = lr })
  ; rtree = (bt-node record { key = record { coler = ry ; number = rn } ; ltree = rl ; rtree = rr }) })
  ; varn = varn ; varl = varl } next exit
  = next record { vart = (bt-node record { key = record { coler = red ; number = number@$\_{1}$@ }
    ; ltree = (bt-node record { key = record { coler = black ; number = ln } ; ltree = ll ; rtree = lr })
    ; rtree = (bt-node record { key = record { coler = black ; number = rn } ; ltree = rl ; rtree = rr }) })
    ; varn = varn ; varl = varl }


-- 右回転
-- 実行時splitへ、それ以外はmerge-treeへ
merge-rotate-right : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
merge-rotate-right node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-right node@record { vart = bt-node record { key = record { coler = coler ; number = number } ; ltree = bt-empty ; rtree = r } ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-right record { vart = 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 }
  ; varn = varn ; varl = varl } next exit
    = next record { vart = 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 }) }
      ; varn = varn ; varl = varl }


-- split
split-branch : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
split-branch node@record{ vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key ; ltree = bt-empty ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key@$\_{1}$@ ; ltree = (bt-node record { key = record { coler = lc ; number = number } ; ltree = bt-empty ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key@$\_{1}$@ ; ltree = (bt-node record { key = record { coler = black ; number = number } ; ltree = (bt-node node@$\_{1}$@) ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key@$\_{1}$@ ; ltree = (bt-node record { key = record { coler = red ; number = number@$\_{1}$@ } ; ltree = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree@$\_{1}$@ }) ; rtree = lr }) ; rtree = rtree } ; varn = varn ; varl = varl } next exit = exit node
split-branch node@record{ vart = bt-node record { key = key@$\_{1}$@
  ; ltree = (bt-node record { key = record { coler = red ; number = number@$\_{1}$@ }
    ; ltree = (bt-node record { key = record { coler = red ; number = number } ; ltree = ltree ; rtree = rtree@$\_{1}$@ })
    ; rtree = lr })
  ; rtree = rtree }
  ; varn = varn ; varl = varl } next exit
  = next node


-- 左回転、exitはsplit_branchへ　nextもsplit_branchへ
merge-rotate-left : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
merge-rotate-left node@record { vart = bt-empty ; varn = varn ; varl = varl }  next exit = exit node
merge-rotate-left node@record { vart = bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = bt-empty } ; varn = varn ; varl = varl } next exit = exit node
merge-rotate-left record { vart = 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 }) } ; varn = varn ; varl = varl } next exit
  = next record { vart = 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}
  　　; varn = varn ; varl = varl }


-- skew exitがsplitへ　nextが左回転
skew-bt : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
skew-bt node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit =  exit node
skew-bt node@record { vart = (bt-node record { key = key ; ltree = ltree ; rtree = bt-empty }) ; varn = varn ; varl = varl } next exit = exit node
skew-bt node@record { vart = (bt-node record { key = key@$\_{1}$@
  ; ltree = ltree@$\_{1}$@
  ; rtree = (bt-node record { key = record { coler = black ; number = number }
  ; ltree = ltree ; rtree = rtree }) }) ; varn = varn ; varl = varl } next exit
    = exit node
skew-bt node@record { vart = (bt-node record { key = key@$\_{1}$@
  ; ltree = ltree@$\_{1}$@
  ; rtree = (bt-node record { key = record { coler = red ; number = number } ; ltree = ltree ; rtree = rtree }) })
  ; varn = varn ; varl = varl } next exit
    = next node


set-node-coler : Env @$\rightarrow$@ Env
set-node-coler node@record { vart = bt-empty ; varn = varn ; varl = varl } = node
set-node-coler record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }
  = record { vart = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }

init-node-coler : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
init-node-coler node@record { vart = bt-empty ; varn = varn ; varl = varl } next exit = exit node
init-node-coler record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit
  = exit record { vart = (bt-node record { key = record { coler = black ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl }

--修正前
merge-tree node@record { vart = vart ; varn = varn ; varl = [] } next exit = exit node
merge-tree node@record { vart = vart ; varn = varn ; varl = (bt-empty @$\text{::}$@ varl@$\_{1}$@) } next exit = exit node

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } @$\text{::}$@ varl@$\_{1}$@) } next exit with <-cmp varn number

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } @$\text{::}$@ varl@$\_{1}$@) } next exit | tri≈ @$\neg$@a b @$\neg$@c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree }) ; varn = number ; varl = varl@$\_{1}$@ }

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } @$\text{::}$@ varl@$\_{1}$@) } next exit | tri> @$\neg$@a @$\neg$@b c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = vart }) ; varn = number ; varl = varl@$\_{1}$@ }

merge-tree record { vart = vart ; varn = varn ; varl = (bt-node record { key = record { coler = coler ; number = number } ; ltree = ltree ; rtree = rtree } @$\text{::}$@ varl@$\_{1}$@) } next exit | tri< a @$\neg$@b @$\neg$@c
  = next record { vart = (bt-node record { key = record { coler = coler ; number = number } ; ltree = vart ; rtree = rtree }) ; varn = number ; varl = varl@$\_{1}$@ }

-- do marge-tree
-- next merge-tree-c
-- exit fin
merge-tree-c :  {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
merge-tree-c env next exit with varl env
... | [] = exit env
... | bt-empty @$\text{::}$@ nl = exit env
... | bt-node xtree @$\text{::}$@ varl with <-cmp (varn env) (number ( key xtree ))
... | tri≈ @$\neg$@a b @$\neg$@c = next (record env { vart = bt-node xtree ; varn = number (key xtree) ; varl = varl })
... | tri> @$\neg$@a @$\neg$@b c = next (record env { vart = (bt-node record xtree{rtree = Env.vart env}) ; varn = number (key xtree) ; varl = varl })
... | tri< a @$\neg$@b @$\neg$@c = next (record env { vart = (bt-node record xtree{ltree = Env.vart env}) ; varn = number (key xtree) ; varl = varl })

-- do brandh
-- next insert-c
-- exit merge-tree
insert-c : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
insert-c env next exit with Env.vart env
... | bt-empty = exit record env{vart = (bt-node (record { key = record { coler = red ; number = Env.varn env }; ltree = bt-empty; rtree = bt-empty }))}
... | bt-node node with <-cmp (Env.varn env) (node.number (tree.key node))
... | tri≈ @$\neg$@a b @$\neg$@c = exit env
... | tri> @$\neg$@a @$\neg$@b c = next record env{vart = (tree.ltree node) ; varl = (bt-node record node{ltree = bt-empty}) @$\text{::}$@ (Env.varl env) }
... | tri< a @$\neg$@b @$\neg$@c = next record env{vart = (tree.rtree node) ; varl = (bt-node record node{rtree = bt-empty}) @$\text{::}$@ (Env.varl env) }

-- insert
bt-insert : {le : Level} {t : Set le} @$\rightarrow$@ Env @$\rightarrow$@ (next : Env @$\rightarrow$@ t) @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
-- mergeへ遷移する
bt-insert (record { vart = bt-empty ; varn = varn ; varl = varl }) next exit
  = exit (record { vart = (bt-node (record { key = record { coler = red ; number = varn }; ltree = bt-empty; rtree = bt-empty })) ; varn = varn ; varl = varl })

-- 場合分けを行う
bt-insert record {vart = (bt-node record { key = record { coler = y; number = x } ; ltree = ltree ; rtree = rtree }) ; varn = n ; varl = varl } next exit with <-cmp x n
bt-insert node@(record { vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } ) next exit | tri≈ @$\neg$@a b @$\neg$@c
  = exit node

bt-insert record {vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit | tri> @$\neg$@a @$\neg$@b c
  = next (record {vart = ltree ; varn = varn ; varl = (bt-node record { key = key ; ltree = bt-empty ; rtree = rtree }) @$\text{::}$@ varl })

bt-insert record {vart = (bt-node record { key = key ; ltree = ltree ; rtree = rtree }) ; varn = varn ; varl = varl } next exit | tri< a @$\neg$@b @$\neg$@c
  = next (record {vart = rtree ; varn = varn ; varl = (bt-node record { key = key ; ltree = ltree ; rtree = bt-empty }) @$\text{::}$@ varl })

-- normal loop without termination
{-@$\#$@ TERMINATING @$\#$@-}
insert : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
insert env exit = bt-insert env (@$\lambda$@ env @$\rightarrow$@ insert env exit ) exit

init-col : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
init-col env exit = init-node-coler env exit exit

{-
{-@$\#$@ TERMINATING @$\#$@-}
merge : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t

{-@$\#$@ TERMINATING @$\#$@-}
split-p : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
split-p env exit = split env (@$\lambda$@ env @$\rightarrow$@ merge env (@$\lambda$@ env @$\rightarrow$@ merge env exit ) ) exit

{-@$\#$@ TERMINATING @$\#$@-}
rotate_right : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
rotate_right env exit = merge-rotate-right env (@$\lambda$@ env @$\rightarrow$@ split-p env (@$\lambda$@ env @$\rightarrow$@ merge env exit ) ) exit

{-@$\#$@ TERMINATING @$\#$@-}
split-b : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
split-b env exit = split-branch env (@$\lambda$@ env @$\rightarrow$@ rotate_right env exit ) @$\lambda$@ env @$\rightarrow$@ merge env exit

{-@$\#$@ TERMINATING @$\#$@-}
rotate_left : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
rotate_left env exit = merge-rotate-left env (@$\lambda$@ env @$\rightarrow$@ split-b env exit ) exit

{-@$\#$@ TERMINATING @$\#$@-}
skew : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
skew env exit = skew-bt env (@$\lambda$@ env @$\rightarrow$@ rotate_left env (@$\lambda$@ env @$\rightarrow$@ split-b env exit ) ) exit

merge env exit = merge-tree env (@$\lambda$@ env @$\rightarrow$@ skew env exit ) exit
-}
-- skewはnextがrotate-right。exitはsplitとなる
-- skewの中にrotate-rightが内包され、実行後は必ずsplitに遷移する
-- それはskewのexitと等しいので同時に記述してやる。
skew' : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
skew' env exit = skew-bt env (@$\lambda$@ env @$\rightarrow$@ merge-rotate-left env exit exit ) exit

split' : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
split' env exit = split-branch env (@$\lambda$@ env @$\rightarrow$@ merge-rotate-right env (@$\lambda$@ env @$\rightarrow$@ split env exit exit ) (@$\lambda$@ env @$\rightarrow$@ split env exit exit ) ) exit


{-
merge' : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
merge' node@record { vart = vart ; varn = varn ; varl = [] } exit = exit node
merge' node@record { vart = vart ; varn = varn ; varl = (x @$\text{::}$@ varl@$\_{1}$@) } exit = merge' (merge-tree node (@$\lambda$@ env @$\rightarrow$@ skew' env (@$\lambda$@ env @$\rightarrow$@ split' env (@$\lambda$@ env @$\rightarrow$@ env) ) ) (@$\lambda$@ env @$\rightarrow$@ env) ) exit
-}

whileTestP : {l : Level} {t : Set l} @$\rightarrow$@ (tree : rbt) @$\rightarrow$@ (n : @$\mathbb{N}$@) @$\rightarrow$@ (Code : Env @$\rightarrow$@ t) @$\rightarrow$@ t
whileTestP tree n next = next (record {vart = tree ; varn = n ; varl = [] } )



{-@$\#$@ TERMINATING @$\#$@-}
mergeP : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
mergeP env exit = merge-tree env (@$\lambda$@ env @$\rightarrow$@ skew' env (@$\lambda$@ env @$\rightarrow$@ split' env (@$\lambda$@ env @$\rightarrow$@ mergeP env exit)) ) exit

{-
mergeP1 : {l : Level} {t : Set l} @$\rightarrow$@ Env @$\rightarrow$@ (exit : Env @$\rightarrow$@ t) @$\rightarrow$@ t
mergeP1 record { vart = vart ; varn = varn ; varl = [] } exit = {!!}
mergeP1 record { vart = vart ; varn = varn ; varl = (x @$\text{::}$@ varl@$\_{1}$@) } exit = {!!}
-}

{-
merge-tree env
  (@$\lambda$@ env @$\rightarrow$@ skew-bt env
    (@$\lambda$@ env @$\rightarrow$@ merge-rotate-left env
      (@$\lambda$@ env @$\rightarrow$@ split-branch env
        (@$\lambda$@ env @$\rightarrow$@  merge-rotate-right env
          (@$\lambda$@ env @$\rightarrow$@ split env (@$\lambda$@ env @$\rightarrow$@ mergeP env exit ) (@$\lambda$@ env @$\rightarrow$@ mergeP env exit ) ) exit)
      (@$\lambda$@ env @$\rightarrow$@ mergeP env exit ) )
    exit )
  exit ) exit
-}

--whileTestP : {l : Level} {t : Set l} @$\rightarrow$@ (c10 : @$\mathbb{N}$@) @$\rightarrow$@ (c11 : @$\mathbb{N}$@) @$\rightarrow$@ (Code : Envc @$\rightarrow$@ t) @$\rightarrow$@ t
--whileTestP c10 n next = next (record {varn = c10 ; vari = 0 ; c10 = c10 } )

--whileTestPCall : (tree : rbt) @$\rightarrow$@ (n : @$\mathbb{N}$@)  @$\rightarrow$@ Env
--whileTestPCall tree n = whileTestP {_} {_} tree n (@$\lambda$@ env @$\rightarrow$@ insert env (@$\lambda$@ env @$\rightarrow$@ merge env (@$\lambda$@ env @$\rightarrow$@ init-col env (@$\lambda$@ env @$\rightarrow$@ env ) ) ) )

whileTestPCall' : (tree : rbt) @$\rightarrow$@ (n : @$\mathbb{N}$@)  @$\rightarrow$@ Env
whileTestPCall' tree n = whileTestP {_} {_} tree n (@$\lambda$@ env @$\rightarrow$@ insert env (@$\lambda$@ env @$\rightarrow$@ mergeP env (@$\lambda$@ env @$\rightarrow$@ init-col env (@$\lambda$@ env @$\rightarrow$@ env ) ) ) )

-- 以下test部分
test1 = whileTestPCall' bt-empty 8
test2 = whileTestPCall' (vart test1) 10
test3 = whileTestPCall' (vart test2) 24
test4 = whileTestPCall' (vart test3) 31
test5 = whileTestPCall' (vart test4) 41
test6 = whileTestPCall' (vart test5) 45
test7 = whileTestPCall' (vart test6) 50
test8 = whileTestPCall' (vart test7) 59
test9 = whileTestPCall' (vart test8) 73
test10 = whileTestPCall' (vart test9) 74
test11 = whileTestPCall' (vart test10) 79