#include <stdio.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_complex_math.h>
#include <memory.h>
#include <math.h>

#include "linalg.h"

#define ERROR(condition, msg, ...) if(condition){fprintf(stderr, msg, ##__VA_ARGS__); exit(1);}
#define FCMP(x, y) gsl_fcmp(x, y, 1e-10)

/*********************************************** temporary storage ********************************************************/

workspace_t *workspace_alloc(int n)
{
	workspace_t *result = (workspace_t*)malloc(sizeof(workspace_t));
	result->n = n;
	result->work_nonsymmv = gsl_eigen_nonsymmv_alloc(n);
	result->work_symmv = gsl_eigen_symmv_alloc(n);
	result->work_sv = gsl_vector_alloc(n);
	result->eval_complex = gsl_vector_complex_alloc(n);
	result->evec_complex = gsl_matrix_complex_alloc(n, n);
	result->eval_real = gsl_vector_alloc(n);
	result->evec_real = gsl_matrix_alloc(n, n);
	result->permutation = gsl_permutation_alloc(n);

	result->tmp_mat = malloc(MAX_TEMP_MATRICES*sizeof(gsl_matrix*));
	for(int i = 0; i < MAX_TEMP_MATRICES; i++)
		result->tmp_mat[i] = gsl_matrix_alloc(3, 3);
	result->tmp_mat_used = 0;
	result->tmp_vec = malloc(MAX_TEMP_MATRICES*sizeof(gsl_vector*));
	for(int i = 0; i < MAX_TEMP_MATRICES; i++)
		result->tmp_vec[i] = gsl_vector_alloc(3);
	result->tmp_vec_used = 0;
	return result;
}

void workspace_free(workspace_t *workspace)
{
	gsl_eigen_nonsymmv_free(workspace->work_nonsymmv);
	gsl_eigen_symmv_free(workspace->work_symmv);
	gsl_vector_free(workspace->work_sv);
	gsl_vector_complex_free(workspace->eval_complex);
	gsl_matrix_complex_free(workspace->evec_complex);
	gsl_vector_free(workspace->eval_real);
	gsl_matrix_free(workspace->evec_real);
	gsl_permutation_free(workspace->permutation);

	for(int i = 0; i < MAX_TEMP_MATRICES; i++)
		gsl_matrix_free(workspace->tmp_mat[i]);
	free(workspace->tmp_mat);
	for(int i = 0; i < MAX_TEMP_VECTORS; i++)
		gsl_vector_free(workspace->tmp_vec[i]);
	free(workspace->tmp_vec);
}

/************************************************** basic operations ********************************************************/

void invert(gsl_matrix *in, gsl_matrix *out, workspace_t *ws)
{
	int s;
	gsl_matrix *tmp = getTempMatrix(ws);

	gsl_matrix_memcpy(tmp, in);
	gsl_linalg_LU_decomp(tmp, ws->permutation, &s);
	gsl_linalg_LU_invert(tmp, ws->permutation, out);

	releaseTempMatrices(ws, 1);
}

void solve(gsl_matrix *A, gsl_vector *b, gsl_vector *result, workspace_t *ws)
{
	int s;
	gsl_matrix *tmp = getTempMatrix(ws);

	gsl_matrix_memcpy(tmp, A);
	gsl_linalg_LU_decomp(tmp, ws->permutation, &s);
	gsl_linalg_LU_solve(tmp, ws->permutation, b, result);

	releaseTempMatrices(ws, 1);
}

void conjugate(gsl_matrix *in, gsl_matrix *conjugator, gsl_matrix *out, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);

	invert(conjugator, out, ws);   // use out to temporarily store inverse conjugator
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, in, out, 0.0, tmp);  // in * conjugator^{-1}
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, conjugator, tmp, 0.0, out);

	releaseTempMatrices(ws, 1);
}

void multiply(gsl_matrix *a, gsl_matrix *b, gsl_matrix *out)
{
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, a, b, 0.0, out);
}

void multiply_right(gsl_matrix *a, gsl_matrix *b, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, a, b, 0.0, tmp);
	gsl_matrix_memcpy(a, tmp);
	releaseTempMatrices(ws, 1);
}

