#include "acutest.h"
#include "vieter_heap_internal.h"
#include <stdlib.h>

#define TEST_SIZE(heap, size) \
    TEST_CHECK(vieter_heap_size(heap) == size); \
    TEST_MSG("Size: %zu, expected: %lu", vieter_heap_size(heap), (uint64_t)size)

void test_init() {
    vieter_heap *heap = vieter_heap_init();
    TEST_CHECK(heap != NULL);
    TEST_SIZE(heap, 0);
    vieter_heap_free(heap);
}

void count_nodes(uint64_t *counter, vieter_heap_node *root) {
    (*counter)++;

    if (root->largest_order != NULL) {
        count_nodes(counter, root->largest_order);
    }

    // This will also traverse the various trees
    if (root->ptr.next_largest_order != NULL) {
        count_nodes(counter, root->ptr.next_largest_order);
    }
}

uint64_t count_nodes_heap(vieter_heap *heap) {
    uint64_t counter = 0;

    if (heap->tree != NULL) {
        count_nodes(&counter, heap->tree);
    }

    return counter;
}

void test_insert() {
    vieter_heap *heap = vieter_heap_init();
    TEST_SIZE(heap, 0);

    void *data;

    for (uint64_t i = 50; i > 0; i--) {
        vieter_heap_insert(heap, i, (void *)i);
        TEST_SIZE(heap, (uint64_t)51 - i);
        TEST_CHECK(count_nodes_heap(heap) == (uint64_t)51 - i);

        data = 0;

        TEST_CHECK(vieter_heap_peek(&data, heap) == vieter_heap_ok);
        TEST_CHECK_(data == (void *)i, "%lX == %lX", (uint64_t)data, i);
    }

    vieter_heap_free(heap);
}

void test_insert_random() {
    srand(1);

    vieter_heap *heap = vieter_heap_init();
    TEST_SIZE(heap, 0);

    uint64_t num = rand();
    uint64_t smallest = num;

    void *data = NULL;

    for (uint64_t i = 0; i < 5000; i++) {
        vieter_heap_insert(heap, num, (void *)num);
        TEST_SIZE(heap, i + 1);
        TEST_CHECK(count_nodes_heap(heap) == (uint64_t)i + 1);

        if (num < smallest) {
            smallest = num;
        }

        TEST_CHECK(vieter_heap_peek(&data, heap) == vieter_heap_ok);
        TEST_CHECK(data == (void *)smallest);

        data = NULL;

        num = rand();
    }

    vieter_heap_free(heap);
}

void test_pop() {
    const uint64_t n = 500;

    vieter_heap *heap = vieter_heap_init();
    TEST_SIZE(heap, 0);

    void *data;

    for (uint64_t i = n; i > 0; i--) {
        vieter_heap_insert(heap, i, (void *)i);
        TEST_SIZE(heap, (uint64_t)n + 1 - i);
        TEST_CHECK(count_nodes_heap(heap) == (uint64_t)n + 1 - i);

        TEST_CHECK(vieter_heap_peek(&data, heap) == vieter_heap_ok);
        TEST_CHECK(data == (void*)i);
    }

    data = NULL;

    for (uint64_t i = 1; i <= n; i++) {
        TEST_CHECK(vieter_heap_pop(&data, heap) == vieter_heap_ok);
        TEST_CHECK(data == (void*)i);
        TEST_SIZE(heap, (uint64_t)n - i);
    }

    vieter_heap_free(heap);
}

int uint64_t_compare(const void *a, const void *b) {
    if ((*(uint64_t *)a) < (*(uint64_t *)b)) {
        return -1;
    } else if ((*(uint64_t *)a) > (*(uint64_t *)b)) {
        return 1;
    } else {
        return 0;
    }
}

void test_pop_random() {
    const uint64_t n = 500;

    srand(0);

    vieter_heap *heap = vieter_heap_init();

    uint64_t *numbers = malloc(n * sizeof(uint64_t));
    uint64_t num;

    for (uint64_t i = 0; i < n; i++) {
        num = rand();
        vieter_heap_insert(heap, num, (void *)num);
        TEST_SIZE(heap, i + 1);
        TEST_CHECK(count_nodes_heap(heap) == i + 1);

        numbers[i] = num;
    }


    qsort(numbers, n, sizeof(uint64_t), uint64_t_compare);

    void *data = NULL;

    for (uint64_t i = 0; i < n; i++) {
        TEST_CHECK(vieter_heap_peek(&data, heap) == vieter_heap_ok);
        TEST_CHECK_(data == (void *)numbers[i], "peek %lx == %lx", (uint64_t)data, numbers[i]);

        data = NULL;

        TEST_CHECK(vieter_heap_pop(&data, heap) == vieter_heap_ok);
        TEST_CHECK_(data == (void *)numbers[i], "pop %lx == %lx", (uint64_t)data, numbers[i]);
        TEST_SIZE(heap, n - i - 1);
        TEST_CHECK(count_nodes_heap(heap) == n - i - 1);

        // Assure each size is also a valid heap after inserting
        vieter_heap_insert(heap, numbers[i], (void *)numbers[i]);
        TEST_SIZE(heap, n - i);
        TEST_CHECK(count_nodes_heap(heap) == n - i);

        data = NULL;

        TEST_CHECK(vieter_heap_pop(&data, heap) == vieter_heap_ok);
        TEST_CHECK_(data == (void *)numbers[i], "pop %lx == %lx", (uint64_t)data, numbers[i]);
        TEST_SIZE(heap, n - i - 1);
        TEST_CHECK(count_nodes_heap(heap) == n - i - 1);
    }

    vieter_heap_free(heap);
    free(numbers);
}

TEST_LIST = {
    {"heap init", test_init},
    {"heap insert", test_insert},
    {"heap insert random", test_insert_random},
    {"heap pop", test_pop},
    {"heap pop random", test_pop_random},
    {NULL, NULL}
};