#include "../../../../gen/l1_5/BufRBTree_SetS64.h"
#include "../../../../gen/l1/VecAndSpan_S64.h"
#include "../../../l1/core/VecU8_as_str.h"
#include "../../../../gen/l1/OptionS64.h"
#include "../../../l1/system/fileio.h"
#include "../../../l1/system/fsmanip.h"
#include "../../../l1/system/creating_child_proc.h"

void check_structure_h_dfs(const VecBufRBTreeNode* tree, U64 x, VecU8* f){
    if (x == 0)
        return;
    if (*VecU8_at(f, x - 1) != 0)
        check(false);
    *VecU8_mat(f, x - 1) = 1;
    check_structure_h_dfs(tree, VecBufRBTreeNode_at(tree, x)->left, f);
    check_structure_h_dfs(tree, VecBufRBTreeNode_at(tree, x)->right, f);
}

U32 check_structure_h_dfs_2(BufRBTreeNode* tree, U64 x){
    if (x == 0) {
        return 0;
    }
    U32 a = check_structure_h_dfs_2(tree, tree[x].left);
    U32 b = check_structure_h_dfs_2(tree, tree[x].right);
    check(a == b);
    return a + (tree[x].color == RBTree_black ? 1 : 0);
}

S64 min_key_in_subtree(const BufRBTree_SetS64* self, U64 x){
    assert(x != 0 && x < self->tree.len);
    S64 ans = self->el.buf[x - 1];
    if (self->tree.buf[x].left != 0) {
        ans = MIN_S64(ans, min_key_in_subtree(self, self->tree.buf[x].left));
    }
    if (self->tree.buf[x].right != 0) {
        ans = MIN_S64(ans, min_key_in_subtree(self, self->tree.buf[x].right));
    }
    return ans;
}

S64 max_key_in_subtree(const BufRBTree_SetS64* self, U64 x){
    assert(x != 0 && x < self->tree.len);
    S64 ans = self->el.buf[x - 1];
    if (self->tree.buf[x].left != 0) {
        ans = MAX_S64(ans, max_key_in_subtree(self, self->tree.buf[x].left));
    }
    if (self->tree.buf[x].right != 0) {
        ans = MAX_S64(ans, max_key_in_subtree(self, self->tree.buf[x].right));
    }
    return ans;
}

void check_only_structure(const BufRBTree_SetS64* self){
    check(self->tree.len == self->el.len + 1);
    check(self->root < self->tree.len);
    if (self->tree.len > 1) {
        check(self->root != 0);
    }
    VecU8 f = VecU8_new_zeroinit(self->tree.len - 1);
    check_structure_h_dfs(&self->tree, self->root, &f);
    for (size_t i = 0; i < self->tree.len - 1; i++) {
        check(*VecU8_at(&f, i));
    }
    VecU8_drop(f);  /* f invalidated */
    for (size_t v = 1; v < self->tree.len; v++) {
        if (v == self->root) {
            check(self->tree.buf[v].parent == 0);
        } else {
            check(self->tree.buf[v].parent != 0);
        }
        check(self->tree.buf[v].parent < self->tree.len);
        check(self->tree.buf[v].left < self->tree.len);
        check(self->tree.buf[v].right < self->tree.len);
        if (self->tree.buf[v].left != 0) {
            check(self->tree.buf[self->tree.buf[v].left].parent == v);
        }
        if (self->tree.buf[v].right != 0) {
            check(self->tree.buf[self->tree.buf[v].right].parent == v);
        }
    }
}

void check_correctness(const BufRBTree_SetS64* self){
    check_only_structure(self);

    /* Checking coloring */
    check(self->tree.buf[0].color == RBTree_black);
    if (self->root != 0)
        check(self->tree.buf[self->root].color == RBTree_black);
    for (size_t v = 1; v < self->tree.len; v++) {
        if (self->tree.buf[v].color == RBTree_red) {
            check(self->tree.buf[self->tree.buf[v].left].color == RBTree_black);
            check(self->tree.buf[self->tree.buf[v].right].color == RBTree_black);
        }
    }
    check_structure_h_dfs_2(self->tree.buf, self->root);

    /* checking keys, but better */
    for (size_t v = 1; v < self->tree.len; v++) {
        if (self->tree.buf[v].left != 0) {
            check(self->el.buf[self->tree.buf[v].left - 1] < self->el.buf[v - 1]);
        }
        if (self->tree.buf[v].right != 0) {
            check(self->el.buf[v - 1] < self->el.buf[self->tree.buf[v].right - 1]);
        }
    }
    for (size_t v = 1; v < self->tree.len; v++) {
        if (self->tree.buf[v].left != 0) {
            check(max_key_in_subtree(self, self->tree.buf[v].left) < self->el.buf[v - 1]);
        }
        if (self->tree.buf[v].right != 0) {
            check(self->el.buf[v - 1] < min_key_in_subtree(self, self->tree.buf[v].right));
        }
    }
}