void multiply_left(gsl_matrix *a, gsl_matrix *b, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);
	gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, a, b, 0.0, tmp);
	gsl_matrix_memcpy(b, tmp);
	releaseTempMatrices(ws, 1);
}

void multiply_many(workspace_t *ws, gsl_matrix *out, int n, ...)
{
	va_list args;
	va_start(args, n);

	gsl_matrix_set_identity(out);

	for(int i = 0; i < n; i++) {
		gsl_matrix *cur = va_arg(args, gsl_matrix *);
		multiply_right(out, cur, ws);
	}

	va_end(args);
}

void cartan_calc(gsl_matrix *g, double *mu, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);

	gsl_matrix_memcpy(tmp, g);
	gsl_linalg_SV_decomp(tmp, ws->evec_real, ws->eval_real, ws->work_sv);

	for(int i = 0; i < ws->n - 1; i++)
		mu[i] = log(gsl_vector_get(ws->eval_real, i) / gsl_vector_get(ws->eval_real, i+1));

	releaseTempMatrices(ws, 1);
}

void initialize(gsl_matrix *g, double *data, int x, int y)
{
	gsl_matrix_view view = gsl_matrix_view_array(data, x, y);
	gsl_matrix_memcpy(g, &view.matrix);
}

void rotation_matrix(gsl_matrix *g, double *vector)
{
	double normalized[3];
	double norm = sqrt(vector[0]*vector[0] + vector[1]*vector[1] + vector[2]*vector[2]);
	for(int i = 0; i < 3; i++)
		normalized[i] = vector[i] / norm;
	gsl_matrix_set_identity(g);
	gsl_matrix_set(g, 0, 0, cos(norm));
	gsl_matrix_set(g, 0, 1, -sin(norm) * normalized[2]);
	gsl_matrix_set(g, 0, 2, +sin(norm) * normalized[1]);
	gsl_matrix_set(g, 1, 0, +sin(norm) * normalized[2]);
	gsl_matrix_set(g, 1, 1, cos(norm));
	gsl_matrix_set(g, 1, 2, -sin(norm) * normalized[0]);
	gsl_matrix_set(g, 2, 0, -sin(norm) * normalized[1]);
	gsl_matrix_set(g, 2, 1, +sin(norm) * normalized[0]);
	gsl_matrix_set(g, 2, 2, cos(norm));
	for(int i = 0; i < 3; i++)
		for(int j = 0; j < 3; j++)
			g->data[i * g->tda + j] += (1 - cos(norm)) * normalized[i] * normalized[j];
}

double trace(gsl_matrix *g)
{
	return gsl_matrix_get(g, 0, 0) + gsl_matrix_get(g, 1, 1) + gsl_matrix_get(g, 2, 2);
}

double determinant(gsl_matrix *g, workspace_t *ws)
{
	int s;
	double result;
	gsl_matrix *tmp = getTempMatrix(ws);

	gsl_matrix_memcpy(tmp, g);
	gsl_linalg_LU_decomp(tmp, ws->permutation, &s);
	result = gsl_linalg_LU_det(tmp, s);

	releaseTempMatrices(ws, 1);
	return result;
}

int jordan_calc(gsl_matrix *g, double *mu, workspace_t *ws)
{
	gsl_eigen_nonsymmv_params(1, ws->work_nonsymmv);
	gsl_eigen_nonsymmv(g, ws->eval_complex, ws->evec_complex, ws->work_nonsymmv);
	gsl_eigen_nonsymmv_sort(ws->eval_complex, ws->evec_complex, GSL_EIGEN_SORT_ABS_DESC);

	int real = 1;
	for(int i = 0; i < 4; i++)
		if(FCMP(GSL_IMAG(gsl_vector_complex_get(ws->eval_complex, i)), 0) != 0)
			real = 0;

	if(!real)
		return 0; // non-real eigenvalues!

	for(int i = 0; i < ws->n - 1; i++) {
		mu[i] = log(GSL_REAL(gsl_vector_complex_get(ws->eval_complex, i)) / GSL_REAL(gsl_vector_complex_get(ws->eval_complex, i+1)));
	}

	return 1;
}

