#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include "matrix.h"
#include "utils.h"

Matrix *matrix_alloc(size_t m, size_t n)
{
    double *xs = calloc(m * n,  sizeof(double));
    assert(xs != NULL && "Allocation of matrix elements failed");
    Matrix *mat = malloc(sizeof(Matrix));
    assert(mat != NULL && "Allocation of matrix failed");
    mat->xs = xs;
    mat->m = m;
    mat->n = n;
    return mat;
}

void matrix_free(Matrix *mat)
{
    free(mat->xs);
    free(mat);
}

void matrix_freen(Matrix **mats)
{
    Matrix *mat, **p;
    p = mats;
    while ((mat = *(p++)))
        matrix_free(mat);
    free(mats);
}

void matrix_print(const Matrix *mat)
{
    matrix_foreach_idx(mat, it, i, j)
    {
        printf("%lf ", *it);
        if (j == mat->n - 1 && i != mat->m - 1) printf("\n");
    }
    printf("\n");
}

/* Expected format of a 3x2 matrix [3 2;1 2;0 0]
 * NOTE: No spaces around ';'
 * NOTE: No proper error handling is done */
Matrix *matrix_from_str(char *str)
{
    assert(str != NULL && "Passed null string");
    char *row, *tmp_row, *entry;

    str = str_delete_at(str, 0);
    str = str_delete_at(str, strlen(str) - 1);

    size_t rows = str_count_occ(str, ';') + 1;
    size_t cols = (str_count_occ(str, ' ') + rows) / rows;

    Matrix *mat_tmp = matrix_alloc(rows, cols);

    for (size_t i = 0; (row = strsep(&str, ";")); ++i)
    {
        tmp_row = strdup(row); 
        for (size_t j = 0; (entry = strsep(&tmp_row, " ")); ++j)
            matrix_at(mat_tmp, i, j)= atof(entry);
    }

    free(tmp_row);
    return mat_tmp;
}

Matrix *matrix_from_arr(double arr[], size_t m, size_t n)
{
    Matrix *mat = matrix_alloc(m, n);

    for (size_t i = 0; i < m * n; ++i)
        mat->xs[i] = arr[i];
    
    return mat;
}

char *matrix_to_str(const Matrix *mat)
{
    assert(mat != NULL);
    return "TODO";
}

Matrix *matrix_id(size_t n)
{
    Matrix *tmp = matrix_alloc(n, n);

    for (size_t i = 0; i < n; ++i)
        matrix_at(tmp, i, i) = 1;

    return tmp;
}

Matrix *matrix_const(size_t m, size_t n, double x)
{
    Matrix *mat = matrix_alloc(m, n);
    matrix_foreach(mat, it) *it = x;
    return mat;
}

void matrix_const1(Matrix *mat, double x)
{
    matrix_foreach(mat, it) *it = x;
}

Matrix *matrix_copy(const Matrix *mat)
{
    Matrix *tmp = matrix_alloc(mat->m, mat->n);
    matrix_foreach_idx(tmp, it, i, j) *it = matrix_at(mat, i, j);
    return tmp;
}

void matrix_copy1(Matrix *A, const Matrix *B)
{
    matrix_foreach_idx(A, a, i, j) *a = matrix_at(B, i, j);
}

Matrix *matrix_swap(const Matrix *mat, MatrixSwapType t, size_t i, size_t j)
{
    Matrix *tmp = matrix_copy(mat);
    size_t lim = (t == MATRIX_SWAP_ROWS ? mat->n : mat->m);

    for (size_t k = 0; k < lim; ++k)
    {
        switch (t)
        {
            case MATRIX_SWAP_ROWS:
                matrix_at(tmp, i, k) = matrix_at(mat, j, k);
                matrix_at(tmp, j, k) = matrix_at(mat, i, k);
                break;
            case MATRIX_SWAP_COLS:
                matrix_at(tmp, k, i) = matrix_at(mat, k, j);
                matrix_at(tmp, k, j) = matrix_at(mat, k, i);
                break;
        }
    }

    return tmp;
}