void save_tree_to_file(const BufRBTree_SetS64* set, SpanU8 name){
    check_only_structure(set);

    VecS64 ranked_node_to_rb_node = VecS64_new_reserved(set->el.len);
    {
        VecS64 bfs_nxt = VecS64_new();
        VecS64 bfs = VecS64_new();
        if (set->root != 0)
            VecS64_append(&bfs_nxt, (S64)set->root);
        while (bfs_nxt.len > 0) {
            VecS64 t = bfs_nxt;
            bfs_nxt = bfs;
            bfs = t;
            for (size_t j = 0; j < bfs.len; j++) {
                S64 cur = *VecS64_at(&bfs, j);
                assert(cur != 0);
                VecS64_append(&ranked_node_to_rb_node, cur);
                if (cur < 0)
                    continue;
                if (set->tree.buf[cur].left != 0)
                    VecS64_append(&bfs_nxt, (S64)set->tree.buf[cur].left);
                else if (set->tree.buf[cur].right != 0)
                    VecS64_append(&bfs_nxt, -cur);
                if (set->tree.buf[cur].right != 0)
                    VecS64_append(&bfs_nxt, (S64)set->tree.buf[cur].right);
                else if (set->tree.buf[cur].left != 0)
                    VecS64_append(&bfs_nxt, -cur);
            }
            bfs.len = 0;
        }
    }
    VecU8 graph = VecU8_new();
    VecU8_append_span(&graph, cstr(
        "digraph rbtree {\n"
        "    fontsize = 20\n"
        "    rankdir = TB\n"
        "    bgcolor = \"lightgray\"\n"
        "    node [fontname = \"Arial\", style=\"rounded,filled\", shape=circle];\n"));
    for (U64 j = 0; j < ranked_node_to_rb_node.len; j++) {
        S64 cur = ranked_node_to_rb_node.buf[j];
        assert(cur != 0);
        if (cur < 0) {
            VecU8_append_vec(&graph, VecU8_fmt("    n_%i [style=invis, label=\"\"]\n", -cur));
            continue;
        }
        U64 i = (U64)cur;
        VecU8_append_vec(&graph, VecU8_fmt(
            "    v%u [label=%i, xlabel=<<FONT color='%s'>%u</FONT>>, color=%s, fontcolor=%s]\n",
            i, set->el.buf[i - 1],
            set->tree.buf[i].color == RBTree_black ? cstr("black") : cstr("red"),  i,
            set->tree.buf[i].color == RBTree_black ? cstr("black") : cstr("red"),
            set->tree.buf[i].color == RBTree_black ? cstr("white") : cstr("black")));
    }
    VecU8_append_span(&graph, cstr("\n"));
    for (U64 j = 0; j < ranked_node_to_rb_node.len; j++) {
        S64 cur = ranked_node_to_rb_node.buf[j];
        if (cur < 0)
            continue;
        U64 i = (U64)cur;
        if (set->tree.buf[i].left != 0) {
            VecU8_append_vec(&graph, VecU8_fmt("    v%u -> v%u [weight=2]\n", i, set->tree.buf[i].left));
        } else if (set->tree.buf[i].right != 0) {
            VecU8_append_vec(&graph, VecU8_fmt("    v%u -> n_%u [style=invis, weight=3]\n", i, i));
        }
        if (set->tree.buf[i].right != 0) {
            VecU8_append_vec(&graph, VecU8_fmt("    v%u -> v%u [weight=2]\n", i, set->tree.buf[i].right));
        } else if (set->tree.buf[i].left != 0) {
            VecU8_append_vec(&graph, VecU8_fmt("    v%u -> n_%u [style=invis, weight=3]\n", i, i));
        }
    }
    VecU8_append_span(&graph, cstr("}\n"));
    mkdir_nofail("GRAPHS");
    VecU8 dot_filename_nt = VecU8_fmt("GRAPHS/GRAPH_%s.gv%c", name, 0);
    write_whole_file_or_abort((CSTR)dot_filename_nt.buf, VecU8_to_span(&graph));
    VecU8_drop(graph);
    VecU8_drop(dot_filename_nt);
    VecU8 command_nt = VecU8_fmt("dot -Tpng GRAPHS/GRAPH_%s.gv -o GRAPHS/GRAPH_%s.png%c", name, name, 0);
    calling_system_func_nofail((CSTR)command_nt.buf);
    VecU8_drop(command_nt);
}