int eigenvectors(gsl_matrix *g, gsl_matrix *evec_real, workspace_t *ws)
{
	gsl_matrix *g_ = getTempMatrix(ws);
	int success = 0;

	gsl_matrix_memcpy(g_, g);
	gsl_eigen_nonsymmv_params(0, ws->work_nonsymmv);
	int r = gsl_eigen_nonsymmv(g_, ws->eval_complex, ws->evec_complex, ws->work_nonsymmv);
	ERROR(r, "gsl_eigen_nonsymmv failed!\n");

	gsl_eigen_nonsymmv_sort(ws->eval_complex, ws->evec_complex, GSL_EIGEN_SORT_ABS_DESC);

	int real = 1;
	for(int i = 0; i < ws->n; i++)
		if(FCMP(GSL_IMAG(gsl_vector_complex_get(ws->eval_complex, i)), 0) != 0)
			real = 0;

	if(!real)
		goto eigenvectors_out;

	for(int i = 0; i < ws->n; i++)
		for(int j = 0; j < ws->n; j++)
			gsl_matrix_set(evec_real, i, j, GSL_REAL(gsl_matrix_complex_get(ws->evec_complex, i, j)));

	success = 1;

eigenvectors_out:
	releaseTempMatrices(ws, 1);
	return success;
}

int real_eigenvalues(gsl_matrix *g, gsl_vector *eval_real, workspace_t *ws)
{
	gsl_matrix *g_ = getTempMatrix(ws);

	gsl_matrix_memcpy(g_, g);
	gsl_eigen_nonsymmv_params(0, ws->work_nonsymmv);
	int r = gsl_eigen_nonsymmv(g_, ws->eval_complex, ws->evec_complex, ws->work_nonsymmv);
	ERROR(r, "gsl_eigen_nonsymmv failed!\n");

	gsl_eigen_nonsymmv_sort(ws->eval_complex, ws->evec_complex, GSL_EIGEN_SORT_ABS_DESC);

	int real = 0;

	for(int i = 0; i < ws->n; i++) {
		if(FCMP(GSL_IMAG(gsl_vector_complex_get(ws->eval_complex, i)), 0) == 0) {// real
			gsl_vector_set(eval_real, real, GSL_REAL(gsl_vector_complex_get(ws->eval_complex, i)));
			real++;
		}
	}

	releaseTempMatrices(ws, 1);
	return real;
}

// only fills in the real eigenvectors and returns their count
int real_eigenvectors(gsl_matrix *g, gsl_matrix *evec_real, workspace_t *ws)
{
	gsl_matrix *g_ = getTempMatrix(ws);

	gsl_matrix_memcpy(g_, g);
	gsl_eigen_nonsymmv_params(0, ws->work_nonsymmv);
	int r = gsl_eigen_nonsymmv(g_, ws->eval_complex, ws->evec_complex, ws->work_nonsymmv);
	ERROR(r, "gsl_eigen_nonsymmv failed!\n");

	gsl_eigen_nonsymmv_sort(ws->eval_complex, ws->evec_complex, GSL_EIGEN_SORT_ABS_DESC);

	int real = 0;

	for(int i = 0; i < ws->n; i++) {
		if(FCMP(GSL_IMAG(gsl_vector_complex_get(ws->eval_complex, i)), 0) == 0) {// real
			for(int j = 0; j < ws->n; j++)
				gsl_matrix_set(evec_real, j, real, GSL_REAL(gsl_matrix_complex_get(ws->evec_complex, j, i)));
			real++;
		}
	}

	releaseTempMatrices(ws, 1);
	return real;
}

void eigensystem_symm(gsl_matrix *g, gsl_vector *eval, gsl_matrix *evec, workspace_t *ws)
{
	gsl_matrix *g_ = getTempMatrix(ws);

	gsl_matrix_memcpy(g_, g);
	int r = gsl_eigen_symmv (g_, eval, evec, ws->work_symmv);
	ERROR(r, "gsl_eigen_symmv failed!\n");

	gsl_eigen_symmv_sort(eval, evec, GSL_EIGEN_SORT_ABS_DESC);

	releaseTempMatrices(ws, 1);
}

