#include "qext.h"

#include <assert.h>

int qext_rank;
mpq_t *qext_coefficient;

#include <stdio.h>
#include <malloc.h>

#define LOOP(i,n) for(int i = 0; i < (n); i++)
#define RANK(x) ((x)->type->rank)

struct qext_type *QT_TRIVIAL = &(struct qext_type) {
	.rank = 1,
	.integer_coeffs = (const int[]) {-1, 1}
};

struct qext_type *QT_SQRT5 = &(struct qext_type) {
	.rank = 2,
	.integer_coeffs = (const int[]) {-5, 0, 1}
};

struct qext_type *QT_GAUSS_SQRT5 = &(struct qext_type) {
	.rank = 4,
	.integer_coeffs = (const int[]) {36, 0, -8, 0, 1}
};

static void qext_init_type(struct qext_type *type)
{
	type->coeffs = malloc(type->rank*sizeof(mpq_t));
	LOOP(i, type->rank) {
		mpq_init(type->coeffs[i]);
		mpq_set_si(type->coeffs[i],
		           -type->integer_coeffs[i],
		           type->integer_coeffs[type->rank]);
	}
}

struct qext_type *qext_newtype(int rank, const int *coeffs)
{
	struct qext_type *type = malloc(sizeof(struct qext_type));
	type->rank = rank;
	type->integer_coeffs = coeffs;
	qext_init_type(type);

	return type;
}

void qext_init(qext_number x, struct qext_type *type)
{
	if(type->coeffs == NULL) // uninitialized default type
		qext_init_type(type);

	x->type = type;
	x->a = malloc(RANK(x)*sizeof(mpq_t));
	LOOP(i, RANK(x)) mpq_init(x->a[i]);
}

void qext_clear(qext_number x)
{
	LOOP(i, RANK(x)) mpq_clear(x->a[i]);
	free(x->a);
}

void qext_set(qext_number x, qext_number y)
{
	LOOP(i, RANK(x)) mpq_set(x->a[i], y->a[i]);
}

void qext_set_int(qext_number x, int y)
{
	mpq_set_si(x->a[0], y, 1);
	for(int i = 1; i < RANK(x); i++)
		mpq_set_ui(x->a[i], 0, 1);
}

void qext_set_q(qext_number x, mpq_t y)
{
	mpq_set(x->a[0], y);
	for(int i = 1; i < RANK(x); i++)
		mpq_set_ui(x->a[i], 0, 1);
}

void qext_add(qext_number result, qext_number x, qext_number y)
{
	LOOP(i, RANK(x)) mpq_add(result->a[i], x->a[i], y->a[i]);
}

void qext_sub(qext_number result, qext_number x, qext_number y)
{
	LOOP(i, RANK(x)) mpq_sub(result->a[i], x->a[i], y->a[i]);
}

void qext_neg(qext_number result, qext_number x)
{
	LOOP(i, RANK(x)) mpq_neg(result->a[i], x->a[i]);
}

void qext_print(qext_number x)
{
	LOOP(i, RANK(x)) {
		if(i == 0)
			gmp_printf("%Qd", x->a[i]);
		else if(i == 1)
			gmp_printf(" + %Qd*a", x->a[i]);
		else
			gmp_printf(" + %Qd*a^%d", x->a[i], i);
	}
}

void qext_snprint(char *str, size_t size, qext_number x)
{
	int len = 0;

	LOOP(i, RANK(x)) {
		if(i == 0)
			len += gmp_snprintf(str+len, size-len, "%Qd", x->a[i]);
		else if(i == 1)
			len += gmp_snprintf(str+len, size-len, " + %Qd*a", x->a[i]);
		else
			len += gmp_snprintf(str+len, size-len, " + %Qd*a^%d", x->a[i], i);
	}
}

int qext_cmp(qext_number x, qext_number y)
{
	int cmp;

	LOOP(i, RANK(x)) {
		cmp = mpq_cmp(x->a[i], y->a[i]);
		if(cmp)
			return cmp;
	}

	return 0;
}