OptionS64 VecS64_find_max_less(const VecS64* set, S64 key){
    OptionS64 max_less = None_S64();
    for (size_t i = 0; i < set->len; i++)
        if (set->buf[i] < key)
            if (max_less.variant == Option_None || max_less.some < set->buf[i])
                max_less = Some_S64(set->buf[i]);
    return max_less;
}

OptionS64 VecS64_find_max_less_or_eq(const VecS64* set, S64 key){
    OptionS64 max_less = None_S64();
    for (size_t i = 0; i < set->len; i++)
        if (set->buf[i] <= key)
            if (max_less.variant == Option_None || max_less.some < set->buf[i])
                max_less = Some_S64(set->buf[i]);
    return max_less;
}


OptionS64 VecS64_find_min_grtr(const VecS64* set, S64 key){
    OptionS64 min_grtr = None_S64();
    for (size_t i = 0; i < set->len; i++)
        if (set->buf[i] > key)
            if (min_grtr.variant == Option_None || min_grtr.some > set->buf[i])
                min_grtr = Some_S64(set->buf[i]);
    return min_grtr;
}


OptionS64 VecS64_find_min_grtr_or_eq(const VecS64* set, S64 key){
    OptionS64 min_grtr = None_S64();
    for (size_t i = 0; i < set->len; i++)
        if (set->buf[i] >= key)
            if (min_grtr.variant == Option_None || min_grtr.some > set->buf[i])
                min_grtr = Some_S64(set->buf[i]);
    return min_grtr;
}

void check_option_int_equiv_node(OptionS64 true_ans, BufRBTree_SetS64* set, U64 my_ans){
    if (true_ans.variant == Option_None) {
        check(my_ans == 0);
    } else {
        check(my_ans != 0 && my_ans < set->tree.len);
        check(set->el.buf[my_ans - 1] == true_ans.some);
    }
}


void insert_only(){
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        check_correctness(&set);
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        bool ret = BufRBTree_SetS64_insert(&set, 42);
        check(ret);
        check_correctness(&set);
        check(set.el.len == 1);
        check(set.root == 1);
        BufRBTree_SetS64_drop(set);
    }
    {
        bool ret;
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        ret = BufRBTree_SetS64_insert(&set, 42);
        check(ret);
        check_correctness(&set);
        ret = BufRBTree_SetS64_insert(&set, 69);
        check(ret);
        check_correctness(&set);
        check(set.el.len == 2);
        BufRBTree_SetS64_drop(set);
    }
    {
        bool ret;
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        ret = BufRBTree_SetS64_insert(&set, 70);
        check(ret);
        check_correctness(&set);
        ret = BufRBTree_SetS64_insert(&set, 50);
        check(ret);
        check_correctness(&set);
        check(set.el.len == 2);
        BufRBTree_SetS64_drop(set);
    }
    {
        bool ret;
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        ret = BufRBTree_SetS64_insert(&set, 1);
        check(ret);
        check_correctness(&set);
        ret = BufRBTree_SetS64_insert(&set, 2);
        check(ret);
        check_correctness(&set);
        ret = BufRBTree_SetS64_insert(&set, 3);
        check(ret);
        check_correctness(&set);
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        for (S64 i = 10; i < 100; i++) {
            bool ret = BufRBTree_SetS64_insert(&set, i);
            check(ret);
            check_correctness(&set);
        }
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        for (S64 i = 99; i >= 10; i--) {
            bool ret = BufRBTree_SetS64_insert(&set, i);
            check(ret);
            check_correctness(&set);
        }
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        for (S64 i = 10; i < 100; i++) {
            bool ret = BufRBTree_SetS64_insert(&set, i);
            check(ret);
            check_correctness(&set);
            bool ret2 = BufRBTree_SetS64_insert(&set, 1000 - i);
            check(ret2);
            check_correctness(&set);
        }
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        // BufRBTree_SetS64_erase_substitute()
        BufRBTree_SetS64_drop(set);
    }
}

void insert_and_then_find(){
    BufRBTree_SetS64 set = BufRBTree_SetS64_new();
    for (S64 p = 0; p < 100; p++) {
        bool ret = BufRBTree_SetS64_insert(&set, 2 * p);
        check(ret);
        check_correctness(&set);
        for (S64 i = 0; i <= p; i++) {
            U64 ret1 = BufRBTree_SetS64_find(&set, 2 * i);
            check(ret1);
            U64 ret2 = BufRBTree_SetS64_find(&set, 2 * i + 1);
            check(ret2 == 0);
        }
    }
    BufRBTree_SetS64_drop(set);
}

