#include <stdlib.h>
#include <string.h>
#include "matrix.h"
#include <stdio.h>

Matrix createMatrix(unsigned int rows, unsigned int cols)
{
    Matrix mat = {.rows = rows, .cols = cols};
    // If one dimension is 0, return both dimensions as 0 and don't init the array/buffer.
    if (rows == 0 || cols == 0)
    {
        mat.rows = 0;
        mat.cols = 0;
        return mat;
    }

    // allocate contiguous and 0 initialized memory
    mat.buffer = calloc(rows * cols, sizeof(MatrixType));

    // check if calloc failed
    if (mat.buffer == NULL)
    {
        clearMatrix(&mat);
        perror("could not allocate memory");
    }

    return mat;
}

// reduce the dimensions to (0, 0) and free the memory
void clearMatrix(Matrix *matrix)
{
    free(matrix->buffer);
    matrix->buffer = NULL;
    matrix->cols = 0;
    matrix->rows = 0;
}

void setMatrixAt(MatrixType value, Matrix matrix, unsigned int rowIdx, unsigned int colIdx)
{
    // do nothing if idx is not in array or matrix buffer is NULL
    if (!(rowIdx < matrix.rows) || !(colIdx < matrix.cols) || matrix.buffer == NULL)
    {
        return;
    }

    matrix.buffer[rowIdx * matrix.cols + colIdx] = value;
}

MatrixType getMatrixAt(const Matrix matrix, unsigned int rowIdx, unsigned int colIdx)
{
    // return UNDEFINED_MATRIX_VALUE if idx is not in array or matrix buffer is NULL
    if (!(rowIdx < matrix.rows) || !(colIdx < matrix.cols) || matrix.buffer == NULL)
    {
        return UNDEFINED_MATRIX_VALUE;
    }

    return matrix.buffer[rowIdx * matrix.cols + colIdx];
};

Matrix add(const Matrix matrix1, const Matrix matrix2)
{
    Matrix resMat = (matrix1.cols > matrix2.cols) ? createMatrix(matrix1.rows, matrix1.cols) : createMatrix(matrix2.rows, matrix2.cols);

    if (resMat.buffer == NULL)
    {
        return createMatrix(0, 0);
    }

    if (matrix1.cols != matrix2.cols)
    {
        if (matrix1.rows != matrix2.rows)
        {
            clearMatrix(&resMat);
            return resMat;
        }
        else if (matrix1.cols == 1)
        {
            // broadcast vector
            for (size_t m = 0; m < matrix2.rows; m++)
            {
                for (size_t n = 0; n < matrix2.cols; n++)
                {
                    setMatrixAt(getMatrixAt(matrix2, m, n) + getMatrixAt(matrix1, m, 0), resMat, m, n);
                }
            }
            return resMat;
        }
        else if (matrix2.cols == 1)
        {
            // broadcast vector
            for (size_t m = 0; m < matrix1.rows; m++)
            {
                for (size_t n = 0; n < matrix1.cols; n++)
                {
                    setMatrixAt(getMatrixAt(matrix1, m, n) + getMatrixAt(matrix2, m, 0), resMat, m, n);
                }
            }
            return resMat;
        }
        else
        {
            clearMatrix(&resMat);
            return resMat;
        }
    }

    for (size_t m = 0; m < matrix1.rows; m++)
    {
        for (size_t n = 0; n < matrix1.cols; n++)
        {
            setMatrixAt(getMatrixAt(matrix1, m, n) + getMatrixAt(matrix2, m, n), resMat, m, n);
        }
    }

    return resMat;
}

Matrix multiply(const Matrix A, const Matrix B)
{
    if (A.cols != B.rows || A.buffer == NULL || B.buffer == NULL)
    {
        return createMatrix(0, 0);
    }

    int rows = A.rows, cols = B.cols;
    Matrix C = createMatrix(rows, cols);

    if (C.buffer == NULL)
    {
        return createMatrix(0, 0);
    }

    // M = Rows, K = Common Dim, N = Cols
    size_t M = A.rows, K = A.cols, N = B.cols;

    for (size_t i = 0; i < M; i++)
    {
        for (size_t k = 0; k < K; k++)

        {
            MatrixType valA = A.buffer[i * K + k];
            for (size_t j = 0; j < N; j++)
            {
                // C[i, j] += A[i, k] * B[k, j];
                // M x N,     M x K,    K x N
                C.buffer[i * N + j] += valA * B.buffer[k * N + j];
            }
        }
    }
    return C;
}