#include "mat.h"

mat_workspace *mat_workspace_init(int n)
{
	mat_workspace *ws = (mat_workspace*)malloc(sizeof(mat_workspace));
	mat_init(ws->tmp_mat, n);
	INIT(ws->tmp_num);
	INIT(ws->tmp_num2);
	return ws;
}

void mat_workspace_clear(mat_workspace *ws)
{
	mat_clear(ws->tmp_mat);
	CLEAR(ws->tmp_num);
	CLEAR(ws->tmp_num2);
	free(ws);
}

void mat_init(mat m, int n)
{
	m->n = n;
	m->x = malloc(n*n*sizeof(NUMBER));
	LOOP(i,n) LOOP(j,n) INIT(m->x[i+j*n]);
}

void mat_get(NUMBER out, mat m, int i, int j)
{
	SET(out, M(m,i,j));
}

void mat_set(mat m, int i, int j, NUMBER x)
{
	SET(M(m,i,j), x);
}

NUMBER *mat_ref(mat m, int i, int j)
{
	return &M(m,i,j);
}

void mat_zero(mat m)
{
	LOOP(i,m->n) LOOP(j,m->n) SET_ZERO(M(m,i,j));
}

void mat_identity(mat m)
{
	LOOP(i,m->n) LOOP(j,m->n) {
		if(i == j)
			SET_ONE(M(m,i,j));
		else
			SET_ZERO(M(m,i,j));
	}
}

void mat_copy(mat to, mat from)
{
	LOOP(i,from->n) LOOP(j,from->n) SET(M(to,i,j), M(from,i,j));
}

void mat_clear(mat m)
{
	LOOP(i,m->n) LOOP(j,m->n) CLEAR(M(m,i,j));
	free(m->x);
}

int mat_same(mat m1, mat m2)
{
	return m1 == m2;
}

static void mat_multiply_outofplace(mat_workspace *ws, mat out, mat in1, mat in2)
{
	int n = out->n;
	NUMBER *tmp = &(ws->tmp_num);

	LOOP(i,n) LOOP(j,n) {
		SET_ZERO(M(out,i,j));
		LOOP(k,n) {
			MUL(*tmp, M(in1,i,k), M(in2,k,j));
			ADD(M(out,i,j), M(out,i,j), *tmp);
		}
	}
}

void mat_multiply(mat_workspace *ws, mat out, mat in1, mat in2)
{
	if(mat_same(out, in1) || mat_same(out, in2)) {
		mat_multiply_outofplace(ws, ws->tmp_mat, in1, in2);
		mat_copy(out, ws->tmp_mat);
	} else {
		mat_multiply_outofplace(ws, out, in1, in2);
	}
}

void mat_det(mat_workspace *ws, NUMBER out, mat in)
{
// let's just assume n = 3 for the moment
	NUMBER *tmp = &(ws->tmp_num);
	int n = 3;

	SET_ZERO(out);

	LOOP(i,n) {
		MUL(*tmp, M(in,0,i), M(in,1,(i+1)%3));
		MUL(*tmp, *tmp, M(in,2,(i+2)%3));
		ADD(out, out, *tmp);

		MUL(*tmp, M(in,0,(i+2)%3), M(in,1,(i+1)%3));
		MUL(*tmp, *tmp, M(in,2,i));
		SUB(out, out, *tmp);
	}
}

static void mat_pseudoinverse_outofplace(mat_workspace *ws, mat out, mat in)
{
	// let's just assume n = 3 for the moment
	NUMBER *tmp = &(ws->tmp_num);
	NUMBER *tmp2 = &(ws->tmp_num2);
	int n = 3;

	LOOP(i,n) LOOP(j,n) {
		MUL(*tmp, M(in,(i+1)%3,(j+1)%3), M(in,(i+2)%3,(j+2)%3));
		MUL(*tmp2, M(in,(i+1)%3,(j+2)%3), M(in,(i+2)%3,(j+1)%3));
		SUB(M(out,j,i), *tmp, *tmp2);
	}
}

void mat_pseudoinverse(mat_workspace *ws, mat out, mat in)
{
	if(mat_same(out, in)) {
		mat_pseudoinverse_outofplace(ws, ws->tmp_mat, in);
		mat_copy(out, ws->tmp_mat);
	} else {
		mat_pseudoinverse_outofplace(ws, out, in);
	}
}

void mat_trace(NUMBER out, mat in)
{
	SET_ZERO(out);
	ADD(out, out, M(in,0,0));
	ADD(out, out, M(in,1,1));
	ADD(out, out, M(in,2,2));
}

void mat_print(mat in)
{
	int n = in->n;

	LOOP(i,n) {
		LOOP(j,n) {
			PRINT(M(in,i,j));
			fputc(j == n - 1 ? '\n' : ' ', stdout);
		}
	}
}