#include <cstdint>
#include <cstdio>
#include <cstdlib>

// set `IRRT_DONT_TYPEDEF_INTS` because `cstdint` has it all
#define IRRT_DONT_TYPEDEF_INTS
#include "irrt_everything.hpp"

void test_fail() {
    printf("[!] Test failed\n");
    exit(1);
}

void __begin_test(const char* function_name, const char* file, int line) {
    printf("######### Running %s @ %s:%d\n", function_name, file, line);
}

#define BEGIN_TEST() __begin_test(__FUNCTION__, __FILE__, __LINE__)

template <typename T>
bool arrays_match(int len, T *as, T *bs) {
    for (int i = 0; i < len; i++) {
        if (as[i] != bs[i]) return false;
    }
    return true;
}

template <typename T>
void debug_print_array(const char* format, int len, T* as) {
    printf("[");
    for (int i = 0; i < len; i++) {
        if (i != 0) printf(", ");
        printf(format, as[i]);
    }
    printf("]");
}

template <typename T>
void assert_arrays_match(const char* label, const char* format, int len, T* expected, T* got) {
    if (!arrays_match(len, expected, got)) {
        printf("expected %s: ", label);
        debug_print_array(format, len, expected);
        printf("\n");
        printf("got %s: ", label);
        debug_print_array(format, len, got);
        printf("\n");
        test_fail();
    }
}

template <typename T>
void assert_values_match(const char* label, const char* format, T expected, T got) {
    if (expected != got) {
        printf("expected %s: ", label);
        printf(format, expected);
        printf("\n");
        printf("got %s: ", label);
        printf(format, got);
        printf("\n");
        test_fail();
    }
}

void test_calc_size_from_shape_normal() {
    // Test shapes with normal values
    BEGIN_TEST();

    int32_t shape[4] = { 2, 3, 5, 7 };
    debug_print_array("%d", 4, shape);
    assert_values_match("size", "%d", 210, ndarray_util::calc_size_from_shape<int32_t>(4, shape));
}

void test_calc_size_from_shape_has_zero() {
    // Test shapes with 0 in them
    BEGIN_TEST();

    int32_t shape[4] = { 2, 0, 5, 7 };
    assert_values_match("size", "%d", 0, ndarray_util::calc_size_from_shape<int32_t>(4, shape));
}

void test_set_strides_by_shape() {
    // Test `set_strides_by_shape()`
    BEGIN_TEST();

    int32_t shape[4] = { 99, 3, 5, 7 };
    int32_t strides[4] = { 0 };
    ndarray_util::set_strides_by_shape(4, strides, shape);

    int32_t expected_strides[4] = { 105, 35, 7, 1 };
    assert_arrays_match("strides", "%u", 4u, expected_strides, strides);
}

void test_ndarray_indices_iter_normal() {
    // Test NDArrayIndicesIter normal behavior
    BEGIN_TEST();

    int32_t shape[3] = { 1, 2, 3 };
    int32_t indices[3] = { 0, 0, 0 };
    auto iter = NDArrayIndicesIter<int32_t> {
        .ndims = 3u,
        .shape = shape,
        .indices = indices
    };

    assert_arrays_match("indices #0", "%u", 3u, iter.indices, (int32_t[3]) { 0, 0, 0 });
    iter.next();
    assert_arrays_match("indices #1", "%u", 3u, iter.indices, (int32_t[3]) { 0, 0, 1 });
    iter.next();
    assert_arrays_match("indices #2", "%u", 3u, iter.indices, (int32_t[3]) { 0, 0, 2 });
    iter.next();
    assert_arrays_match("indices #3", "%u", 3u, iter.indices, (int32_t[3]) { 0, 1, 0 });
    iter.next();
    assert_arrays_match("indices #4", "%u", 3u, iter.indices, (int32_t[3]) { 0, 1, 1 });
    iter.next();
    assert_arrays_match("indices #5", "%u", 3u, iter.indices, (int32_t[3]) { 0, 1, 2 });
    iter.next();
    assert_arrays_match("indices #6", "%u", 3u, iter.indices, (int32_t[3]) { 0, 0, 0 }); // Loops back
    iter.next();
    assert_arrays_match("indices #7", "%u", 3u, iter.indices, (int32_t[3]) { 0, 0, 1 });
}

