app "rbtree-ck"
    packages { base: "platform" }
    imports [base.Task]
    provides [ main ] to base


Color : [ Red, Black ]

Tree a b : [ Leaf, Node Color (Tree a b) a b (Tree a b) ]

Map : Tree I64 Bool

ConsList a : [ Nil, Cons a (ConsList a) ]

makeMap : I64, I64 -> ConsList Map
makeMap = \freq, n ->
    makeMapHelp freq n Leaf Nil

makeMapHelp : I64, I64, Map, ConsList Map -> ConsList Map
makeMapHelp = \freq, n, m, acc ->
    when n is
        0 -> Cons m acc
        _ ->
            powerOf10 =
                (n % 10 |> resultWithDefault 0) == 0


            m1 = insert m n powerOf10

            isFrequency =
                (n % freq |> resultWithDefault 0) == 0

            x = (if isFrequency then (Cons m1 acc) else acc)
            makeMapHelp freq (n-1) m1 x

fold : (a, b, omega -> omega), Tree a b, omega -> omega
fold = \f, tree, b ->
    when tree is
        Leaf  -> b
        Node _ l k v r -> fold f r (f k v (fold f l b))

resultWithDefault : Result a e, a -> a
resultWithDefault = \res, default ->
    when res is
        Ok v -> v
        Err _ -> default

main : Task.Task {} []
main =
    ms : ConsList Map
    ms = makeMap 5 5 # 42_000_00

    when ms is
        Cons head _ ->
            val = fold (\_, v, r -> if v then r + 1 else r) head 0
            val
                |> Str.fromInt
                |> Task.putLine

        Nil ->
            Task.putLine "fail"

insert : Map, I64, Bool -> Map
insert = \t, k, v -> if isRed t then setBlack (ins t k v) else ins t k v


setBlack : Tree a b -> Tree a b
setBlack = \tree ->
    when tree is
        Node _ l k v r -> Node Black l k v r
        _ -> tree

isRed : Tree a b -> Bool
isRed = \tree ->
    when tree is
        Node Red _ _ _ _ -> True
        _ -> False

lt = \x, y -> x < y

ins : Map, I64, Bool -> Map
ins = \tree, kx, vx ->
    when tree is
        Leaf ->
            Node Red Leaf kx vx Leaf

        Node Red a ky vy b ->
            if lt kx ky then
                Node Red (ins a kx vx) ky vy b
            else if lt ky kx then
                Node Red a ky vy (ins b kx vx)
            else
                Node Red a ky vy (ins b kx vx)

        Node Black a ky vy b ->
            if lt kx ky then
              (if isRed a then balance1 (Node Black Leaf ky vy b) (ins a kx vx) else Node Black (ins a kx vx) ky vy b)
            else if lt ky kx then
              (if isRed b then balance2 (Node Black a ky vy Leaf) (ins b kx vx) else Node Black a ky vy (ins b kx vx))
            else
                Node Black a kx vx b

balance1 : Map, Map -> Map 
balance1 = \tree1, tree2 ->
    when tree1 is
        Leaf -> Leaf
        Node _ _ kv vv t ->
            when tree2 is
                Node _ (Node Red l kx vx r1) ky vy r2 -> Node Red (Node Black l kx vx r1) ky vy (Node Black r2 kv vv t)
                Node _ l1 ky vy (Node Red l2 kx vx r) -> Node Red (Node Black l1 ky vy l2) kx vx (Node Black r kv vv t)
                Node _ l  ky vy r                     -> Node Black (Node Red l ky vy r) kv vv t
                Leaf -> Leaf

balance2 : Map, Map -> Map 
balance2 = \tree1, tree2 ->
    when tree1 is
        Leaf -> Leaf
        Node _ t kv vv _ ->
            when tree2 is
                Node _ (Node Red l kx1 vx1 r1) ky vy r2  -> Node Red (Node Black t kv vv l) kx1 vx1 (Node Black r1 ky vy r2)
                Node _ l1 ky vy (Node Red l2 kx2 vx2 r2) -> Node Red (Node Black t kv vv l1) ky vy (Node Black l2 kx2 vx2 r2)
                Node _ l ky vy r                         -> Node Black t kv vv (Node Red l ky vy r)
                Leaf -> Leaf