#include "enumerate.h"

#include <stdlib.h>

#define LOOP(i,n) for(int i = 0; i < (n); i++)

void enumerate_coxeter_group(group_t *group, mat *gen, mat *matrices)
{
	TYPE t = GETTYPE(gen[0]->x[0]);
	mat_workspace *ws;

	ws = mat_workspace_init(3, t);

	mat_identity(matrices[0]);
	for(int i = 1; i < group->size; i++) {
		if(!group->elements[i].inverse)
			continue;

		int parent = group->elements[i].parent->id;
		int letter = group->elements[i].letter;

		mat_multiply(ws, matrices[i], matrices[parent], gen[letter]);
	}

	mat_workspace_clear(ws);
}

static int compare_tracedata(const void *a_, const void *b_)
{
        int d = 0;

        struct tracedata **a = (struct tracedata **)a_;
        struct tracedata **b = (struct tracedata **)b_;

        d = CMP((*a)->tr,(*b)->tr);
        if(d == 0) {
                d = CMP((*a)->trinv, (*b)->trinv);
        }

        return d;
}

static int compare_tracedata_id(const void *a, const void *b)
{
        int ida = (*(struct tracedata **)a)->id;
        int idb = (*(struct tracedata **)b)->id;

        return ida > idb ? 1 : ida < idb ? -1 : 0;
}

int enumerate_coxeter_group_traces(group_t *group, mat *gen, struct tracedata **traces_out)
{
	TYPE t = GETTYPE(gen[0]->x[0]);
	int n = group->size;

	mat *matrices = malloc(n*sizeof(mat));
	struct tracedata *traces = malloc(n*sizeof(struct tracedata));
	struct tracedata **distinct_traces = malloc(n*sizeof(struct tracedata*));

	LOOP(i, n) mat_init(matrices[i], 3, t);
	enumerate_coxeter_group(group, gen, matrices);

	LOOP(i, n) {
		INIT(traces[i].tr, t);
		INIT(traces[i].trinv, t);
	}
	LOOP(i, n) {
		if(!group->elements[i].inverse)
			continue;
		mat_trace(traces[i].tr, matrices[i]);
		mat_trace(traces[i].trinv, matrices[group->elements[i].inverse->id]);
		traces[i].id = i;
		distinct_traces[i] = &traces[i];
	}

	qsort(distinct_traces, n, sizeof(struct tracedata*), compare_tracedata);

	int nuniq = 0;
	LOOP(i, n) {
		if(i == 0 || compare_tracedata(&distinct_traces[i], &distinct_traces[nuniq-1]) != 0) {
			distinct_traces[nuniq] = distinct_traces[i];
			nuniq++;
		} else {
			int oldlength = group->elements[distinct_traces[nuniq-1]->id].length;
			int newlength = group->elements[distinct_traces[i]->id].length;
			if(newlength < oldlength)
				distinct_traces[nuniq-1]->id = distinct_traces[i]->id;
		}
	}

	qsort(distinct_traces, nuniq, sizeof(struct tracedata*), compare_tracedata_id);

	struct tracedata *unique_traces = malloc(nuniq*sizeof(struct tracedata));
	LOOP(i, nuniq) {
		INIT(unique_traces[i].tr, t);
		INIT(unique_traces[i].trinv, t);
		SET(unique_traces[i].tr, distinct_traces[i]->tr);
		SET(unique_traces[i].trinv, distinct_traces[i]->trinv);
		unique_traces[i].id = distinct_traces[i]->id;
	}

	LOOP(i, n) mat_clear(matrices[i]);
	LOOP(i, n) {
		CLEAR(traces[i].tr);
		CLEAR(traces[i].trinv);
	}
	free(matrices);
	free(traces);
	free(distinct_traces);

	*traces_out = unique_traces;
	return nuniq;
}

void enumerate_tracedata_clear(struct tracedata *traces, int n)
{
	LOOP(i, n) {
		CLEAR(traces[i].tr);
		CLEAR(traces[i].trinv);
	}
	free(traces);
}

// returns squares of absolute values
int solve_characteristic_polynomial_d(mps_context *solv, mps_monomial_poly *poly, double tr_real, double tr_imag, double trinv_real, double trinv_imag, double *eigenvalues_real, double *eigenvalues_imag)
{
//        mpq_t coeff1, coeff2, zero;
        cplx_t *roots;
        double radii[3];
        double *radii_p[3];
        mps_boolean errors;
        int result = 0;

        /*
        mpq_inits(coeff1, coeff2, zero, NULL);
        mpq_set(coeff1, trinv);
        mpq_sub(coeff2, zero, tr);
        */

        mps_monomial_poly_set_coefficient_d(solv, poly, 0, -1, 0);
        mps_monomial_poly_set_coefficient_d(solv, poly, 1, trinv_real, trinv_imag);
        mps_monomial_poly_set_coefficient_d(solv, poly, 2, -tr_real, -tr_imag);
//        mps_monomial_poly_set_coefficient_q(solv, poly, 1, coeff1, zero);
//        mps_monomial_poly_set_coefficient_q(solv, poly, 2, coeff2, zero);
        mps_monomial_poly_set_coefficient_d(solv, poly, 3, 1, 0);

        mps_context_set_input_poly(solv, (mps_polynomial*)poly);
        mps_mpsolve(solv);

        roots = cplx_valloc(3);
        for(int i = 0; i < 3; i++)
                radii_p[i] = &(radii[i]);
        mps_context_get_roots_d(solv, &roots, radii_p);
        errors = mps_context_has_errors(solv);

        if(errors) {
                result = 1;
        } else {
                for(int i = 0; i < 3; i++) {
                        eigenvalues_real[i] = cplx_Re(roots[i]);
                        eigenvalues_imag[i] = cplx_Im(roots[i]);
//                        if(fabs(cplx_Im(roots[i])) > radii[i]) // non-real root
//                                result = 2;
                }
        }

        cplx_vfree(roots);
//        mpq_clears(coeff1, coeff2, zero, NULL);

        return result;
}