#include "fracture.h"
#include <gsl/gsl_sf_erf.h>
#include <gsl/gsl_sf_laguerre.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multifit_nlinear.h>

struct model_params {
	double nu;
	double ac;
	double mc;
	double ac02;
	double ac03;
	double ac12;
	double ac13;
	double bc02;
	double bc03;
	double bc12;
	double bc13;
	double Ac20;
	double Ac21;
	double Ac22;
	double Ac23;
	double Ac30;
	double Ac31;
	double Ac32;
	double Ac33;
	double Bc20;
	double Bc21;
	double Bc22;
	double Bc23;
	double Bc30;
	double Bc31;
	double Bc32;
	double Bc33;
};

struct data {
	uint_t nc;
	uint_t *Ls_c;
	double *betas_c;
	double *vals_c2;
	double *vals_c3;
};

double Jc(uint_t n, uint_t o, struct model_params *par, double x) {
	double nu, ac, mc, ac0n, bc0n, *Acn, yc, sum;

	nu = par->nu;
	ac = par->ac;
	mc = par->mc;
	ac0n = *(&(par->ac02) + (n - 2) * sizeof(double));
	bc0n = *(&(par->bc02) + (n - 2) * sizeof(double));
	Acn = &(par->Ac20) + (n - 2) * o * sizeof(double);

	yc = (log(x) - mc) / ac;

	sum = 0;

	for (uint_t i = 0; i < o; i++) {
		sum += Acn[i] * gsl_sf_laguerre_n(i, 0, yc);
	}

	return exp(ac0n + bc0n * (2 - gsl_sf_erf(yc)) + exp(-gsl_pow_2(yc)) * sum);
}

double Kc(uint_t n, uint_t o, struct model_params *par, double x) {
	double nu, ac, mc, ac1n, bc1n, *Bcn, yc, sum;

	nu = par->nu;
	ac = par->ac;
	mc = par->mc;
	ac1n = *(&(par->ac12) + (n - 2) * sizeof(double));
	bc1n = *(&(par->bc12) + (n - 2) * sizeof(double));
	Bcn = &(par->Bc20) + (n - 2) * o * sizeof(double);

	yc = (log(x) - mc) / ac;

	sum = 0;

	for (uint_t i = 0; i < o; i++) {
		sum += Bcn[i] * gsl_sf_laguerre_n(i, 0, yc);
	}

	return exp(ac1n + bc1n * (2 - gsl_sf_erf(yc)) + exp(-gsl_pow_2(yc)) * sum);
}

double sc(uint_t n, uint_t o, struct model_params *par, uint_t L, double beta) {
	double nu, bLnu, deltanu, sigmanu, tau, Lntau, Ldelta;

	nu = par->nu;
	deltanu = 1.5;
	sigmanu = 48. / 91;
	tau = 187. / 91;

	bLnu = beta * exp(log(L) / nu);
	Lntau = exp((n + 1 - tau) / sigmanu * log(L));
	Ldelta = exp(-deltanu * log(L));

	return Lntau * (Jc(n, o, par, bLnu) + Ldelta * Kc(n, 0, par, bLnu));
}

int f_moms(const gsl_vector *x, void *params, gsl_vector *f) {
	struct data *dat = (struct data *)params;

	for (uint_t i = 0; i < dat->nc; i++) {
		double f2i = sc(2, 4, (struct model_params *)x->data, dat->Ls_c[i], dat->betas_c[i]);
		double F2i = dat->vals_c2[i];

		f2i = log(f2i);
		F2i = log(F2i);

		gsl_vector_set(f, i, f2i - F2i);

		double f3i = sc(3, 4, (struct model_params *)x->data, dat->Ls_c[i], dat->betas_c[i]);
		double F3i = dat->vals_c3[i];

		f3i = log(f3i);
		F3i = log(F3i);

		gsl_vector_set(f, dat->nc + i, f3i - F3i);
	}

	return GSL_SUCCESS;
}

void mom(uint_t len, uint_t n, uint32_t *data, double result[2]) {
	uint_t total = 0;
	double mom = 0;
	double mom_err = 0;

	for (uint_t i = 0; i < len; i++) {
		uint32_t datai = data[i];
		double in = pow(i, n);

		total += datai;
		mom += in * datai;
		mom_err += gsl_pow_2(in) * datai; 
	}

	double momf = mom / total;
	double momf_err = momf * sqrt(mom_err / gsl_pow_2(mom) + 1 / total);

	result[0] = momf;
	result[1] = momf_err;
}

void
callback(const size_t iter, void *params,
         const gsl_multifit_nlinear_workspace *w)
{
  gsl_vector *f = gsl_multifit_nlinear_residual(w);

  fprintf(stderr, "iter %2zu: |f(x)| = %.4f\n",
          iter,
          gsl_blas_dnrm2(f));
}


