#include "coxeter.h"
#include "enumerate.h"
#include "generators.h"
#include "mat.h"
#include "qext.h"

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <time.h>

#define LOOP(i,n) for(int i = 0; i < (n); i++)
#define SWAP(t,x,y) do { t _tmp = (x); (x) = (y); (y) = _tmp; } while (0);
//#define INFO(msg, ...) fprintf(stderr, "[%10.3f] " msg, runtime(), ##__VA_ARGS__)
#define INFO(msg, ...)

/*
  Number of elements
  up to length 0:        1
  up to length 1:        4
  up to length 2:       10
  up to length 3:       22
  up to length 4:       46
  up to length 5:       91
  up to length 6:      175
  up to length 7:      334
  up to length 8:      634
  up to length 9:     1198
  up to length 10:    2260
  up to length 11:    4261
  up to length 12:    8029
  up to length 13:   15124
  up to length 14:   28486
  up to length 15:   53650
  up to length 16:  101038
  up to length 17:  190279
  up to length 18:  358339
  up to length 19:  674830
  up to length 20: 1270846
  up to length 21: 2393266
  up to length 22: 4507012
  up to length 23: 8487625
*/

static double gaussian_sqrt5_real(NUMBER x)
{
	double result = 0.0;

	mpq_t tmp;
	mpq_init(tmp);

	// a_0 + sqrt(5)a_1 + 4a_2 + 2sqrt(5)a_3
	mpq_set_si(tmp, 4, 1);
	mpq_mul(tmp, tmp, x->a[2]);
	mpq_add(tmp, tmp, x->a[0]);
	result = mpq_get_d(tmp);
	mpq_set_si(tmp, 2, 1);
	mpq_mul(tmp, tmp, x->a[3]);
	mpq_add(tmp, tmp, x->a[1]);
	result += mpq_get_d(tmp)*sqrt(5);

	mpq_clear(tmp);

	return result;
}

static double gaussian_sqrt5_imag(NUMBER x)
{
	double result = 0.0;

	mpq_t tmp;
	mpq_init(tmp);

	// a_1 + 2sqrt(5)a_2 + 14a_3
	mpq_set_si(tmp, 14, 1);
	mpq_mul(tmp, tmp, x->a[3]);
	mpq_add(tmp, tmp, x->a[1]);
	result = mpq_get_d(tmp);
	mpq_set_si(tmp, 2, 1);
	mpq_mul(tmp, tmp, x->a[2]);
	result += mpq_get_d(tmp)*sqrt(5);

	mpq_clear(tmp);

	return result;
}

static int read_idlist(const char *filename, int *list)
{
	FILE *f = fopen(filename, "r");
	if(f == NULL) {
		fprintf(stderr, "Could not open %s\n", filename);
		exit(1);
	}

    char *line = NULL;
    size_t len = 0;
    ssize_t read;
    int n = 0;

    while ((read = getline(&line, &len, f)) != -1)
	    list[n++] = atoi(line);

    if (line)
        free(line);

	fclose(f);

	return n;
}

static int compare_int(const void *a, const void *b)
{
    const int *aa = a;
    const int *bb = b;
    return (*aa > *bb) - (*aa < *bb);
}

static double runtime()
{
	static struct timespec starttime;
	static int started = 0;

	if(!started) {
		clock_gettime(CLOCK_MONOTONIC, &starttime);
		started = 1;
	}

	struct timespec curtime;
	double diff;
	clock_gettime(CLOCK_MONOTONIC, &curtime);
	return (curtime.tv_sec - starttime.tv_sec) + (curtime.tv_nsec - starttime.tv_nsec) / 1e9;
}

enum mode {
	MODE_HELP,
	MODE_EIGENVALUES,
	MODE_SUMMARY,
	MODE_TRACE_IDS
};

