#include "acutest.h"
#include "vieter_tree_internal.h"
#include "vieter_tree_node.h"
#include "vieter_tree_balancing.h"

// This uses the examples from wikipedia
// https://en.wikipedia.org/wiki/Tree_rotation

void test_rotate_right() {
    vieter_tree_node *a = vieter_tree_node_init();
    vieter_tree_node *b = vieter_tree_node_init();
    vieter_tree_node *p = vieter_tree_node_init();

    vieter_tree_node_set_child(p, a, false);
    vieter_tree_node_set_child(p, b, true);

    vieter_tree_node *c = vieter_tree_node_init();
    vieter_tree_node *q = vieter_tree_node_init();

    vieter_tree_node_set_child(q, p, false);
    vieter_tree_node_set_child(q, c, true);

    vieter_tree_node *new_root = vieter_tree_node_rotate(q, true);

    TEST_CHECK(new_root == p);
    TEST_CHECK(new_root->children[0] == a);
    TEST_CHECK(new_root->children[1] == q);
    TEST_CHECK(new_root->children[1]->children[0] == b);
    TEST_CHECK(new_root->children[1]->children[1] == c);

    vieter_tree_node_free(a);
    vieter_tree_node_free(b);
    vieter_tree_node_free(p);
    vieter_tree_node_free(c);
    vieter_tree_node_free(q);
}

void test_rotate_left() {
    vieter_tree_node *b = vieter_tree_node_init();
    vieter_tree_node *c = vieter_tree_node_init();
    vieter_tree_node *q = vieter_tree_node_init();

    vieter_tree_node_set_child(q, b, false);
    vieter_tree_node_set_child(q, c, true);

    vieter_tree_node *a = vieter_tree_node_init();
    vieter_tree_node *p = vieter_tree_node_init();

    vieter_tree_node_set_child(p, a, false);
    vieter_tree_node_set_child(p, q, true);

    vieter_tree_node *new_root = vieter_tree_node_rotate(p, false);

    TEST_CHECK(new_root == q);
    TEST_CHECK(new_root->children[0] == p);
    TEST_CHECK(new_root->children[1] == c);
    TEST_CHECK(new_root->children[0]->children[0] == a);
    TEST_CHECK(new_root->children[0]->children[1] == b);

    vieter_tree_node_free(a);
    vieter_tree_node_free(b);
    vieter_tree_node_free(p);
    vieter_tree_node_free(c);
    vieter_tree_node_free(q);
}

TEST_LIST = {
    {"tree rotate right", test_rotate_right},
    {"tree rotate left", test_rotate_left},
    {NULL, NULL}
};