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

#include "clm.h"

const clm_Matrix INVALID_MATRIX = {.rows = 0, .cols = 0, .values = NULL};
const clm_Vector INVALID_VECTOR = {.length = 0, .values = NULL};

clm_Matrix clm_createMatrix(unsigned int rows, unsigned int cols) {
	printf("CREATING MATRIX\n");
	clm_Matrix mat;
	mat.rows = rows;
	mat.cols = cols;
	mat.values = calloc(rows * cols, sizeof(float));
	mat.transposed = false;
	return mat;
}

clm_Matrix clm_createMatrixRandom(unsigned int rows, unsigned int cols) {
	clm_Matrix mat = clm_createMatrix(rows, cols);

	for(unsigned int i = 0; i < rows; i++) {
		for(unsigned int j = 0; j < cols; j++) {
			mat.values[i * cols + j] = ((float) rand() / RAND_MAX) * 2.0f - 1.0f;
		}
	}

	return mat;
}

clm_Matrix clm_matrixZero(clm_Matrix mat) {
	for(unsigned int i = 0; i < mat.rows * mat.cols; i++) {
		mat.values[i] = 0;
	}

	return mat;
}

void clm_freeMatrix(clm_Matrix mat) {
	free(mat.values);
}

clm_Matrix clm_matrixCopy(clm_Matrix mat) {
	clm_Matrix copy = clm_createMatrix(mat.rows, mat.cols);

	memcpy(copy.values, mat.values, mat.rows * mat.cols * sizeof(float));

	return copy;
}

clm_Matrix clm_matrixAddScalar(clm_Matrix mat, float scalar) {
	for(unsigned int i = 0; i < mat.cols * mat.rows; i++) {
		mat.values[i] += scalar;
	}

	return mat;
}

clm_Matrix clm_matrixAddMatrix(clm_Matrix mat, clm_Matrix other) {
	if(mat.cols != other.cols || mat.rows != other.rows) {
		printf("Failed to add matrices (got %dx%d and %dx%d)\n", mat.cols, mat.rows, other.cols, other.rows);
		return INVALID_MATRIX;
	}

	for(unsigned int i = 0; i < mat.cols * mat.rows; i++) {
		mat.values[i] += other.values[i];
	}

	return mat;
}

clm_Matrix clm_matrixSubtractMatrix(clm_Matrix mat, clm_Matrix other) {
	if(mat.cols != other.cols || mat.rows != other.rows) {
		printf("Failed to sub matrices\n");
		return INVALID_MATRIX;
	}

	for(unsigned int i = 0; i < mat.rows * mat.cols; i++) {
		mat.values[i] -= other.values[i];
	}

	return mat;
}

clm_Matrix clm_matrixTranspose(clm_Matrix mat) {
	clm_Matrix tr;
	tr.cols = mat.rows;
	tr.rows = mat.cols;
	tr.values = mat.values;
	tr.transposed = !mat.transposed;
	return tr;
}

clm_Matrix clm_matrixMultiplyMatrix(clm_Matrix a, clm_Matrix b, clm_Matrix out) {
	if(a.cols != b.rows) {
		printf("Cannot multiply matrices (got %dx%d and %dx%d)\n", a.rows, a.cols, b.rows, b.cols);
		return INVALID_MATRIX;
	}

	if(out.rows != a.rows || out.cols != b.cols) {
		printf("Cannot multiply matrices: output invalid shape (expected %dx%d, got %dx%d)\n", a.rows, b.cols, out.rows, out.cols);
		return INVALID_MATRIX;
	}

	for(unsigned int i = 0; i < out.rows; i++) {
		for(unsigned int j = 0; j < out.cols; j++) {
			float sum = 0;
			for(unsigned int k = 0; k < a.cols; k++) {
				sum += a.values[i * a.cols + k] * b.values[k * b.cols + j];
			}
			out.values[i * out.cols + j] = sum;
		}
	}

	return out;
}

clm_Matrix clm_matrixMultiplyMatrixALLOC(clm_Matrix a, clm_Matrix b) {
	clm_Matrix out = clm_createMatrix(a.rows, b.cols);
	clm_matrixMultiplyMatrix(a, b, out);
	return out;
}

clm_Matrix clm_matrixMultiplyMatrixElements(clm_Matrix mat, clm_Matrix other) {
	if(mat.rows != other.rows || mat.cols != other.cols) {
		printf("Cannot multiply matrices element-wise\n");
		return INVALID_MATRIX;
	}

	for(unsigned int i = 0; i < mat.cols * mat.rows; i++) {
		mat.values[i] *= other.values[i];
	}

	return mat;
}

clm_Matrix clm_matrixMultiplyScalar(clm_Matrix mat, float scalar) {
	for(unsigned int i = 0; i < mat.cols * mat.rows; i++) {
		mat.values[i] *= scalar;
	}

	return mat;
}

clm_Matrix clm_matrixSigmoid(clm_Matrix mat) {
	for(unsigned int i = 0; i < mat.rows; i++) {
		for(unsigned int j = 0; j < mat.cols; j++) {
			matrixAt(mat, i, j) = 1 / (1 + exp(-matrixAt(mat, i, j)));
		}
	}

	return mat;
}

clm_Matrix clm_matrixDSigmoid(clm_Matrix mat) {
	for(unsigned int i = 0; i < mat.rows; i++) {
		for(unsigned int j = 0; j < mat.cols; j++) {
			float v = matrixAt(mat, i, j);
			matrixAt(mat, i, j) = v * (1 - v);
		}
	}

	return mat;
}

clm_Matrix clm_matrixFromArray(float *array, unsigned int length) {
	clm_Matrix matrix = clm_createMatrix(length, 1);
	memcpy(matrix.values, array, length * sizeof(float));
	return matrix;
}

clm_Matrix clm_matrixWrapArray(float *array, unsigned int length) {
	clm_Matrix mat;
	mat.rows = length;
	mat.cols = 1;
	mat.values = array;
	return mat;
}

bool clm_matrixIsInvalid(clm_Matrix mat) {
	return mat.values == NULL;
}

clm_Vector clm_vectorCreate(unsigned int length) {
	clm_Vector vector;
	vector.length = length;
	vector.values = calloc(length, sizeof(float));
	return vector;
}

bool clm_vectorIsInvalid(clm_Vector vec) {
	return vec.values != NULL;
}

clm_Linear clm_linearCreateRandom(unsigned int inputs, unsigned int outputs) {
	clm_Linear linear;
	linear.weights = clm_createMatrixRandom(outputs, inputs);
	linear.bias = clm_createMatrixRandom(outputs, 1);

	linear.output = clm_createMatrix(outputs, 1);
	linear.error = clm_createMatrix(outputs, 1);
	linear.weightsError = clm_createMatrix(outputs, inputs);
	return linear;
}

void clm_freeVector(clm_Vector vector) {
	free(vector.values);
}

void clm_freeLinear(clm_Linear linear) {
	clm_freeMatrix(linear.weights);
	clm_freeMatrix(linear.bias);
}

void clm_matrixPrint(clm_Matrix mat) {
	printf("[\n");
	for(unsigned int i = 0; i < mat.rows; i++) {
		for(unsigned int j = 0; j < mat.cols; j++) {
			printf("%7.3f", matrixAt(mat, i, j));
		}
		printf("\n");
	}
	printf("]\n");
}

void clm_matrixPrintShape(clm_Matrix mat) {
	printf("(%dx%d)\n", mat.rows, mat.cols);
}