#include "coxeter.h"
#include "linalg.h"
#include "mat.h"
#include "enumerate_triangle_group.h"
#include "parallel.h"

#include <time.h>

#define SWAP(t,x,y) do { t _tmp = (x); (x) = (y); (y) = _tmp; } while (0);

//#define DEBUG(msg, ...) fprintf(stderr, "[%003d%10.3f] " msg, mpi_rank(0), runtime(), ##__VA_ARGS__)
#define DEBUG(msg, ...)
#define INFO(msg, ...) fprintf(stderr, "[%003d%10.3f] " msg, mpi_rank(0), runtime(), ##__VA_ARGS__)

struct result {
	int id;
	mpq_t tr;
	mpq_t trinv;
	double x;
	double y;
	double slope;
};

// we want as much as possible to be node data, except if it is only known to the main node
// (command line arguments) or should only be computed once (id list)

struct global_data {
	// command line arguments
	unsigned int nmax;
	unsigned int p1, p2, p3;
	unsigned int sstart, send, sdenom;
	unsigned int qstart, qend, qdenom;

	unsigned int *id_list;
	unsigned int id_list_length;
};

struct node_data {
	group_t *group;
	mat* matrices;
	struct result *invariants;
	struct result **distinct_invariants;
	int distinct_invariants_length;
	mps_context *solver;
};

struct input_data {
	unsigned int snum, sden;
	unsigned int qnum, qden;
};

struct output_data {
	int max_slope_id;
	double max_slope;
};

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

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

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

	return d;
}

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

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

	return (*a)->id - (*b)->id;
}

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

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

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

	return d;
}

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

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

	double slopea = (*a)->x / (*a)->y;
	double slopeb = (*b)->x / (*b)->y;

	return slopea > slopeb ? -1 : slopea < slopeb ? 1 : 0;
}

int compute_invariants(group_t *group, mat *matrices, struct result **invariants, int *n, int unique)
{
	mpq_t tmp;
	mps_context *solver;
	mps_monomial_poly *poly;
	int index;
	int ntraces = *n, nuniq;
	int retval;
	double evs[3];
	char buf[100];

//	DEBUG("Compute traces\n");
	for(int i = 0; i < ntraces; i++) {
		int id = invariants[i]->id;
		int invid = group->elements[id].inverse->id;
		mat_trace(invariants[i]->tr, matrices[id]);
		mat_trace(invariants[i]->trinv, matrices[invid]);
	}

	if(!unique)
		nuniq = ntraces;
	else {
//		DEBUG("Get unique traces\n");
		qsort(invariants, ntraces, sizeof(struct result*), compare_result);

		nuniq = 0;
		for(int i = 0; i < ntraces; i++) {
			if(i == 0 || compare_result(&invariants[i], &invariants[nuniq-1]) != 0) {
				invariants[nuniq] = invariants[i];
				nuniq++;
			} else {
				int oldlength = group->elements[invariants[nuniq-1]->id].length;
				int newlength = group->elements[invariants[i]->id].length;
				if(newlength < oldlength)
					invariants[nuniq-1]->id = invariants[i]->id;
			}
		}
	}

	DEBUG("Solve characteristic polynomials\n");
	solver = mps_context_new();
	poly = mps_monomial_poly_new(solver, 3);
	mps_context_set_output_prec(solver, 20); // relative precision
	mps_context_set_output_goal(solver, MPS_OUTPUT_GOAL_APPROXIMATE);

	for(int i = 0; i < nuniq; i++) {
		retval = solve_characteristic_polynomial(solver, poly, invariants[i]->tr, invariants[i]->trinv, evs);

		if(retval == 1) {
			fprintf(stderr, "Error! Could not solve polynomial.\n");
			continue;
		} else if(retval == 2) {
			continue;
		}

		if(fabs(evs[0]) < fabs(evs[1]))
			SWAP(double, evs[0], evs[1]);
		if(fabs(evs[1]) < fabs(evs[2]))
			SWAP(double, evs[1], evs[2]);
		if(fabs(evs[0]) < fabs(evs[1]))
			SWAP(double, evs[0], evs[1]);

		double x = log(fabs(evs[0]));
		double y = -log(fabs(evs[2]));

		invariants[i]->x = x;
		invariants[i]->y = y;
		invariants[i]->slope = y/x;
	}
	mps_context_free(solver);

	qsort(invariants, nuniq, sizeof(struct result*), compare_result_by_id);

	*n = nuniq;
	return 0;
}

long check_memory_usage(mat *matrices, int n)
{
	mpq_t x;
	long total;

	for(int i = 0; i < n; i++)
	{
		LOOP(j,3) LOOP(k,3) {
			total += mpq_numref(M(matrices[i], j, k))->_mp_size;
			total += mpq_denref(M(matrices[i], j, k))->_mp_size;
		}
	}

	return total;
}