int main(int argc, char *argv[]) {
	struct data *d = (struct data *)malloc(sizeof(struct data));
	uint_t nc = argc - 1;
	uint_t *Ls_c = (uint_t *)malloc(nc * sizeof(uint_t));
	double *betas_c = (double *)malloc(nc * sizeof(double));
	double *vals_c2 = (double *)malloc(nc * sizeof(double));
	double *weights_c2 = (double *)malloc(nc * sizeof(double));
	double *vals_c3 = (double *)malloc(nc * sizeof(double));
	double *weights_c3 = (double *)malloc(nc * sizeof(double));

	double chisq, chisq0;

	d->nc = nc;
	d->Ls_c = Ls_c;
	d->betas_c = betas_c;
	d->vals_c2 = vals_c2;
	d->vals_c3 = vals_c3;

	for (uint_t i = 0; i < nc; i++) {
		char *fn = argv[1 + i];
		char *buff = malloc(20 * sizeof(char));
		uint_t pos = 0; uint_t c = 0;
		while (c < 4) {
			if (fn[pos] == '_') {
				c++;
			}
			pos++;
		}
		c = 0;
		while (fn[pos] != '_') {
			buff[c] = fn[pos];
			pos++;
			c++;
		}
		buff[c] = '\0';
		Ls_c[i] = atoi(buff);
		c = 0;
		pos++;
		while (fn[pos] != '_') {
			buff[c] = fn[pos];
			pos++;
			c++;
		}
		buff[c] = '\0';
		betas_c[i] = atof(buff);
		free(buff);

		uint_t dist_len = 4 * pow(Ls_c[i], 2);
		uint32_t *dist = malloc(dist_len * sizeof(uint32_t));
		FILE *dist_file = fopen(fn, "rb");
		fread(dist, sizeof(uint32_t), dist_len, dist_file);
		fclose(dist_file);
		{
			double mom2[2];
			mom(dist_len, 2, dist, mom2);
			vals_c2[i] = mom2[0];
			weights_c2[i] = mom2[1];

			double mom3[2];
			mom(dist_len, 3, dist, mom3);
			vals_c3[i] = mom3[0];
			weights_c3[i] = mom3[1];
		}
		free(dist);
	}

	const gsl_multifit_nlinear_type *T = gsl_multifit_nlinear_trust;
	gsl_multifit_nlinear_workspace *w;
	gsl_multifit_nlinear_fdf fdf;
	gsl_multifit_nlinear_parameters fdf_params = gsl_multifit_nlinear_default_parameters();
	uint_t n = 2 * nc;
	uint_t p = 27;

	double x_init[27] = { 1.56, .3, 2, -6, -10, -10, -10, 10, 15, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
	double weights[n];
	for (uint_t i = 0; i < nc; i++) {
		weights[i] = 1/gsl_pow_2(weights_c2[i] / vals_c2[i]);
		weights[nc + i] = 1/gsl_pow_2(weights_c3[i] / vals_c3[i]);
	}

	gsl_vector_view x = gsl_vector_view_array(x_init, p);
	gsl_vector_view wts = gsl_vector_view_array(weights, n);

	gsl_vector *f;

	fdf.f = f_moms;
	fdf.df = NULL;
	fdf.fvv = NULL;
	fdf.n = n;
	fdf.p = p;
	fdf.params = d;

	fdf_params.trs = gsl_multifit_nlinear_trs_lm;

	w = gsl_multifit_nlinear_alloc(T, &fdf_params, n, p);
	gsl_multifit_nlinear_winit(&x.vector, &wts.vector, &fdf, w);
	f = gsl_multifit_nlinear_residual(w);
	gsl_blas_ddot(f, f, &chisq0);

	const double xtol = 0.0;
	const double gtol = 0.0;
	const double ftol = 0.0;

	int info;
	int status = gsl_multifit_nlinear_driver(100, xtol, gtol, ftol, callback, NULL, &info, w);
	gsl_blas_ddot(f, f, &chisq);

	fprintf(stderr, "summary from method '%s/%s'\n",
			          gsl_multifit_nlinear_name(w),
								          gsl_multifit_nlinear_trs_name(w));

	fprintf(stderr, "number of iterations: %zu\n",
			          gsl_multifit_nlinear_niter(w));
	  fprintf(stderr, "function evaluations: %zu\n", fdf.nevalf);
			  fprintf(stderr, "reason for stopping: %s\n",
						          (info == 1) ? "small step size" : "small gradient");
				  fprintf(stderr, "initial |f(x)| = %g\n", sqrt(chisq0));
					  fprintf(stderr, "final   |f(x)| = %g\n", sqrt(chisq));

	free(Ls_c);
	free(betas_c);
	free(vals_c2);
	free(weights_c2);
	free(vals_c3);
	free(weights_c3);

	return 0;
}