(with_import "./match.kp"
(let (
    ; This is based on https://www.cs.cornell.edu/courses/cs3110/2020sp/a4/deletion.pdf
    ; and the figure references refer to it
    ; Insert is taken from the same paper, but is origional to Okasaki, I belive

    ; The tree has been modified slightly to take in a comparison function
    ; and override if insert replaces or not to allow use as a set or as a map

    ; I think this is actually pretty cool - instead of having a bunch of seperate ['B]
    ; be our leaf node, we use ['B] with all nils. This allows us to not use -B, as
    ; both leaf and non-leaf 'BB has the same structure with children! Also, we make
    ; sure to use empty itself so we don't make a ton of empties...
    empty ['B nil nil nil]
    E  empty
    EE ['BB nil nil nil]

    size (rec-lambda recurse (t) (match t
                                        ~E  0
                                        [c a x b] (+ 1 (recurse a) (recurse b))))

    generic-foldl (rec-lambda recurse (f z t) (match t
                                                     ~E z
                                                     [c a x b]  (recurse f (f (recurse f z a) x) b)))

    generic-contains? (rec-lambda recurse (t cmp v found not-found) (match t
                                            ~E                            (not-found)
                                            [c a x b] (match (cmp v x) '< (recurse a cmp v found not-found)
                                                                       '= (found x)
                                                                       '> (recurse b cmp v found not-found))))

    blacken (lambda (t) (match t
                            ['R a x b] ['B a x b]
                            t          t))
    balance (lambda (t) (match t
                            ; figures 1 and 2
                            ['B ['R ['R a x b] y c] z d] ['R ['B a x b] y ['B c z d]]
                            ['B ['R a x ['R b y c]] z d] ['R ['B a x b] y ['B c z d]]
                            ['B a x ['R ['R b y c] z d]] ['R ['B a x b] y ['B c z d]]
                            ['B a x ['R b y ['R c z d]]] ['R ['B a x b] y ['B c z d]]
                            ; figure 8, double black cases
                            ['BB ['R a x ['R b y c]] z d] ['B ['B a x b] y ['B c z d]]
                            ['BB a x ['R ['R b y c] z d]] ['B ['B a x b] y ['B c z d]]
                            ; already balenced
                            t                            t))
    generic-insert (lambda (t cmp v replace) (let (
        ins (rec-lambda ins (t) (match t
                                    ~E                            ['R t v t]
                                    [c a x b] (match (cmp v x) '< (balance [c (ins a) x b])
                                                               '= (if replace [c a v b]
                                                                              t)
                                                               '> (balance [c a x (ins b)]))))
    ) (blacken (ins t))))

    rotate (lambda (t) (match t
      ; case 1, fig 6
      ['R ['BB a x b] y ['B c z d]]          (balance ['B ['R ['B a x b] y c] z d])
      ['R ['B a x b] y ['BB c z d]]          (balance ['B a x ['R b y ['B c z d]]])
      ; case 2, figure 7
      ['B ['BB a x b] y ['B c z d]]          (balance ['BB ['R ['B a x b] y c] z d])
      ['B ['B a x b] y ['BB c z d]]          (balance ['BB a x ['R b y ['B c z d]]])
      ; case 3, figure 9
      ['B ['BB a w b] x ['R ['B c y d] z e]] ['B (balance ['B ['R ['B a w b] x c] y d]) z e]
      ['B ['R a w ['B b x c]] y ['BB d z e]] ['B a w (balance ['B b x ['R c y ['B d z e]]])]
      ; fall through
      t                                     t))

    redden (lambda (t) (match t
                        ['B a x b] (if (and (= 'B (idx a 0)) (= 'B (idx b 0))) ['R a x b]
                                                                               t)
                        t                                                      t))

    min_delete (rec-lambda recurse (t) (match t
                    ~E                                      (error "min_delete empty tree")
                    ['R ~E x ~E]                            [x E]
                    ['B ~E x ~E]                            [x EE]
                    ['B ~E x ['R a y b]]                    [x ['B a y b]]
                    [c a x b]   (let ((v ap) (recurse a))   [v (rotate [c ap x b])])))
    generic-delete (lambda (t cmp v) (let (
        del (rec-lambda del (t v) (match t
                                    ; figure 3
                                    ~E                                          t
                                    ; figure 4
                                    ['R ~E x ~E]            (match (cmp v x) '= E
                                                                             _  t)
                                    ['B ['R a x b] y ~E]    (match (cmp v y) '< (rotate ['B (del ['R a x b] v) y ~E])
                                                                             '= ['B a x b]
                                                                             '> t)
                                    ; figure 5
                                    ['B ~E x ~E]            (match (cmp v x) '= EE
                                                                             _  t)
                                    [c a x b]               (match (cmp v x) '< (rotate [c (del a v) x b])
                                                                             '= (let ([vp bp] (min_delete b))
                                                                                (rotate [c a vp bp]))
                                                                             '> (rotate [c a x (del b v)]))))
    ) (del (redden t) v)))


    set-cmp (lambda (a b) (cond (< a b) '<
                                (= a b) '=
                                true    '>))
    set-empty empty
    set-foldl generic-foldl
    set-insert           (lambda (t x)   (generic-insert     t set-cmp x false))
    set-contains?        (lambda (t x)   (generic-contains?  t set-cmp x       (lambda (f) true)      (lambda () false)))
    set-remove           (lambda (t x)   (generic-delete     t set-cmp x))

    map-cmp (lambda (a b) (let (ak (idx a 0)
                                bk (idx b 0))
                                (cond (< ak bk) '<
                                      (= ak bk) '=
                                      true      '>)))
    map-empty empty
    map-insert           (lambda (t k v) (generic-insert     t map-cmp [k v] true))
    map-contains-key?    (lambda (t k)   (generic-contains?  t map-cmp [k nil] (lambda (f) true)      (lambda () false)))
    map-get              (lambda (t k)   (generic-contains?  t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () (error (str "didn't find key " k " in map " t)))))
    map-get-or-default   (lambda (t k d) (generic-contains?  t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () d)))
    map-get-with-default (lambda (t k d) (generic-contains?  t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () (d))))
    map-remove           (lambda (t k)   (generic-delete     t map-cmp [k nil]))

    ; This could be 2x as efficent by being implmented on generic instead of map,
    ; as we wouldn't have to traverse once to find and once to insert
    multimap-empty map-empty
    multimap-insert           (lambda (t k v) (map-insert t k (set-insert (map-get-or-default t k set-empty) v)))
    multimap-get              (lambda (t k)   (map-get-or-default t k set-empty))
)
    (provide set-empty set-foldl set-insert set-contains? set-remove
             map-empty map-insert map-contains-key? map-get map-get-or-default map-get-with-default map-remove
             multimap-empty multimap-insert multimap-get
             size)
))