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

void setUp(void) {
    // set stuff up here
}

void tearDown(void) {
    // clean stuff up here
}

void test_matrix_create(void) {
    Matrix *m = matrix_create(3, 3);
    TEST_ASSERT_NOT_NULL(m);
    TEST_ASSERT_EQUAL_INT(3, m->rows);
    TEST_ASSERT_EQUAL_INT(3, m->cols);
    matrix_free(m);
}

void test_matrix_multiply(void) {
    Matrix *m1 = matrix_create(2, 3);
    Matrix *m2 = matrix_create(3, 2);
    
    double m1_data[2][3] = {
        {1, 2, 3},
        {4, 5, 6}
    };
    double m2_data[3][2] = {
        {7, 8},
        {9, 10},
        {11, 12}
    };
    matrix_initialise(m1, m1_data);
    matrix_initialise(m2, m2_data);

    Matrix *m3 = matrix_multiply(m1, m2);
    
    TEST_ASSERT_NOT_NULL(m3);

    TEST_ASSERT_EQUAL_INT(2, m3->rows);
    TEST_ASSERT_EQUAL_INT(2, m3->cols);

    TEST_ASSERT_EQUAL_FLOAT(58, m3->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(64, m3->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(139, m3->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(154, m3->data[1][1]);

    matrix_free(m1);
    matrix_free(m2);
    matrix_free(m3);
}

void test_matrix_add(void) {
    Matrix *m1 = matrix_create(2, 2);
    Matrix *m2 = matrix_create(2, 2);
    
    double m1_data[2][2] = {
        {1, 2},
        {3, 4}
    };
    double m2_data[2][2] = {
        {5, 6},
        {7, 8}
    };
    matrix_initialise(m1, m1_data);
    matrix_initialise(m2, m2_data);

    Matrix *m3 = matrix_add(m1, m2);
    
    TEST_ASSERT_NOT_NULL(m3);

    TEST_ASSERT_EQUAL_INT(2, m3->rows);
    TEST_ASSERT_EQUAL_INT(2, m3->cols);

    TEST_ASSERT_EQUAL_FLOAT(6, m3->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(8, m3->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(10, m3->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(12, m3->data[1][1]);

    matrix_free(m1);
    matrix_free(m2);
    matrix_free(m3);
}

void test_matrix_subtract(void) {
    Matrix *m1 = matrix_create(2, 2);
    Matrix *m2 = matrix_create(2, 2);
    
    double m1_data[2][2] = {
        {1, 2},
        {3, 4}
    };
    double m2_data[2][2] = {
        {5, 6},
        {7, 8}
    };
    matrix_initialise(m1, m1_data);
    matrix_initialise(m2, m2_data);

    Matrix *m3 = matrix_subtract(m1, m2);
    
    TEST_ASSERT_NOT_NULL(m3);

    TEST_ASSERT_EQUAL_INT(2, m3->rows);
    TEST_ASSERT_EQUAL_INT(2, m3->cols);

    TEST_ASSERT_EQUAL_FLOAT(-4, m3->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(-4, m3->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(-4, m3->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(-4, m3->data[1][1]);

    matrix_free(m1);
    matrix_free(m2);
    matrix_free(m3);
}

void test_matrix_print(void) {
    Matrix *m = matrix_create(3, 3);
    for (unsigned int i = 0; i < m->rows; i++) {
        for (unsigned int j = 0; j < m->cols; j++) {
            m->data[i][j] = 100 * i * m->cols + j;
        }
    }
    matrix_print(m, stdout);
    matrix_free(m);
}

static double alwaysZeroTransform(int i, int j, double val) {
    return 0;
}

static double customTransform(int i, int j, double val) {
    return 3 * i + j;
}

void test_matrix_transform(void) {
    Matrix *m = matrix_create(3, 3);
    for (unsigned int i = 0; i < m->rows; i++) {
        for (unsigned int j = 0; j < m->cols; j++) {
            m->data[i][j] = 100 * i * m->cols + j;
        }
    }
    matrix_transform(m, alwaysZeroTransform);
    for (unsigned int i = 0; i < m->rows; i++) {
        for (unsigned int j = 0; j < m->cols; j++) {
            TEST_ASSERT_EQUAL_FLOAT(0, m->data[i][j]);
        }
    }
    matrix_transform(m, customTransform);
    for (unsigned int i = 0; i < m->rows; i++) {
        for (unsigned int j = 0; j < m->cols; j++) {
            TEST_ASSERT_EQUAL_FLOAT(3 * i + j, m->data[i][j]);
        }
    }
    matrix_free(m);
}

void test_matrix_transpose(void) {
    Matrix *m = matrix_create(2, 3);
    
    double m_data[2][3] = {
        {1, 2, 3},
        {4, 5, 6}
    };
    matrix_initialise(m, m_data);

    Matrix *mt = matrix_transpose(m);
    
    TEST_ASSERT_NOT_NULL(mt);
    TEST_ASSERT_EQUAL_INT(3, mt->rows);
    TEST_ASSERT_EQUAL_INT(2, mt->cols);

    TEST_ASSERT_EQUAL_FLOAT(1, mt->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(4, mt->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(2, mt->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(5, mt->data[1][1]);
    TEST_ASSERT_EQUAL_FLOAT(3, mt->data[2][0]);
    TEST_ASSERT_EQUAL_FLOAT(6, mt->data[2][1]);

    matrix_free(m);
    matrix_free(mt);
}

void test_matrix_dot(void) {
    Matrix *m1 = matrix_create(2, 2);
    Matrix *m2 = matrix_create(2, 2);
    
    double m1_data[2][2] = {
        {1, 2},
        {3, 4}
    };
    double m2_data[2][2] = {
        {2, 0.5},
        {1, 2}
    };
    matrix_initialise(m1, m1_data);
    matrix_initialise(m2, m2_data);

    Matrix *m3 = matrix_dot(m1, m2);
    
    TEST_ASSERT_NOT_NULL(m3);

    TEST_ASSERT_EQUAL_INT(2, m3->rows);
    TEST_ASSERT_EQUAL_INT(2, m3->cols);

    TEST_ASSERT_EQUAL_FLOAT(2, m3->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(1, m3->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(3, m3->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(8, m3->data[1][1]);

    matrix_free(m1);
    matrix_free(m2);
    matrix_free(m3);
}


void test_matrix_scale(void) {
    Matrix *m = matrix_create(2, 2);
    
    double m_data[2][2] = {
        {1, 2},
        {3, 4}
    };
    matrix_initialise(m, m_data);

    matrix_scale(m, 2.5);
    
    TEST_ASSERT_EQUAL_FLOAT(2.5, m->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(5.0, m->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(7.5, m->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(10.0, m->data[1][1]);

    matrix_free(m);
}

void test_matrix_multiplyp(void) {
    Matrix *m1 = matrix_create(2, 2);
    Matrix *m2 = matrix_create(2, 2);
    
    double m1_data[2][2] = {
        {1, 2},
        {3, 4}
    };

    double m2_data[2][2] = {
        {2, 0.5},
        {1, 2}
    };
    matrix_initialise(m1, m1_data);
    matrix_initialise(m2, m2_data);

    matrix_multiplyp(m1, m2);

    TEST_ASSERT_EQUAL_FLOAT(4, m1->data[0][0]);
    TEST_ASSERT_EQUAL_FLOAT(4.5, m1->data[0][1]);
    TEST_ASSERT_EQUAL_FLOAT(10, m1->data[1][0]);
    TEST_ASSERT_EQUAL_FLOAT(9.5, m1->data[1][1]);

    matrix_free(m1);
    matrix_free(m2);
}

int main(void) {
    UNITY_BEGIN();
    
    RUN_TEST(test_matrix_create);
    RUN_TEST(test_matrix_print);
    RUN_TEST(test_matrix_multiply);
    RUN_TEST(test_matrix_add);
    RUN_TEST(test_matrix_subtract);
    RUN_TEST(test_matrix_transform);
    RUN_TEST(test_matrix_transpose);
    RUN_TEST(test_matrix_dot);
    RUN_TEST(test_matrix_scale);
    RUN_TEST(test_matrix_multiplyp);

    return UNITY_END();
}