void qext_mul(qext_number out, qext_number x, qext_number y)
{
	mpq_t result[2*RANK(x)-1];
	mpq_t tmp;

	mpq_init(tmp);
	LOOP(i, 2*RANK(x)-1) {
		mpq_init(result[i]);
		mpq_set_ui(result[i], 0, 1);
	}

	// rk*rk multiplications + (rk-1)*rk multiplications with fixed coefficients
	// degree 1: 1+0
	// degree 2: 4+2
	// degree 4: 16+12, including 6 times 0*

	LOOP(i, RANK(x)) LOOP(j, RANK(x)) {
		mpq_mul(tmp, x->a[i], y->a[j]);
		mpq_add(result[i+j], result[i+j], tmp);
	}

	for(int i = RANK(x)-2; i >= 0; i--) {
		LOOP(j, RANK(x)) {
			mpq_mul(tmp, x->type->coeffs[j], result[RANK(x)+i]);
			mpq_add(result[i+j], result[i+j], tmp);
		}
	}

	LOOP(i, RANK(x)) mpq_set(out->a[i], result[i]);

	LOOP(i, 2*RANK(x)-1) mpq_clear(result[i]);
	mpq_clear(tmp);
}

void qext_inv(qext_number out, qext_number in)
{
	qext_number one;
	qext_init(one, in->type);
	qext_set_int(one, 1);
	qext_div(out, one, in);
	qext_clear(one);
}

void qext_div(qext_number out, qext_number x, qext_number y)
{
	int rk = RANK(x);
	struct qext_type *type = x->type;

	mpq_t tmp;
	mpq_t result[rk];
	mpq_t matrix[rk*rk];

	mpq_init(tmp);
	LOOP(i, rk) mpq_init(result[i]);
	LOOP(i, rk*rk) mpq_init(matrix[i]);

	// initialize a matrix which represents multiplication by y
	// that means the i-th column is y*a^i
	LOOP(i, rk)
		mpq_set(matrix[i*rk], y->a[i]);
	LOOP(j, rk-1) { // columns
		mpq_mul(matrix[j+1], matrix[(rk-1)*rk+j], type->coeffs[0]);
		LOOP(i, rk-1) { // rows
			mpq_mul(matrix[(i+1)*rk+(j+1)], matrix[(rk-1)*rk+j], type->coeffs[i+1]);
			mpq_add(matrix[(i+1)*rk+(j+1)], matrix[(i+1)*rk+(j+1)], matrix[i*rk+j]);
		}
	}

	// initialize result as x
	LOOP(i, rk) mpq_set(result[i], x->a[i]);

	// use Gaussian elimination to solve the system matrix * out = result
	LOOP(j, rk)
	{
		// find nonzero entry
		int row = j;
		while(row < rk && mpq_sgn(matrix[row*rk+j]) == 0) row++;

		// if the entire column is zero, the matrix was not invertible
		// this could happen if the polynomial is not irreducible / we're not in a field
		assert(row != rk);

		// permute rows
		if(row != j) {
			mpq_swap(result[row], result[j]);
			LOOP(k, rk) {
				mpq_swap(matrix[row*rk+k], matrix[j*rk+k]);
			}
		}

		// normalize
		LOOP(k, rk)
			if(k != j)
				mpq_div(matrix[j*rk+k], matrix[j*rk+k], matrix[j*rk+j]);
		mpq_div(result[j], result[j], matrix[j*rk+j]);
		mpq_set_ui(matrix[j*rk+j], 1, 1);

		// subtract
		LOOP(i, rk) {
			if(i == j)
				continue;
			mpq_mul(tmp, matrix[i*rk+j], result[j]);
			mpq_sub(result[i], result[i], tmp);
			for(int k = j+1; k < rk; k++) {
				mpq_mul(tmp, matrix[i*rk+j], matrix[j*rk+k]);
				mpq_sub(matrix[i*rk+k], matrix[i*rk+k], tmp);
			}
			mpq_set_ui(matrix[i*rk+j], 0, 1); // it isn't strictly necessary to do this as we won't use this entry again, but helps debugging
		}
	}

	LOOP(i, rk) mpq_set(out->a[i], result[i]);

	mpq_clear(tmp);
	LOOP(i, rk) mpq_clear(result[i]);
	LOOP(i, rk*rk) mpq_clear(matrix[i]);
}