void test_ndarray_fill_generic() {
    // Test ndarray fill_generic
    BEGIN_TEST();

    // Choose a type that's neither int32_t nor uint64_t (candidates of SizeT) to spice it up
    // Also make all the octets non-zero, to see if `memcpy` in `fill_generic` is working perfectly.
    uint16_t fill_value = 0xFACE;

    uint16_t in_data[6] = { 100, 101, 102, 103, 104, 105 }; // Fill `data` with values that != `999`
    int32_t in_itemsize = sizeof(uint16_t);
    const int32_t in_ndims = 2;
    int32_t in_shape[in_ndims] = { 2, 3 };
    int32_t in_strides[in_ndims] = {};
    NDArray<int32_t> ndarray = {
        .data = (uint8_t*) in_data,
        .itemsize = in_itemsize,
        .ndims = in_ndims,
        .shape = in_shape,
        .strides = in_strides,
    };
    ndarray.set_strides_by_shape();
    ndarray.fill_generic((uint8_t*) &fill_value); // `fill_generic` here

    uint16_t expected_data[6] = { fill_value, fill_value, fill_value, fill_value, fill_value, fill_value };
    assert_arrays_match("data", "0x%hX", 6, expected_data, in_data);
}

void test_ndarray_set_to_eye() {
    // Test `set_to_eye` behavior (helper function to implement `np.eye()`)
    BEGIN_TEST();

    double in_data[9] = { 99.0, 99.0, 99.0, 99.0, 99.0, 99.0, 99.0, 99.0, 99.0 };
    int32_t in_itemsize = sizeof(double);
    const int32_t in_ndims = 2;
    int32_t in_shape[in_ndims] = { 3, 3 };
    int32_t in_strides[in_ndims] = {};
    NDArray<int32_t> ndarray = {
        .data = (uint8_t*) in_data,
        .itemsize = in_itemsize,
        .ndims = in_ndims,
        .shape = in_shape,
        .strides = in_strides,
    };
    ndarray.set_strides_by_shape();

    double zero = 0.0;
    double one = 1.0;
    ndarray.set_to_eye(1, (uint8_t*) &zero, (uint8_t*) &one);

    assert_values_match("in_data[0]", "%f", 0.0, in_data[0]);
    assert_values_match("in_data[1]", "%f", 1.0, in_data[1]);
    assert_values_match("in_data[2]", "%f", 0.0, in_data[2]);
    assert_values_match("in_data[3]", "%f", 0.0, in_data[3]);
    assert_values_match("in_data[4]", "%f", 0.0, in_data[4]);
    assert_values_match("in_data[5]", "%f", 1.0, in_data[5]);
    assert_values_match("in_data[6]", "%f", 0.0, in_data[6]);
    assert_values_match("in_data[7]", "%f", 0.0, in_data[7]);
    assert_values_match("in_data[8]", "%f", 0.0, in_data[8]);
}

void test_slice_1() {
    // Test `slice(5, None, None).indices(100) == slice(5, 100, 1)`
    BEGIN_TEST();

    UserSlice<int> user_slice = {
        .start_defined = 1,
        .start = 5,
        .stop_defined = 0,
        .step_defined = 0,
    };

    auto slice = user_slice.indices(100);
    assert_values_match("start", "%d", 5, slice.start);
    assert_values_match("stop", "%d", 100, slice.stop);
    assert_values_match("step", "%d", 1, slice.step);
}

void test_slice_2() {
    // Test `slice(400, 999, None).indices(100) == slice(100, 100, 1)`
    BEGIN_TEST();

    UserSlice<int> user_slice = {
        .start_defined = 1,
        .start = 400,
        .stop_defined = 0,
        .step_defined = 0,
    };

    auto slice = user_slice.indices(100);
    assert_values_match("start", "%d", 100, slice.start);
    assert_values_match("stop", "%d", 100, slice.stop);
    assert_values_match("step", "%d", 1, slice.step);
}

void test_slice_3() {
    // Test `slice(-10, -5, None).indices(100) == slice(90, 95, 1)`
    BEGIN_TEST();

    UserSlice<int> user_slice = {
        .start_defined = 1,
        .start = -10,
        .stop_defined = 1,
        .stop = -5,
        .step_defined = 0,
    };

    auto slice = user_slice.indices(100);
    assert_values_match("start", "%d", 90, slice.start);
    assert_values_match("stop", "%d", 95, slice.stop);
    assert_values_match("step", "%d", 1, slice.step);
}

void test_slice_4() {
    // Test `slice(None, None, -5).indices(100) == (99, -1, -5)`
    BEGIN_TEST();

    UserSlice<int> user_slice = {
        .start_defined = 0,
        .stop_defined = 0,
        .step_defined = 1,
        .step = -5
    };

    auto slice = user_slice.indices(100);
    assert_values_match("start", "%d", 99, slice.start);
    assert_values_match("stop", "%d", -1, slice.stop);
    assert_values_match("step", "%d", -5, slice.step);
}

int main() {
    test_calc_size_from_shape_normal();
    test_calc_size_from_shape_has_zero();
    test_set_strides_by_shape();
    test_ndarray_indices_iter_normal();
    test_ndarray_fill_generic();
    test_ndarray_set_to_eye();
    test_slice_1();
    test_slice_2();
    test_slice_3();
    test_slice_4();
    return 0;
}