#include "../../../../gen/l1_5/RBTree_SetS64.h"
#include "../../../../gen/l1/VecAndSpan_U8.h"
#include "../../../../gen/l1/VecAndSpan_S64.h"
#include "../../../../gen/l1/OptionS64.h"

#include "../../../l1/core/VecU8_as_str.h"
#include "../../../l1/system/fileio.h"
#include "../../../l1/system/fsmanip.h"
#include "../../../l1/system/creating_child_proc.h"

typedef RBTreeNode_S64* RefRBTreeNode_S64;

#include "../../../../gen/l1/eve/ds_test/VecRefRBTreeNode_S64.h"

S64 min_key_in_subtree(RBTreeNode_S64* x, RBTreeNode* NIL){
    check(x != NULL && &x->base != NIL);
    S64 ans = x->key;
    if (x->base.left != NIL)
        ans = MIN_S64(ans, min_key_in_subtree((RBTreeNode_S64*)x->base.left, NIL));
    if (x->base.right != NIL)
        ans = MIN_S64(ans, min_key_in_subtree((RBTreeNode_S64*)x->base.right, NIL));
    return ans;
}

S64 max_key_in_subtree(RBTreeNode_S64* x, RBTreeNode* NIL){
    check(x != NULL && &x->base != NIL);
    S64 ans = x->key;
    if (x->base.left != NIL)
        ans = MAX_S64(ans, max_key_in_subtree((RBTreeNode_S64*)x->base.left, NIL));
    if (x->base.right != NIL)
        ans = MAX_S64(ans, max_key_in_subtree((RBTreeNode_S64*)x->base.right, NIL));
    return ans;
}

int check_structure_h_dfs(RBTree_SetS64* set, RBTreeNode_S64* x){
    check(x != NULL);
    if (&x->base == set->NIL)
        return 0;
    int a = x->base.left == set->NIL ? 0 : check_structure_h_dfs(set, (RBTreeNode_S64*)x->base.left);
    int b = x->base.right == set->NIL ? 0 : check_structure_h_dfs(set, (RBTreeNode_S64*)x->base.right);
    check(a == b);
    if (&x->base == set->root) {
        check(x->base.parent == set->NIL);
    } else {
        check(x->base.parent != set->NIL);
    }
    if (x->base.left != set->NIL)
        check(x->base.left->parent == &x->base);
    if (x->base.right != set->NIL)
        check(x->base.right->parent == &x->base);

    if (x->base.left != set->NIL)
        check(max_key_in_subtree((RBTreeNode_S64*)x->base.left, set->NIL) < x->key);
    if (x->base.right != set->NIL)
        check(x->key < min_key_in_subtree((RBTreeNode_S64*)x->base.right, set->NIL));

    if (x->base.color == RBTREE_RED) {
        check(x->base.left->color == RBTREE_BLACK);
        check(x->base.right->color == RBTREE_BLACK);
    }

    return a + (x->base.color == RBTREE_BLACK ? 1 : 0);
}

