#include "unity.h"
#include "../src/network.h"

void setUp(void) {}

void tearDown(void) {}

void test_network_create(void) {
    int dimensions[] = {2, 3, 1};
    Network *network = network_create(3, dimensions);
    TEST_ASSERT_NOT_NULL(network);
    TEST_ASSERT_EQUAL_INT(3, network->layerCount);
    TEST_ASSERT_EQUAL_INT(2, network->dimensions[0]);
    TEST_ASSERT_EQUAL_INT(3, network->dimensions[1]);
    TEST_ASSERT_EQUAL_INT(1, network->dimensions[2]);
    TEST_ASSERT_NOT_NULL(network->weights);
    TEST_ASSERT_NOT_NULL(network->biases);
    TEST_ASSERT_NOT_NULL(network->outputs);
    for (int i = 0; i < 2; i++) {
        TEST_ASSERT_NOT_NULL(network->weights[i]);
    }
    network_free(network);
}

void test_network_predict(void) {
    // XOR function. Assume it already trained, could it predict the output?
    int dimensions[] = {2, 2, 1};
    Network *network = network_create(3, dimensions);

    // A correct (trained) weights and biases for the XOR function:
    double weights1[2][2] = {{20, 20}, {-20, -20}};
    double weights2[1][2] = {{20, 20}};
    double biases1[2][1] = {{-10}, {30}};
    double biases2[1][1] = {{-30}};

    matrix_initialise(network->weights[0], weights1);
    matrix_initialise(network->weights[1], weights2);
    matrix_initialise(network->biases[0], biases1);
    matrix_initialise(network->biases[1], biases2);

    // Inputs and targets for the XOR function
    double inputs[4][2] = {{0, 0}, {0, 1}, {1, 0}, {1, 1}};
    double targets[4][1] = {{0}, {1}, {1}, {0}};

    for (int i = 0; i < 4; i++) {
        Matrix *input = matrix_create(2, 1);
        matrix_initialise(input, (double[2][1]){inputs[i][0], inputs[i][1]});

        Matrix *target = matrix_create(1, 1);
        matrix_initialise(target, (double[1][1]){targets[i][0]});

        network_predict(network, input);
        // The result is correct if it is within 0.0001 of the target
        TEST_ASSERT_FLOAT_WITHIN(0.0001, targets[i][0], network_output(network)[0]);
        
        matrix_free(input);
        matrix_free(target);
    }
}

void check_matrix_equality(Matrix *m1, Matrix *m2) {
    TEST_ASSERT_EQUAL_INT(m1->rows, m2->rows);
    TEST_ASSERT_EQUAL_INT(m1->cols, m2->cols);
    for (int i = 0; i < m1->rows; i++) {
        for (int j = 0; j < m1->cols; j++) {
            TEST_ASSERT_EQUAL_DOUBLE(m1->data[i][j], m2->data[i][j]);
        }
    }
}

void test_network_loadstore(void) {
    int dimensions[] = {2, 3, 1};
    Network *network = network_create(3, dimensions);
    matrix_initialise(network->weights[0], (double[3][2]){{1, 2}, {3, 4}, {5, 6}});
    matrix_initialise(network->weights[1], (double[1][3]){{7, 8, 9}});
    matrix_initialise(network->biases[0], (double[3][1]){{1}, {2}, {3}});
    network_store(network, "my_network");
    Network *loaded_network = network_load("my_network");
    TEST_ASSERT_NOT_NULL(loaded_network);
    TEST_ASSERT_EQUAL_INT(3, loaded_network->layerCount);
    TEST_ASSERT_EQUAL_INT(2, loaded_network->dimensions[0]);
    TEST_ASSERT_EQUAL_INT(3, loaded_network->dimensions[1]);
    TEST_ASSERT_EQUAL_INT(1, loaded_network->dimensions[2]);
    TEST_ASSERT_NOT_NULL(loaded_network->weights);
    TEST_ASSERT_NOT_NULL(loaded_network->biases);
    TEST_ASSERT_NOT_NULL(loaded_network->outputs);
    for (int i = 0; i < 2; i++) {
        TEST_ASSERT_NOT_NULL(loaded_network->weights[i]);
        check_matrix_equality(network->weights[i], loaded_network->weights[i]);
    }
    for (int i = 0; i < 2; i++) {
        TEST_ASSERT_NOT_NULL(loaded_network->biases[i]);
        check_matrix_equality(network->biases[i], loaded_network->biases[i]);
    }
    network_free(network);
    network_free(loaded_network);

}

int main(void) {
    UNITY_BEGIN();
    RUN_TEST(test_network_create);
    RUN_TEST(test_network_predict);
    RUN_TEST(test_network_loadstore);
    return UNITY_END();
}