// returns number of positive directions and matrix transforming TO diagonal basis
int diagonalize_symmetric_form(gsl_matrix *A, gsl_matrix *cob, workspace_t *ws)
{
	gsl_matrix *A_ = getTempMatrix(ws);

	gsl_matrix_memcpy(A_, A);
	int r = gsl_eigen_symmv (A_, ws->eval_real, cob, ws->work_symmv);
	ERROR(r, "gsl_eigen_symmv failed!\n");

	gsl_eigen_symmv_sort(ws->eval_real, cob, GSL_EIGEN_SORT_VAL_ASC);

	gsl_matrix_transpose(cob);

	int positive = 0;

	for(int i = 0; i < ws->n; i++) {
		if(gsl_vector_get(ws->eval_real, i) > 0)
			positive++;

		for(int j = 0; j < ws->n; j++)
			*gsl_matrix_ptr(cob, i, j) *= sqrt(fabs(gsl_vector_get(ws->eval_real, i)));
	}

	releaseTempMatrices(ws, 1);
	return positive;
}

// computes a matrix in SL(3, R) which projectively transforms (e1, e2, e3, e1+e2+e3) to the 4 given vectors
void projective_frame(gsl_vector **vertices, gsl_matrix *result, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);
	gsl_vector *coeff = getTempVector(ws);
	int s;
	double det, scale;

	for(int i = 0; i < 3; i++)
		for(int j = 0; j < 3; j++)
			gsl_matrix_set(tmp, i, j, gsl_vector_get(vertices[j], i));

	gsl_linalg_LU_decomp(tmp, ws->permutation, &s);
	gsl_linalg_LU_solve(tmp, ws->permutation, vertices[3], coeff);
	det = gsl_linalg_LU_det(tmp, s);

	for(int i = 0; i < 3; i++)
		det *= gsl_vector_get(coeff, i);
	scale = 1/cbrt(det);

	for(int i = 0; i < 3; i++)
		for(int j = 0; j < 3; j++)
			gsl_matrix_set(result, i, j, scale*gsl_vector_get(vertices[j], i)*gsl_vector_get(coeff, j));

	releaseTempMatrices(ws, 1);
	releaseTempVectors(ws, 1);
}

void rotation_frame(gsl_matrix *rotation, gsl_matrix *result, workspace_t *ws)
{
	gsl_matrix *tmp = getTempMatrix(ws);
	gsl_matrix *rot_basis = getTempMatrix(ws);

	gsl_matrix_memcpy(tmp, rotation);
	gsl_eigen_nonsymmv_params(0, ws->work_nonsymmv);
	int r = gsl_eigen_nonsymmv(tmp, ws->eval_complex, ws->evec_complex, ws->work_nonsymmv);
	ERROR(r, "gsl_eigen_nonsymmv failed!\n");

	double arg, minarg = 5; // greater than pi
	int minidx;
	for(int i = 0; i < 3; i++) {
		arg = gsl_complex_arg(gsl_vector_complex_get(ws->eval_complex, i));
		if(abs(arg) < minarg)
		{
			minidx = i;
			minarg = abs(arg);
		}
	}
	ERROR(FCMP(minarg, 0.0) != 0, "rotation_frame() failed! No eigenvalue was 1.\n");

	for(int i = 0; i < 3; i++) {
		gsl_complex x = gsl_matrix_complex_get(ws->evec_complex, i, (minidx+1)%3);
		gsl_complex y = gsl_matrix_complex_get(ws->evec_complex, i, (minidx+2)%3);
		gsl_complex z = gsl_matrix_complex_get(ws->evec_complex, i, minidx);
		gsl_matrix_set(result, i, 0, GSL_REAL(x)+GSL_REAL(y));
		gsl_matrix_set(result, i, 1, GSL_IMAG(x)-GSL_IMAG(y));
		gsl_matrix_set(result, i, 2, GSL_REAL(z));
	}

	releaseTempMatrices(ws, 2);
}