RBTree(k, v) := [
    Empty, 
    Node({
        colour: [R, B], 
        left: RBTree(k, v),
        key: k,
        value: v,
        right: RBTree(k, v)
    })
].{


    Intermediate(a) := [
        Done(a), ToDo(a)
    ].{
        apply : Intermediate(a), (a -> a) -> Intermediate(a)
        apply = |inter, f| {
            match inter {
                Done(x) => Done(f(x))
                ToDo(x) => ToDo(f(x))
            }
        }

        unit : a -> Intermediate(a)
        unit = |x| {
            ToDo(x)
        }

        done : a -> Intermediate(a)
        done = |x| Done(x)

        and_then : Intermediate(a), (a -> Intermediate(a)) -> Intermediate(a)
        and_then = |inter, f| {
            match inter {
                Done(x) => Done(x)
                ToDo(x) => f(x)
            }
        }

        from_inter : Intermediate(a) -> a
        from_inter = |inter| {
            match inter {
                Done(x) => x
                ToDo(x) => x
            }
        }
    }

    Inner(k, v) : {
        colour: [R, B], 
        left: RBTree(k, v),
        key: k,
        value: v,
        right: RBTree(k, v)
    }

    InProgress(k, v) : Intermediate(RBTree(k, v))
        
    empty : () -> RBTree(k, v) where [
        k.is_eq: k, k -> Bool,
        k.is_lt: k, k -> Bool,
        k.is_gt: k, k -> Bool
    ]
    empty = || {
        Empty
    }

    insert : RBTree(k, v), k, v -> RBTree(k, v) where [
        k.is_eq: k, k -> Bool,
        k.is_lt: k, k -> Bool,
        k.is_gt: k, k -> Bool
    ]
    insert = |tree, new_k, new_v| {
        make_black = |node| {
            match node {
                RBTree.Node(inner) => RBTree.Node({..inner, colour: B})
                RBTree.Empty => RBTree.Empty
            }
        }

        ins : RBTree(k, v) -> Intermediate(RBTree(k, v))
        ins = |into| {
            match into {
                Empty => RBTree.Node({colour: R,
                               left: Empty,
                               key: new_k,
                               value: new_v,
                               right: Empty
                               })->RBTree.Intermediate.unit()
                RBTree.Node({left, key} as before) if new_k < key => {
                    # Seperate steps to debug
                    ins(left).apply(|new_left| RBTree.Node({..before, left: new_left})).and_then(balance)

                    # Switching the below lines causes an error
                    # (malloc: heap corruption, free list damaged)
                    # b = a.apply(|new_left| RBTree.Node({colour: before.colour, key: before.key, value: before.value, right: before.right, left: new_left}))
                    # b = a
                    # b = match a {
                    #     Intermediate.Done(new_left) => {
                    #         Intermediate.Done(RBTree.Node({colour: before.colour, key: before.key, value: before.value, right: before.right, left: new_left}))
                    #     }
                    #     Intermediate.ToDo(new_left) => {
                    #         # ..before, left: new_left
                    #         Intermediate.ToDo(RBTree.Node({..before, value: before.value, left: new_left}))
                    #     }
                    # }
                    # bef = before
                    # b = a.apply(|new_left| RBTree.Node({..bef, left: new_left}))
                    # c = b.and_then(balance)
                    # c
                }
                RBTree.Node({colour, left, key, right}) if new_k == key => {
                    RBTree.Node({colour: colour, left: left, key: key, value: new_v, right: right})->RBTree.Intermediate.unit()
                }
                RBTree.Node({key, right} as before) if new_k > key => {
                    # balance(colour, left, key, value, ins(right))
                    ins(right).apply(|new_right| RBTree.Node({..before, right: new_right})).and_then(balance)


                    # b = a.apply(|new_right| {
                    #     # RBTree.node(before)
                    #     # RBTree.empty()
                    #     # Empty
                    #     # RBTree.Node(before)
                    #     # RBTree.Node({colour: before.colour, key: before.key, value: before.value, right: new_right, left: before.left})
                    # })
                    # b = a.apply(|new_right| RBTree.node({colour: before.colour, key: before.key, value: before.value, right: new_right, left: before.left}) )
                    # b = a.apply(|new_right| before)

                }
            }
        }

        tree->ins()->RBTree.Intermediate.from_inter()->make_black()
    }

    member : RBTree(k, v), k -> Bool where [
        k.is_eq: k, k -> Bool,
        k.is_lt: k, k -> Bool,
        k.is_gt: k, k -> Bool
    ]
    member = |tree, other_key| {
        match tree.get(other_key) {
            Just(_) => True
            Nothing => False
        }
    }

    get : RBTree(k, v), k -> [Just(v), Nothing] where [
        k.is_eq: k, k -> Bool,
        k.is_lt: k, k -> Bool,
        k.is_gt: k, k -> Bool
    ]
    get = |tree, k| {
        match tree {
            Empty => Nothing
            RBTree.Node({key, value}) if key == k => {
                Just(value)
            }
            RBTree.Node({left, key}) if k < key => {
                left.get(k)
            }
            RBTree.Node({key, right}) if k > key => {
                right.get(k)
            }
        }
    }

    values : RBTree(_, v) -> List(v)
    values = |tree| {
        match tree {
            Empty => {
                []
            }
            RBTree.Node({left, value, right}) => {
                ltree = left.values()
                rtree = right.values()
                ltree.append(value).concat(rtree)
            }
        }
    }

    keys : RBTree(k, _) -> List(k)
    keys = |tree| {
        match tree {
            Empty => {
                []
            }
            RBTree.Node({left, key, right}) => {
                ltree = left.keys()
                rtree = right.keys()
                ltree.append(key).concat(rtree)
            }
        }
    }

    delete = |tree, query| {

        make_black : RBTree(k, v) -> Intermediate(RBTree(k, v))
        make_black = |t| {
            match t {
                RBTree.Node({colour: R} as inner) => {
                    RBTree.Node({..inner, colour: B})->Intermediate.done()
                }
                anything_else => {
                    # In the context of deletion, make_black is used either
                    # when there is only one successor, or in balancing up-tree.
                    # If the current node is red, then no RR violation can be
                    # introduced by changing to black, and nodes are only
                    # blackened to account for (and fix) a BH violation.
                    # So if we try and blacken a node and find that it is
                    # already black, then it is because we needed
                    # to increase the black height by 1 to account for a black
                    # deletion, but the node was already black and so more 
                    # work must be done
                    anything_else->Intermediate.unit()
                }
            }
        }
        
        balanceD : RBTree(k, v) -> Intermediate(RBTree(k, v))
        balanceD = |t| {
            match t {
                RBTree.Node({left: RBTree.Node({colour: R, left: RBTree.Node({colour: R} as ll)} as l)} as p) => {
                    #     BR (p)            l
                    #    /                 / \
                    #   R (l)     =>     ll   p
                    #  /
                    # R (ll)
                    RBTree.Node({
                        colour: p.colour,               # Maintain the colour
                        key: l.key,                     # Data from the first left
                        value: l.value,
                        left: RBTree.Node({..ll, colour: B}),   # New left is just the left grandchild
                        right: Node({..p, colour: B, left: l.right}),  # New right is blackened parent
                    })->Intermediate.done()
                }
                Node({right: Node({colour: R, right: Node({colour: R} as rr)} as r)} as p) => {
                    # BR (p)            r
                    #  \               / \
                    #   R (r)     =>  p   rr
                    #    \
                    #     R (rr)
                    Node({
                        colour: p.colour,               # Maintain the colour
                        key: r.key,                     # Data from the first right
                        value: r.value,
                        left: Node({..p, colour: B, right: r.left}),   # New left is just the blackened parent
                        right: Node({..rr, colour: B}),  # New right is just the right grandchild
                    })->Intermediate.done()
                }
                Node({left: Node({colour: R, right: Node({colour: R} as lr)} as l)} as p) => {
                    #     BR (p)         lr
                    #    /              /  \
                    #   R (l)     =>   l    p
                    #    \
                    #     R (lr)
                    Node({
                        colour: p.colour,               # Maintain the colour
                        key: lr.key,                     # Data from grandchild
                        value: lr.value,
                        left: Node({..l, colour: B, right: lr.left}),   # New left is the child
                        right: Node({..p, colour: B, left: lr.right}),  # New right is the parent
                    })->Intermediate.done()
                }
                Node({right: Node({colour: R, left: Node({colour: R} as rl)} as r)} as p) => {
                    #     BR (p)         rl
                    #       \           /  \
                    #       R (r) =>   p    r
                    #      /
                    #     R (rl)
                    Node({
                        colour: p.colour,               # Maintain the colour
                        key: rl.key,                     # Data from grandchild
                        value: rl.value,
                        left: Node({..p, colour: B, right: rl.left}),   # New left is the child
                        right: Node({..r, colour: B, left: rl.right}),  # New right is the parent
                    })->Intermediate.done()
                }
                anything_else => {
                    make_black(anything_else)
                }
            }
        }

        eqL : RBTree(k, v) -> Intermediate(RBTree(k, v))
        eqL = |t| {
            # Left child is definitely short by 1
            # so the task is to balance the heights of
            # the siblings, either by raising the left
            # or reducing the right
            match t {
                Node({right: Node({colour: B} as right_inner)} as inner) => {
                    # Right sibling is black so make it red then balance.
                    # The whole tree is now short, and there can be RR(R) violations
                    # which are handled by balanceD
                    Node({..inner, right: Node({..right_inner, colour: R})})->balanceD()
                }
                Node({right: Node({colour: R} as right_inner)} as inner) => {
                    # Right sibling is red
                    # Rotate the nodes so that the right sibling is now a parent
                    # of the common parent, and the old parent is now a red
                    # left child.
                    # The rotation leaves a black height violation inside the new
                    # left subtree, except that new left subtree has a red root.
                    # Because there were no RR violations before, we know
                    # that the child of the new left subtree must be black, and
                    # so the prior eqL case will apply (black right sibling under (red) parent)
                    # The prior case is guaranteed to increase the black height
                    # when called on a red node as an outcome of calling balanceD
                    RBTree.Node(
                        {colour: R, left: inner.left, key: inner.key, value: inner.value, right: right_inner.left}
                    )->eqL().apply(|new_left| {
                        # Then (using result of the eqL, whether todo or done) put the
                        # new left (the old parent) inside the new parent (the old sibling)
                        RBTree.Node({colour: B, left: new_left, key: right_inner.key, value: right_inner.value, right: right_inner.right})
                    })
                }
                Empty => {
                    expect False # Would be nice to give a message?
                    Empty->Intermediate.done()
                }
                RBTree.Node({right: Empty}) => {
                    expect False
                    t->Intermediate.done()
                }

            }

        }

        eqR : RBTree(k, v) -> Intermediate(RBTree(k, v))
        eqR = |t| {
            # TODO Fix left and right in this fn (currently copied from eqL except the name)
            match t {
                RBTree.Node({left: RBTree.Node({colour: B} as left_data)} as parent_data) => {
                    RBTree.Node({..parent_data, left: RBTree.Node({..left_data, colour: R})})->balanceD()
                }
                RBTree.Node({left: RBTree.Node({colour: R} as left_data)} as parent_data) => {
                    uneq_right = RBTree.Node(
                        {colour: R, left: left_data.right, key: parent_data.key, value: parent_data.value, right: parent_data.right}
                    )
                    uneq_right->eqL().apply(|new_right| {
                        new_parent = RBTree.Node({colour: B, left: left_data.left, key: left_data.key, value: left_data.value, right: new_right})
                        new_parent
                    })
                }
                Empty => {
                    expect False # Would be nice to give a message?
                    Empty->Intermediate.done()
                }
                RBTree.Node({left: Empty}) => {
                    expect False
                    t->Intermediate.done()
                }
            }
        }

        delMin : Inner(k, v) -> (Intermediate(RBTree(k, v)), {key: k, value: v})
        delMin = |inner| {
            # This deletes the smallest value of the tree and returns it
            # alongside the new tree
            match inner {
                {colour: R, left: Empty, key, value, right} => {
                    (right->Intermediate.done(), {key, value})
                }
                {colour: B, left: Empty, key, value, right} => {
                    (right->make_black(), {key, value})
                }
                {colour, left: RBTree.Node(left_inner), key, value, right} => {
                    (new_left, deleted) = left_inner->delMin()
                    (
                        new_left.apply(|l| {
                            RBTree.Node({colour, left: l, key, value, right})
                        }).and_then(eqL), # Note this only applies if the left
                        # tree is ToDo, which indicates "definitely short"
                        deleted
                    )
                }
            }
        }
        delCurr : RBTree(k, v) -> Intermediate(RBTree(k, v))
        delCurr = |t| {
            match t {
                RBTree.Empty => {
                    # This happens when the key does not exist
                    # and we have descended to an empty node
                    # searching for it
                    RBTree.Empty->Intermediate.done()
                }
                RBTree.Node({colour: R, left, right: RBTree.Empty}) => {
                    # Deleting when there is no successor to the right
                    # is easy as we know there can't be a red with a red child
                    # and black height will be maintained
                    left->Intermediate.done()
                }
                RBTree.Node({colour: B, left, right: RBTree.Empty}) => {
                    # Deleting a black when there is no successor to the right
                    # is easy as we know the left child must be red (otherwise
                    # there would be a black height violation due to the empty
                    # right), and the red left child can be blackened
                    left->make_black()
                }
                RBTree.Node({colour, left, right: RBTree.Node(inner)}) => {
                    # The right node has a value
                    # so we need to delete the smallest from the right
                    # and replace the current value with it
                    (new_right, {key, value}) = inner->delMin()
                    # Having removed the successor, we push down the 'new' right
                    # tree below it to rotate, and do it inside an `apply` because
                    # it could be todo or done
                    new_right.apply(|r| {
                        RBTree.Node({colour, left, key, value, right: r})
                        # I think delMin preserves the RR constraint, so we only
                        # need to worry about black height on the children (specifically
                        # the new right hand side after rotating the smallest up), and
                        # eqR takes care of that
                    }).and_then(eqR)
                }
            }
        }
        del : RBTree(k, v) -> Intermediate(RBTree(k, v)) where [
            k.is_lt : k, k -> Bool,
            k.is_eq : k, k -> Bool
        ]
        del = |t| { 
            match t {
                RBTree.Empty => {
                    t->Intermediate.done()
                }
                RBTree.Node({ left, key } as inner) if query < key => {
                    del(left).apply(|new_left| {
                        RBTree.Node({..inner, left: new_left})
                    }).and_then(eqL)
                }
                RBTree.Node({ key }) if query == key => {
                    delCurr(t)
                }
                RBTree.Node({ right, key } as inner) if key < query => {
                    del(right).apply(|new_right| {
                        RBTree.Node({..inner, right: new_right})
                    }).and_then(eqR)
                }
            }
        }
    }
}