void matrix_swap1(Matrix *mat, MatrixSwapType t, size_t i, size_t j)
{
    double tmp;
    size_t lim = (t == MATRIX_SWAP_ROWS ? mat->n : mat->m);

    for (size_t k = 0; k < lim; ++k)
    {
        switch (t)
        {
            case MATRIX_SWAP_ROWS:
                tmp = matrix_at(mat, i, k);
                matrix_at(mat, i, k) = matrix_at(mat, j, k);
                matrix_at(mat, j, k) = tmp;
                break;
            case MATRIX_SWAP_COLS:
                tmp = matrix_at(mat, k, i);
                matrix_at(mat, k, i) = matrix_at(mat, k, j);
                matrix_at(mat, k, j) = tmp;
                break;
        }
    }
}

double matrix_trace(const Matrix *mat)
{
    assert(matrix_is_square(mat) && "Trace is not defined for non-square matrices");
    double sum = 0;

    for (size_t i = 0; i < mat->n; ++i)
        sum += matrix_at(mat, i, i);

    return sum;
}

Matrix *matrix_transpose(const Matrix *mat)
{
    assert(mat != NULL && "Matrix points to null");
    assert(mat->xs != NULL && "Matrix elements point to null");

    Matrix *T = matrix_alloc(mat->n, mat->m);
    matrix_foreach_idx(T, t, i, j) *t = matrix_at(mat, j, i);
    return T;
}

Matrix *matrix_add(const Matrix *A, const Matrix *B)
{
    assert(A->m == B->m && A->n == B->n && "Dimension mismatch");

    Matrix *C = matrix_alloc(A->m, A->n);
    matrix_foreach_idx(C, c, i, j) *c = matrix_at(A, i, j) + matrix_at(B, i, j);

    return C;
}

void matrix_add1(Matrix *A, const Matrix *B)
{
    matrix_foreach_idx(A, a, i, j) *a += matrix_at(B, i, j);
}

Matrix *matrix_scale(const Matrix *A, double x)
{
    Matrix *tmp = matrix_alloc(A->m, A->n);
    matrix_foreach_idx(tmp, it, i, j) *it = matrix_at(A, i, j) * x;

    return tmp;
}

void matrix_scale1(Matrix *A, double x)
{
    matrix_foreach(A, a) *a *= x;
}

Matrix *matrix_sub(const Matrix *A, const Matrix *B)
{
    assert(A->m == B->m && A->n == B->n && "Dimension mismatch");

    Matrix *C = matrix_alloc(A->m, A->n);
    matrix_foreach_idx(C, c, i, j) *c = matrix_at(A, i, j) - matrix_at(B, i, j);

    return C;
}

void matrix_sub1(Matrix *A, const Matrix *B)
{
    matrix_foreach_idx(A, a, i, j) *a -= matrix_at(B, i, j);
}

Matrix *matrix_mult(const Matrix *A, const Matrix *B)
{
    assert(A->n == B->m && "Dimension mismatch");
    Matrix *C = matrix_alloc(A->m, B->n);

    matrix_foreach_idx(C, c, i, j) {
        for (size_t k = 0; k < A->n; ++k)
            *c += matrix_at(A, i, k) * matrix_at(B, k, j);
    }

    return C;
}

Matrix *matrix_rand(size_t m, size_t n, int bound_l, int bound_u, MatrixType type)
{
    Matrix *mat = matrix_alloc(m, n);
    matrix_rand1(mat, bound_l, bound_u, type);
    return mat;
}