int main(int argc, char *argv[])
{
	char buf[100];
	int mode;

	runtime(); // start timer

	// parse arguments
	if(argc < 2 || strcmp(argv[1], "help") == 0)
		mode = MODE_HELP;
	else if(strcmp(argv[1], "evs") == 0)
		mode = MODE_EIGENVALUES;
	else if(strcmp(argv[1], "summary") == 0)
		mode = MODE_SUMMARY;
	else if(strcmp(argv[1], "trace_ids") == 0)
		mode = MODE_TRACE_IDS;
	else
		mode = MODE_HELP;

	if(mode == MODE_HELP) {
		fprintf(stderr, "Usage: %s <help|evs|summary|trace_ids> [arguments]\n", argv[0]);
		fprintf(stderr, "%s help                           display this page\n", argv[0]);
		fprintf(stderr, "%s evs <n> <q1> <q2> <q3> <treal> <timag>        enumerate group and output unique (log) eigenvalue triples\n", argv[0]);
		fprintf(stderr, "%s summary <n> <q1> <q2> <q3> <treal> <timag>    only output max slope etc.\n", argv[0]);
		fprintf(stderr, "%s trace_ids <n> <q1> <q2> <q3> <treal> <timag>  list of ids of unique traces\n", argv[0]);
		return 0;
	}

	if(argc < 8) {
		fprintf(stderr, "Not enough arguments!\n");
		return 0;
	}

	int n = atoi(argv[2]);
	int nlist, nuniq;

	int q[3];
	mpq_t treal, timag;
	q[0] = atoi(argv[3]);
	q[1] = atoi(argv[4]);
	q[2] = atoi(argv[5]);
	mpq_inits(treal, timag, NULL);
	mpq_set_str(treal, argv[6], 10);
	mpq_set_str(timag, argv[7], 10);
	mpq_canonicalize(treal);
	mpq_canonicalize(timag);

	// enumerate group
	INFO("generate group\n");
	group_t *group = coxeter_init_triangle(5, 5, 5, n);

	// read list of elements we need to compute, or just fill it with all elements
	INFO("prepare list\n");
	int *idlist = malloc(n*sizeof(int));
	char *id_file = getenv("IDLIST");
	if(id_file != NULL) {
		nlist = read_idlist(id_file, idlist);
		// sort and symmetrize the list
		qsort(idlist, nlist, sizeof(int), compare_int);
		int ninverses = 0;
		LOOP(i, nlist) {
			int id = idlist[i];
			int invid = group->elements[id].inverse->id;
			if(!bsearch(&invid, idlist, nlist, sizeof(int), compare_int)) {
				idlist[nlist+ninverses] = invid;
				ninverses++;
			}
		}
		nlist += ninverses;
		qsort(idlist, nlist, sizeof(int), compare_int);
	} else {
		// just list all elements which have inverses
		nlist = 0;
		LOOP(i, n) {
			if(group->elements[i].inverse)
				idlist[nlist++] = i;
		}
	}

	// get generator matrices
	INFO("make generators\n");
	mat gen[3];
	LOOP(i, 3) mat_init(gen[i], 3, QT_GAUSS_SQRT5);
	generators_triangle_reflection_group_555_complex(gen, q[0], q[1], q[2], treal, timag);

	// compute the traces of all elements in idlist
	INFO("enumerate traces\n");
	struct tracedata *traces;
	nuniq = enumerate_coxeter_group_traces(group, gen, &traces, idlist, nlist, 1);

	// compute eigenvalues out of traces
	INFO("compute eigenvalues\n");
	mps_context *solver = mps_context_new();
	mps_monomial_poly *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);

	double ev_real[3], ev_imag[3], ev_abs2[3];
	double max_slope = 0;
	int max_slope_id = 0;

	LOOP(i, nuniq) {
		solve_characteristic_polynomial_d(solver, poly,
		                                  gaussian_sqrt5_real(traces[i].tr),
		                                  gaussian_sqrt5_imag(traces[i].tr),
		                                  gaussian_sqrt5_real(traces[i].trinv),
		                                  gaussian_sqrt5_imag(traces[i].trinv),
		                                  ev_real, ev_imag);

		LOOP(j, 3) ev_abs2[j] = ev_real[j]*ev_real[j] + ev_imag[j]*ev_imag[j];

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

		if(mode == MODE_TRACE_IDS) {
			// we only want to record unordered pairs here
			if(CMP(traces[i].tr, traces[i].trinv) >= 0)
				printf("%d\n", traces[i].id);
			continue;
		}

		if(log(ev_abs2[0]) < 1e-6) // we regard this as a finite order element
			continue;

		double slope = - log(ev_abs2[0]) / log(ev_abs2[2]);
		if(slope > max_slope) {
			max_slope = slope;
			max_slope_id = traces[i].id;
		}

		if(mode == MODE_EIGENVALUES) {
			printf("%d %f %f %f\n",
			       traces[i].id, log(ev_abs2[0])/2, log(ev_abs2[1])/2, log(ev_abs2[2])/2);
		}
	}

	// output summary
	coxeter_snprint(buf, sizeof(buf), &group->elements[max_slope_id]);
	if(mode == MODE_SUMMARY) {
		gmp_fprintf(stdout, "%f %f %d %d %d %f %s\n", mpq_get_d(treal), mpq_get_d(timag), n, nlist, nuniq, max_slope, buf);
	} else if(mode == MODE_EIGENVALUES) {
		gmp_fprintf(stderr, "q = %.2f + i*%.2f\tn = %d\tnlist = %d\tnuniq = %d\tMaximal slope: %f at %s\n", mpq_get_d(treal), mpq_get_d(timag), n, nlist, nuniq, max_slope, buf);
	} else if(mode == MODE_TRACE_IDS) {
		gmp_fprintf(stderr, "q = %.2f + i*%.2f\tElements: %d\tTraces: %d\n", mpq_get_d(treal), mpq_get_d(timag), n, nuniq);
	}

	// clean up
	INFO("clean up\n");
//	mps_monomial_poly_free(solver, MPS_POLYNOMIAL(poly));
	mps_context_free(solver);
	enumerate_tracedata_clear(traces, nuniq);
	free(idlist);
	LOOP(i, 3) mat_clear(gen[i]);
	coxeter_clear(group);
	mpq_clears(treal, timag, NULL);

	return 0;
}