void check_correctness(RBTree_SetS64* set){
    check(set->root->color == RBTREE_BLACK);
    check_structure_h_dfs(set, (RBTreeNode_S64*)set->root);
}


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

    VecU8 graph = vcstr(
        "digraph rbtree {\n"
        "    fontsize = 20\n"
        "    rankdir = TB\n"
        "    bgcolor = \"lightgray\"\n"
        "    node [fontname = \"Arial\", style=\"rounded,filled\", shape=circle];\n\n");

    VecRefRBTreeNode_S64 bfs_nxt = VecRefRBTreeNode_S64_new();
    VecRefRBTreeNode_S64 bfs = VecRefRBTreeNode_S64_new();
    if (set->root != set->NIL) {
        check(set->root->parent == set->NIL);
        VecRefRBTreeNode_S64_append(&bfs_nxt, (RBTreeNode_S64*)set->root);
    }
    while (bfs_nxt.len > 0) {
        VecRefRBTreeNode_S64 t = bfs_nxt;
        bfs_nxt = bfs;
        bfs = t;
        for (size_t j = 0; j < bfs.len; j++) {
            RBTreeNode_S64* cur = *VecRefRBTreeNode_S64_at(&bfs, j);
            assert(cur != NULL);
            if (cur->base.left != set->NIL) {
                check(cur->base.left->parent == &cur->base);
                VecRefRBTreeNode_S64_append(&bfs, (RBTreeNode_S64*)cur->base.left);
            }

            if (cur->base.right != set->NIL) {
                check(cur->base.right->parent == &cur->base);
                VecRefRBTreeNode_S64_append(&bfs, (RBTreeNode_S64*)cur->base.right);
            }

            VecU8_append_vec(&graph, VecU8_fmt(
            "    v%i [label=%i, color=%s, fontcolor=%s]\n",
            cur->key, cur->key,
            cur->base.color == RBTREE_BLACK ? cstr("black") : cstr("red"),
            cur->base.color == RBTREE_BLACK ? cstr("white") : cstr("black")));

            if (cur->base.left != set->NIL)
                VecU8_append_vec(&graph, VecU8_fmt("    v%i -> v%i [label=left]\n",
                    cur->key, ((RBTreeNode_S64*)cur->base.left)->key));
            if (cur->base.right != set->NIL)
                VecU8_append_vec(&graph, VecU8_fmt("    v%i -> v%i [label=right]\n",
                    cur->key, ((RBTreeNode_S64*)cur->base.right)->key));
        }
        bfs.len = 0;
    }
    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);
}



void insert_only(){
    {
        RBTree_SetS64 set = RBTree_SetS64_new();
        check_correctness(&set);
        RBTree_SetS64_drop(set);
    }
    {
        RBTree_SetS64 set = RBTree_SetS64_new();
        for (S64 p = 0; p < 100; p++) {
            bool ret = RBTree_SetS64_insert(&set, 2 * p);
            check(ret);
            check_correctness(&set);
            // save_tree_to_file(&set, cstr("last_correct"));
            for (S64 i = 0; i <= p; i++) {
                RBTreeNode_S64* ret1 = RBTree_SetS64_find(&set, 2 * i);
                check(ret1 != NULL);
                RBTreeNode_S64* ret2 = RBTree_SetS64_find(&set, 2 * i + 1);
                check(ret2 == NULL);
            }
        }
        RBTree_SetS64_drop(set);
    }
    {
        RBTree_SetS64 set = RBTree_SetS64_new();
        for (S64 p = -99; p <= 0; p++) {
            bool ret = RBTree_SetS64_insert(&set, 2 * p);
            check(ret);
            check_correctness(&set);
            // save_tree_to_file(&set, cstr("last_correct"));
            for (S64 i = -99; i <= p; i++) {
                RBTreeNode_S64* ret1 = RBTree_SetS64_find(&set, 2 * i);
                check(ret1 != NULL);
                RBTreeNode_S64* ret2 = RBTree_SetS64_find(&set, 2 * i + 1);
                check(ret2 == NULL);
            }
        }
        RBTree_SetS64_drop(set);
    }
    for (int s = 0; s < 1000; s++) {
        srand(s);
        VecS64 have = VecS64_new();
        RBTree_SetS64 set = RBTree_SetS64_new();
        for (int m = 0; m < 1000; m++) {
            S64 x = (S64)rand();
            bool in = false;
            for (size_t j = 0; j < have.len; j++) {
                if (have.buf[j] == x)
                    in = true;
            }
            bool iret = RBTree_SetS64_insert(&set, x);
            if (in) {
                check(!iret);
            } else {
                check(iret);
                VecS64_append(&have, x);
            }
        }
        RBTree_SetS64_drop(set);
        VecS64_drop(have);
    }
}