balance : RBTree(k, v) -> RBTree.Intermediate(RBTree(k, v))
balance = |tree_node| {
    new = |colour, (left, (k, v), right)| {
        RBTree.Node({colour: colour, left: left, key: k, value: v, right: right})
    }

    # Pulling it out to make it easy to pattern match
    {colour: this_colour, left: l_tree, right: r_tree, key: key, value: value} = match tree_node {
        # Balance on empty should never happen, but for exhaustiveness:
        Empty => {
            return RBTree.Intermediate.done(Empty)
        }
        RBTree.Node(inner) => inner
    }

    (left, centre, right) = match (this_colour, l_tree, r_tree) {
        (B, RBTree.Node({left: RBTree.Node(grandchild)} as child), _) if (child.colour == R and grandchild.colour == R) => {
            new_left = (grandchild.left, (grandchild.key, grandchild.value), grandchild.right)
            new_centre = (child.key, child.value)
            new_right = (child.right, (key, value), r_tree)
            (new_left, new_centre, new_right)
        }
        (B, RBTree.Node({right: RBTree.Node(grandchild)} as child), _) if (child.colour == R and grandchild.colour == R) => {
            new_left = (child.left, (child.key, child.value), grandchild.left)
            new_centre = (grandchild.key, grandchild.value)
            new_right = (grandchild.right, (key, value), r_tree)
            (new_left, new_centre, new_right)
        }
        (B, _, RBTree.Node({left: RBTree.Node(grandchild)} as child)) if (child.colour == R and grandchild.colour == R) => {
            new_left = (l_tree, (key, value), grandchild.left)
            new_centre = (grandchild.key, grandchild.value)
            new_right = (grandchild.right, (child.key, child.value), child.right)
            (new_left, new_centre, new_right)
        }
        (B, _, RBTree.Node({right: RBTree.Node(grandchild)} as child)) if (child.colour == R and grandchild.colour == R) => {
            new_left = (l_tree, (key, value), child.left)
            new_centre = (child.key, child.value)
            new_right = (grandchild.left, (grandchild.key, grandchild.value), grandchild.right)
            (new_left, new_centre, new_right)
        }
        (B, other_left, other_right) => {
            # Cannot cause any further issues
            # This is an optimisation to prevent further checking up the tree
            return new(this_colour, (other_left, (key, value), other_right))->RBTree.Intermediate.done()
        }
        (R, other_left, other_right) => {
            return new(this_colour, (other_left, (key, value), other_right))->RBTree.Intermediate.unit()
        }
    }
    left_node = new(B, left)
    right_node = new(B, right)
    # Produces red, so can have issues with its parents
    new(R, (left_node, centre, right_node))->RBTree.Intermediate.unit()
}