void insert_and_delete(){
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        check(BufRBTree_SetS64_insert(&set, 1))
        check_correctness(&set);
        check(BufRBTree_SetS64_find(&set, 1));
        check(BufRBTree_SetS64_erase(&set, 1));
        check_correctness(&set);
        check(!BufRBTree_SetS64_find(&set, 1));
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        check(BufRBTree_SetS64_insert(&set, 1))
        check_correctness(&set);
        check(BufRBTree_SetS64_insert(&set, 2))
        check_correctness(&set);
        check(BufRBTree_SetS64_erase(&set, 2));
        check_correctness(&set);
        check(BufRBTree_SetS64_find(&set, 1));
        check(!BufRBTree_SetS64_find(&set, 2));
        check(BufRBTree_SetS64_erase(&set, 1));
        check_correctness(&set);
        check(!BufRBTree_SetS64_find(&set, 1));
        check(!BufRBTree_SetS64_find(&set, 2));
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        check(BufRBTree_SetS64_insert(&set, 1))
        check_correctness(&set);
        check(BufRBTree_SetS64_insert(&set, 2))
        check_correctness(&set);
        check(BufRBTree_SetS64_erase(&set, 1));
        check_correctness(&set);
        check(!BufRBTree_SetS64_find(&set, 1));
        check(BufRBTree_SetS64_find(&set, 2));

        check(!BufRBTree_SetS64_erase(&set, 1));
        check(!BufRBTree_SetS64_find(&set, 1));
        check(BufRBTree_SetS64_find(&set, 2));
        check(BufRBTree_SetS64_erase(&set, 2));
        check_correctness(&set);
        check(!BufRBTree_SetS64_find(&set, 1));
        check(!BufRBTree_SetS64_find(&set, 2));
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        for (S64 i = -120; i < 1000; i += 10) {
            check(BufRBTree_SetS64_insert(&set, i))
            check_correctness(&set);
        }
        // save_tree_to_file(&set, cstr("tree"));
        BufRBTree_SetS64_drop(set);
    }
    {
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        check(BufRBTree_SetS64_insert(&set, 1))
        check_correctness(&set);
        check(BufRBTree_SetS64_insert(&set, 2))
        check_correctness(&set);
        check(BufRBTree_SetS64_insert(&set, 3))
        check_correctness(&set);
        // save_tree_to_file(&set, cstr("trio"));
        check(BufRBTree_SetS64_erase(&set, 2));
        check_correctness(&set);
        // save_tree_to_file(&set, cstr("duo"));
        BufRBTree_SetS64_drop(set);
    }
}


void MutSpanS64_shuffle(MutSpanS64 span){
    if (span.len <= 1)
        return;
    for (size_t i = span.len-1; i > 0; i--) {
        size_t j = rand() % (i + 1);
        S64 temp = span.data[i];
        span.data[i] = span.data[j];
        span.data[j] = temp;
    }
}

VecS64 gen_cool_sequence(S64 i){
    VecS64 res = VecS64_new_reserved(i);
    for (S64 j = 0; j < i; j++) {
        VecS64_append(&res, (j + 1) * 10);
    }
    MutSpanS64_shuffle(VecS64_to_mspan(&res));
    return res;
}

void chaos(){
    for (int i = 1; i < 100; i++) {
        srand(i);
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        VecS64 seq_to_insert = gen_cool_sequence(i);
        for (int j = 0; j < i; j++) {
            S64 to_insert = seq_to_insert.buf[j];
            BufRBTree_SetS64_insert(&set, to_insert);
            check_correctness(&set);
        }
        assert(set.el.len == (size_t)i);
        // save_tree_to_file(&set, cstr("last_good"));
        S64 to_delete = (rand() % i + 1) * 10;
        printf("Started deleting from i=%d\n", i);
        check(BufRBTree_SetS64_erase(&set, to_delete));
        printf("Deleted from i=%d\n", i);
        check_only_structure(&set);
        // save_tree_to_file(&set, cstr("chucked_tree"));
        check_correctness(&set);
        VecS64_drop(seq_to_insert);
        BufRBTree_SetS64_drop(set);
    }
}