void matrix_rand1(Matrix *mat, int bound_l, int bound_u, MatrixType type)
{
    assert(bound_l <= bound_u && "Lower bound is expected to be less or equal than upper bound");

    if (bound_l == bound_u)
        matrix_const1(mat, bound_l);

    bool zero;
    matrix_foreach_idx(mat, it, i, j)
    {
        zero =
            (type == MATRIX_DIAG && i != j)         ||
            (type == MATRIX_TRIAG_UPPER && i > j)   ||
            (type == MATRIX_TRIAG_LOWER && i < j)   ||
            (type == MATRIX_TRIAG_SUPPER && i >= j) ||
            (type == MATRIX_TRIAG_SLOWER && i <= j);

        *it = zero ? 0 : int_rand(bound_l, bound_u);
    }
}

bool matrix_eq(const Matrix *A, const Matrix *B, double tol)
{
    assert(A->m == B->m && A->n == B->n && "Dimension mismatch");
    matrix_foreach_idx(A, a, i, j)
            if (ABS(*a - matrix_at(B, i, j)) > tol) return false;
    return true;
}

double matrix_norm_frob(const Matrix *mat)
{
    return sqrt(matrix_trace(matrix_mult(matrix_transpose(mat), mat)));
}

Matrix **matrix_LR(const Matrix *A, const Matrix *b)
{
    assert(matrix_is_square(A) && "LR decomposition only works for square matrices");
    assert(A->n == b->m && "Dimension mismatch");
    assert(matrix_is_colvec(b) && "Column vector b expected");

    Matrix *L = matrix_id(A->n);
    Matrix *A_prev = matrix_copy(A);
    Matrix *A_curr = matrix_alloc(A->n, A->n);

    for (size_t k = 1; k < A->n; ++k)
    {
        for (size_t i = 0; i < k; ++i)
        {
            for (size_t j = 0; j < A->n; ++j)
                matrix_at(A_curr, i, j) = matrix_at(A_prev, i, j);
        }

        assert(matrix_at(A_prev, k - 1, k - 1) != 0 && "Pivoting strategy needed");

        for (size_t i = k; i < A->n; ++i)
        {
            matrix_at(L, i, k - 1) = matrix_at(A_prev, i, k - 1) / matrix_at(A_prev, k - 1, k - 1);
            for (size_t j = k; j < A->n; ++j)
            {
                matrix_at(A_curr, i, j) = matrix_at(A_prev, i, j) - matrix_at(L, i, k - 1) * matrix_at(A_prev, k - 1, j);
            }
        }

        A_prev = matrix_copy(A_curr);
        if (k != A->n - 1)
            A_curr = matrix_alloc(A->n, A->n);
    }
    
    matrix_free(A_prev);
    Matrix **result = malloc(sizeof(Matrix) * 2);
    result[0] = L;
    result[1] = A_curr;
    result[2] = NULL;
    return result;
}

Matrix *matrix_forwardel(const Matrix *L, const Matrix *b)
{
    assert(L->n == b->m && "Dimension mismatch");
    assert(matrix_is_colvec(b) && "Column vector b expected");

    double sum;
    Matrix *y = matrix_alloc(L->n, 1);

    for (size_t i = 0; i < y->m; ++i)
    {
        sum = matrix_at(b, i, 0);
        for (size_t j = 0; j < i; ++j)
            sum -= matrix_at(L, i, j) * matrix_at(y, j, 0);
        matrix_at(y, i, 0) = sum / matrix_at(L, i, i);
    }

    return y;
}

Matrix *matrix_backsubst(const Matrix *R, const Matrix *y)
{
    assert(R->n == y->m && "Dimension mismatch");
    assert(matrix_is_colvec(y) && "Column vector y expected");
    Matrix *x = matrix_alloc(R->n, 1);

    double sum;

    for (int i = x->m; i >= 0; --i)
    {
        sum = matrix_at(y, i, 0);
        for (size_t j = i + 1; j < x->m; ++j)
            sum -= matrix_at(R, i, j) * matrix_at(x, j, 0);
        matrix_at(x, i, 0) = sum / matrix_at(R, i, i);
    }

    return x;
}