void destroy_node(const void *_g, void *_n)
{
	struct global_data *g = (struct global_data *)_g;
	struct node_data *n = (struct node_data *)_n;

	for(int i = 0; i < g->nmax; i++) {
		mpq_clear(n->invariants[i].tr);
		mpq_clear(n->invariants[i].trinv);
	}
	free(n->invariants);
	free(n->distinct_invariants);
	for(int i = 0; i < g->nmax; i++)
		mat_clear(n->matrices[i]);
	free(n->matrices);
	coxeter_clear(n->group);
}

int init_node(const void *_g, void *_n)
{
	struct global_data *g = (struct global_data *)_g;
	struct node_data *n = (struct node_data *)_n;

	DEBUG("Allocate\n");
	g->id_list = (int*)(g+1);  // pointers get scrambled by transmission, reconstruct
	n->matrices = malloc(g->nmax*sizeof(mat));
	for(int i = 0; i < g->nmax; i++)
		mat_init(n->matrices[i], 3);
	n->invariants = malloc(g->nmax*sizeof(struct result));
	n->distinct_invariants = malloc(g->nmax*sizeof(struct result)); // we won't need that many, but just in case
	for(int i = 0; i < g->nmax; i++) {
		mpq_init(n->invariants[i].tr);
		mpq_init(n->invariants[i].trinv);
		n->invariants[i].id = i;
	}

	// order of the triangle reflection generators: a, b, c
	// order of the rotation orders: bc, ac, ab
	DEBUG("Generate group\n");
	n->group = coxeter_init_triangle(g->p1, g->p2, g->p3, g->nmax);

	return 0;
}

int process_output(group_t *group, mat *matrices, struct result **invariants, int invariants_length, struct output_data *out)
{
	out->max_slope = 0;
	for(int i = 0; i < invariants_length; i++) {
		double x = invariants[i]->x;
		double y = invariants[i]->y;

		if(y/x > out->max_slope + 1e-12 && (x > 0.1 || y > 0.1)) {
			out->max_slope_id = invariants[i]->id;
			out->max_slope = y/x;
		}
	}
}

int do_computation(const void *_g, void *_n, const void *_in, void *_out)
{
	struct global_data *g = (struct global_data *)_g;
	struct node_data *n = (struct node_data *)_n;
	struct input_data *in = (struct input_data *)_in;
	struct output_data *out = (struct output_data *)_out;

	mpq_t s, q;

	mpq_inits(s, q, NULL);
	mpq_set_ui(s, in->snum, in->sden);
	mpq_set_ui(q, in->qnum, in->qden);

	INFO("Computing max slope element for s = %d/%d and q = %d/%d.\n",
	      in->snum, in->sden,
	      in->qnum, in->qden);

	for(int i = 0; i < n->group->size; i++)
		n->group->elements[i].need_to_compute = 0;
	n->group->elements[0].need_to_compute = 1;

	int needed_elements = 1;
	for(int i = 0; i < g->id_list_length; i++)
	{
		int id = g->id_list[i];
		n->distinct_invariants[i] = &n->invariants[id];
		groupelement_t *cur = &n->group->elements[id];
		while(cur->need_to_compute == 0) {
			cur->need_to_compute = 1;
			needed_elements++;
			cur = cur->parent->parent;  // also need to compute its even-length ancestors
		}
		cur = n->group->elements[id].inverse;
		while(cur->need_to_compute == 0) {
			cur->need_to_compute = 1;
			needed_elements++;
			cur = cur->parent->parent;
		}
	}
	n->distinct_invariants_length = g->id_list_length;

	DEBUG("Need to compute %d elements to get %d traces up to reflection length %d\n",
	      needed_elements, g->id_list_length, n->group->elements[n->group->size-1].length);

	DEBUG("Compute matrices\n");
	enumerate(n->group, n->matrices, g->p1, g->p2, g->p3, s, q);

	DEBUG("Compute invariants\n");
	compute_invariants(
		n->group, n->matrices,
		n->distinct_invariants, &n->distinct_invariants_length, 1);

	DEBUG("Find max slopes\n");
	process_output(n->group, n->matrices, n->distinct_invariants, n->distinct_invariants_length, out);

	mpq_clears(s, q, NULL);

	return 0;
}