void s_4_t_42(){
    RBTreeNode_S64* v0 = safe_malloc(sizeof(RBTreeNode_S64));
    RBTreeNode_S64* v10 = safe_malloc(sizeof(RBTreeNode_S64));
    RBTreeNode_S64* v20 = safe_malloc(sizeof(RBTreeNode_S64));
    RBTreeNode_S64* v30 = safe_malloc(sizeof(RBTreeNode_S64));
    RBTreeNode* NIL = safe_calloc(1, sizeof(RBTreeNode));
    *v0 = (RBTreeNode_S64){ .base.parent = &v20->base, .base.left = NIL, .base.right = &v10->base,
        .base.color = RBTREE_BLACK, .key = 0};
    *v10 = (RBTreeNode_S64){ .base.parent = &v0->base, .base.left = NIL, .base.right = NIL,
        .base.color = RBTREE_RED, .key = 10};
    *v20 = (RBTreeNode_S64){ .base.parent = NIL, .base.left = &v0->base, .base.right = &v30->base,
        .base.color = RBTREE_BLACK, .key = 20};
    *v30 = (RBTreeNode_S64){ .base.parent = &v20->base, .base.left = NIL, .base.right = NIL,
        .base.color = RBTREE_BLACK, .key = 30};
    RBTree_SetS64 set = {.root = &v20->base, .NIL = NIL};
    check_correctness(&set);
    // save_tree_to_file(&set, cstr("SOROK_ONE"));
    check(RBTree_SetS64_erase(&set, 20));
    // save_tree_to_file(&set, cstr("I_AM_LOOSING_MY_MIND"));
    check_correctness(&set);
    RBTree_SetS64_drop(set);
}

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, RBTreeNode_S64* my_ans){
    if (true_ans.variant == Option_None) {
        check(my_ans == NULL);
    } else {
        check(my_ans != NULL && my_ans->key == true_ans.some);
    }
}


void stress(){
    printf("Are you ready for real stress?\n");
    for (int s = 2; s < 1000; s++) {
        srand(s);
        RBTree_SetS64 set = RBTree_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(RBTree_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(RBTree_SetS64_erase(&set, v));
                check_correctness(&set);
            }
            // VecU8 name = VecU8_fmt("t_%u", (U64)t);
            // save_tree_to_file(&set, VecU8_to_span(&name));
            /* We did something */
            for (size_t j = 0; j < real_set.len; j++) {
                check(RBTree_SetS64_find(&set, real_set.buf[j]) != NULL);
            }
            for (size_t j = 0; j < complementary_set.len; j++) {
                check(RBTree_SetS64_find(&set, complementary_set.buf[j]) == NULL);
            }

            {   /* Tree detour forward */
                U64 count = 0;
                RBTreeNode_S64* cur = RBTree_SetS64_find_min(&set);
                while (cur) {
                    count++;
                    cur = RBTree_SetS64_find_next(&set, cur);
                }
                check(count == real_set.len);
            }
            {   /* Tree detour backward */
                U64 count = 0;
                RBTreeNode_S64* cur = RBTree_SetS64_find_max(&set);
                while (cur) {
                    count++;
                    cur = RBTree_SetS64_find_prev(&set, cur);
                }
                check(count == real_set.len);
            }
            /* 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);
                    RBTreeNode_S64* my_ans = RBTree_SetS64_find_max_less(&set, key);
                    check_option_int_equiv_node(true_ans, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_max_less_or_eq(&real_set, key);
                    RBTreeNode_S64* my_ans = RBTree_SetS64_find_max_less_or_eq(&set, key);
                    check_option_int_equiv_node(true_ans, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_min_grtr(&real_set, key);
                    RBTreeNode_S64* my_ans = RBTree_SetS64_find_min_grtr(&set, key);
                    check_option_int_equiv_node(true_ans, my_ans);
                }
                {
                    OptionS64 true_ans = VecS64_find_min_grtr_or_eq(&real_set, key);
                    RBTreeNode_S64* my_ans = RBTree_SetS64_find_min_grtr_or_eq(&set, key);
                    check_option_int_equiv_node(true_ans, my_ans);
                }
            }
        }
        VecS64_drop(real_set);
        VecS64_drop(complementary_set);
        RBTree_SetS64_drop(set);
        printf("Seed s=%d passed test\n", s);
    }
}

int main(){
    insert_only();
    stress();
    s_4_t_42();
    return 0;
}