void triple_black(){
    BufRBTree_SetS64 set;
    {
        VecBufRBTreeNode tree = VecBufRBTreeNode_new_zeroinit(4);
        tree.buf[1] = (BufRBTreeNode){.parent = 3};
        tree.buf[2] = (BufRBTreeNode){.parent = 3};
        tree.buf[3] = (BufRBTreeNode){.left = 2, .right = 1};
        VecS64 el = VecS64_new_reserved(3);
        VecS64_append(&el, 300);
        VecS64_append(&el, 100);
        VecS64_append(&el, 200);
        set = (BufRBTree_SetS64){.tree = tree, .root = 3, .el = el};
    }
    check_correctness(&set);
    check(BufRBTree_SetS64_erase(&set, 300));
    check_correctness(&set);
    BufRBTree_SetS64_drop(set);
}

void absolute_chaos(){
    for (int i = 1; i < 100; i++) {
        srand(i);
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        VecS64 seq_to_insert = gen_cool_sequence(i);
        for (int j = 0; j < i; j++) {
            S64 to_insert = seq_to_insert.buf[j];
            check(BufRBTree_SetS64_insert(&set, to_insert));
            check_correctness(&set);
        }
        printf("Filled chaos with i = %d\n", i);
        // save_tree_to_file(&set, cstr("i_3"));
        assert(set.el.len == (size_t)i);
        VecS64 seq_to_erase = gen_cool_sequence(i);
        for (int j = 0; j < i; j++) {
            S64 to_erase = seq_to_erase.buf[j];
            check(BufRBTree_SetS64_erase(&set, to_erase));
            check_correctness(&set);
        }
        printf("Emptied chaos with i = %d\n", i);
        assert(set.el.len == 0);
        VecS64_drop(seq_to_insert);
        VecS64_drop(seq_to_erase);
        BufRBTree_SetS64_drop(set);
    }
}

void stress(){
    printf("Prepare for some real stress\n");
    for (int s = 2; s < 1000; s++) {
        srand(s);
        BufRBTree_SetS64 set = BufRBTree_SetS64_new();
        VecS64 real_set = VecS64_new();
        U32 n = (U32)s;
        VecS64 complementary_set = VecS64_new_reserved(n);
        for (size_t i = 0; i < n; i++) {
            VecS64_append(&complementary_set, (S64)i * 10);
        }
        for (int t = 0; t < 1000; t++) {
            /* Do something */
            bool do_insertion = (rand() % 9) >= 4;
            if (real_set.len == 0)
                do_insertion = 1;
            if (complementary_set.len == 0)
                do_insertion = 0;
            if (do_insertion) {
                assert(complementary_set.len > 0);
                size_t j = rand() % complementary_set.len;
                S64 v = VecS64_unordered_pop(&complementary_set, j);
                VecS64_append(&real_set, v);
                check(BufRBTree_SetS64_insert(&set, v));
                check_correctness(&set);
            } else {
                assert(real_set.len > 0);
                size_t j = rand() % real_set.len;
                S64 v = VecS64_unordered_pop(&real_set, j);
                VecS64_append(&complementary_set, v);
                check(BufRBTree_SetS64_erase(&set, v));
                check_correctness(&set);
            }
            /* We did something */
            check(real_set.len == set.el.len);
            for (size_t j = 0; j < real_set.len; j++) {
                check(BufRBTree_SetS64_find(&set, real_set.buf[j]) != 0);
            }
            for (size_t j = 0; j < complementary_set.len; j++) {
                check(BufRBTree_SetS64_find(&set, complementary_set.buf[j]) == 0);
            }
            /* Checking weird methods that nobody needs */
            for (int tt = 0; tt < 15; tt++) {
                S64 key = (S64)(rand() % n) * 10;
                {
                    OptionS64 true_ans = VecS64_find_max_less(&real_set, key);
                    U64 my_ans = BufRBTree_SetS64_find_max_less(&set, key);
                    check_option_int_equiv_node(true_ans, &set, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_max_less_or_eq(&real_set, key);
                    U64 my_ans = BufRBTree_SetS64_find_max_less_or_eq(&set, key);
                    check_option_int_equiv_node(true_ans, &set, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_min_grtr(&real_set, key);
                    U64 my_ans = BufRBTree_SetS64_find_min_grtr(&set, key);
                    check_option_int_equiv_node(true_ans, &set, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_min_grtr_or_eq(&real_set, key);
                    U64 my_ans = BufRBTree_SetS64_find_min_grtr_or_eq(&set, key);
                    check_option_int_equiv_node(true_ans, &set, my_ans);
                }
            }
        }
        VecS64_drop(real_set);
        VecS64_drop(complementary_set);
        BufRBTree_SetS64_drop(set);
        printf("Seed s=%d passed test\n", s);
    }
}

int main() {
    insert_only();
    insert_and_then_find();
    insert_and_delete();
    chaos();
    triple_black();
    absolute_chaos();
    stress();
    return 0;
}