int main(int argc, char *argv[])
{
	char buf[100];
	char buf2[100];
	struct global_data *g;
	struct node_data n;

	start_timer();

	// parse command line arguments
	if(argc < 11) {
		fprintf(stderr, "Usage: %s <N> <p1> <p2> <p3> <s start> <s end> <s denom> <q start> <q end> <q denom> [restart]\n", argv[0]);
		exit(1);
	}
	int nmax = atoi(argv[1]);
	g = (struct global_data*)malloc(sizeof(struct global_data) + nmax*sizeof(int));
	g->id_list = (int*)(g+1);
	g->nmax = nmax;
	g->p1 = atoi(argv[2]);
	g->p2 = atoi(argv[3]);
	g->p3 = atoi(argv[4]);
	g->sstart = atoi(argv[5]);
	g->send = atoi(argv[6]);
	g->sdenom = atoi(argv[7]);
	g->qstart = atoi(argv[8]);
	g->qend = atoi(argv[9]);
	g->qdenom = atoi(argv[10]);

	// initialize
	parallel_context *ctx = parallel_init();
	parallel_set_datasize_and_callbacks(ctx, init_node, do_computation, destroy_node,
	                                    sizeof(struct global_data) + g->nmax*sizeof(int),
	                                    sizeof(struct node_data),
	                                    sizeof(struct input_data),
	                                    sizeof(struct output_data));
	if(ctx->mpi_mode == 1 && ctx->rank != 0) {
		// worker mode
		parallel_work(ctx);
		parallel_destroy(ctx);
		exit(0);
	}

	init_node(g, &n);

	// use very generic values for the pilot run unless sstart=send and qstart=qend
	struct input_data pilot_input;
	struct output_data pilot_output;
	if(g->sstart == g->send && g->qstart == g->qend) {
		pilot_input.snum = g->sstart;
		pilot_input.sden = g->sdenom;
		pilot_input.qnum = g->qstart;
		pilot_input.qden = g->qdenom;
		DEBUG("Single run for s = %d/%d, q = %d/%d\n", g->sstart, g->sdenom, g->qstart, g->qdenom);
	} else {
		pilot_input.snum = 4;
		pilot_input.sden = 100;
		pilot_input.qnum = 7;
		pilot_input.qden = 100;
		DEBUG("Initial run for s = %d/%d, q = %d/%d\n", 4, 100, 7, 100);
	}

	g->id_list_length = 0;
	for(int i = 0; i < n.group->size; i++)
		if(n.group->elements[i].length % 2 == 0 && n.group->elements[i].inverse)
			g->id_list[g->id_list_length++] = i;

	do_computation(g, &n, &pilot_input, &pilot_output);

	for(int i = 0; i < n.distinct_invariants_length; i++)
		g->id_list[i] = n.distinct_invariants[i]->id;
	g->id_list_length = n.distinct_invariants_length;

	if(g->sstart != g->send || g->qstart != g->qend) {

		struct input_data *inputs = malloc((g->send - g->sstart + 1)*(g->qend - g->qstart + 1)*sizeof(struct input_data));
		struct output_data *outputs = malloc((g->send - g->sstart + 1)*(g->qend - g->qstart + 1)*sizeof(struct input_data));

		int njobs = 0;
		for(int sloop = g->sstart; sloop <= g->send; sloop++) {
			for(int qloop = g->qstart; qloop <= g->qend; qloop++) {
				inputs[njobs].sden = g->sdenom;
				inputs[njobs].qden = g->qdenom;
				inputs[njobs].snum = sloop;
				inputs[njobs].qnum = qloop;
				njobs++;
			}
		}

		if(argc >= 12)
			parallel_run(ctx, g, inputs, outputs, njobs, argv[11]);
		else
			parallel_run(ctx, g, inputs, outputs, njobs, NULL);

//				DEBUG("Loop for s = %d/%d, q = %d/%d\n", sloop, g->sdenom, qloop, g->qdenom);

		for(int i = 0; i < njobs; i++)
		{
			gmp_printf("%d/%d %d/%d %d %s %f\n",
			           inputs[i].snum, inputs[i].sden, inputs[i].qnum, inputs[i].qden,
			           outputs[i].max_slope_id,
			           print_word(&n.group->elements[outputs[i].max_slope_id], buf),
			           outputs[i].max_slope);
		}

		free(inputs);
		free(outputs);
	} else {
		// output
		for(int i = 0; i < n.distinct_invariants_length; i++) {
			// exclude tr = trinv = 2/1/0/-1/3
			mpq_t tmp;
			mpq_init(tmp);
			mpq_set_si(tmp, 2, 1);
			if(mpq_cmp(n.distinct_invariants[i]->tr, tmp) == 0 &&
			   mpq_cmp(n.distinct_invariants[i]->trinv, tmp) == 0)
				continue;
			mpq_set_si(tmp, 1, 1);
			if(mpq_cmp(n.distinct_invariants[i]->tr, tmp) == 0 &&
			   mpq_cmp(n.distinct_invariants[i]->trinv, tmp) == 0)
				continue;
			mpq_set_si(tmp, 0, 1);
			if(mpq_cmp(n.distinct_invariants[i]->tr, tmp) == 0 &&
			   mpq_cmp(n.distinct_invariants[i]->trinv, tmp) == 0)
				continue;
			mpq_set_si(tmp, -1, 1);
			if(mpq_cmp(n.distinct_invariants[i]->tr, tmp) == 0 &&
			   mpq_cmp(n.distinct_invariants[i]->trinv, tmp) == 0)
				continue;
			mpq_set_si(tmp, 3, 1);
			if(mpq_cmp(n.distinct_invariants[i]->tr, tmp) == 0 &&
			   mpq_cmp(n.distinct_invariants[i]->trinv, tmp) == 0)
				continue;
			mpq_clear(tmp);

			double slope = n.distinct_invariants[i]->y/n.distinct_invariants[i]->x;

			gmp_printf("%d %s %f\n",
			           n.distinct_invariants[i]->id,
			           print_word(&n.group->elements[n.distinct_invariants[i]->id], buf),
			           slope);
		}

	}

	destroy_node(g, &n);
	free(g);
	parallel_destroy(ctx);
}