#include <stdlib.h>
#include <stdio.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multifit_nlin.h>

#define N 40

void print_state (size_t iter, gsl_multifit_fdfsolver * s);

struct data {
	size_t n;
	float * y;
	float * sigma;
};

int sinb_f (const gsl_vector *x, void *data, gsl_vector *f)
{
	size_t n = ((struct data *)data)->n;
	float *y = ((struct data *)data)->y;
	float *sigma = ((struct data *) data)->sigma;

	float a = gsl_vector_get (x, 0);
	float b = gsl_vector_get (x, 1);
	float c = gsl_vector_get (x, 2);
	float d = gsl_vector_get (x, 3);

	size_t i;

	for (i = 0; i < n; i++)
	{
		/* Model Yi = a * sin(b/i + c) + d */
		float t = i;
		float Yi = a * sin (b/t + c) + d;
		gsl_vector_set (f, i, (Yi - y[i])/sigma[i]);
	}

	return GSL_SUCCESS;
}

int sinb_df (const gsl_vector * x, void *data, gsl_matrix * J)
{
	size_t n = ((struct data *)data)->n;
	float *sigma = ((struct data *) data)->sigma;

	float a = gsl_vector_get (x, 0);
	float b = gsl_vector_get (x, 1);
	float c = gsl_vector_get (x, 2);
	float d = gsl_vector_get (x, 3);

	size_t i;

	for (i = 0; i < n; i++)
	{
		/* Jacobian matrix J(i,j) = dfi / dxj,		*/
		/* where fi = (Yi - yi)/sigma[i],		*/
		/*       Yi = a * sin(b/i + c) + d		*/
		/* and the xj are the parameters (A,lambda,b)	*/

		float t = i;
		gsl_matrix_set (J, i, 0,     sinf(b/t + c) / sigma[i]);
		gsl_matrix_set (J, i, 1, a * cosf(b/t + c) / sigma[i] / t);
		gsl_matrix_set (J, i, 2, a * cosf(b/t + c) / sigma[i]);
		gsl_matrix_set (J, i, 3, 1		   / sigma[i]);
	}

	return GSL_SUCCESS;
}

int sinb_fdf (const gsl_vector *x, void *data, gsl_vector *f, gsl_matrix *J)
{
	sinb_f (x, data, f);
	sinb_df (x, data, J);

	return GSL_SUCCESS;
}

int main (void)
{
	const gsl_multifit_fdfsolver_type *T;
	gsl_multifit_fdfsolver *s;
	int status;
	unsigned int i, iter = 0;
	const size_t n = N;
	const size_t p = 4;

	gsl_matrix *covar = gsl_matrix_alloc (p, p);
	float y[N], sigma[N];
	struct data d = { n, y, sigma};
	gsl_multifit_function_fdf f;
	double x_init[4] = {1.0, 1.0, 1.0, 1.0};
	gsl_vector_view x = gsl_vector_view_array (x_init, p);
	const gsl_rng_type * type;
	gsl_rng * r;

	gsl_rng_env_setup();

	type = gsl_rng_default;
	r = gsl_rng_alloc (type);

	f.f = &sinb_f;
	f.df = &sinb_df;
	f.fdf = &sinb_fdf;
	f.n = n;
	f.p = p;
	f.params = &d;

	/* This is the data to be fitted */
	for (i = 0; i < n; i++)
	{
		float t = i;
		y[i] = 5 * sin (4/t + 3) + 2 + gsl_ran_gaussian (r, 0.3);
		sigma[i] = 0.1;
		printf ("data: %u %g %g\n", i, y[i], sigma[i]);
	};

	T = gsl_multifit_fdfsolver_lmsder;
	s = gsl_multifit_fdfsolver_alloc (T, n, p);
	gsl_multifit_fdfsolver_set (s, &f, &x.vector);

	print_state (iter, s);

	do
	{
		iter++;
		status = gsl_multifit_fdfsolver_iterate (s);

		printf ("status = %s\n", gsl_strerror (status));

		print_state (iter, s);

		if (status)
			break;

		status = gsl_multifit_test_delta (s->dx, s->x, 1e-4, 1e-4);
	}
	while (status == GSL_CONTINUE && iter < 500);

	gsl_multifit_covar (s->J, 0.0, covar);

#define FIT(i) gsl_vector_get(s->x, i)
#define ERR(i) sqrt(gsl_matrix_get(covar,i,i))

	{
		float chi = gsl_blas_dnrm2(s->f);
		float dof = n - p;
		float c = GSL_MAX_DBL(1, chi / sqrt(dof));

		printf("chisq/dof = %g\n",  pow(chi, 2.0) / dof);

		printf ("a = %.5f +/- %.5f\n", FIT(0), c*ERR(0));
		printf ("b = %.5f +/- %.5f\n", FIT(1), c*ERR(1));
		printf ("c = %.5f +/- %.5f\n", FIT(2), c*ERR(2));
		printf ("d = %.5f +/- %.5f\n", FIT(3), c*ERR(3));
	}

	printf ("status = %s\n", gsl_strerror (status));

	gsl_multifit_fdfsolver_free (s);
	gsl_matrix_free (covar);
	gsl_rng_free (r);

	return 0;
}

void print_state (size_t iter, gsl_multifit_fdfsolver * s)
{
	printf ("iter: %3u x = % 15.8f % 15.8f % 15.8f % 15.8f"
		"|f(x)| = %g\n",
		iter,
		gsl_vector_get (s->x, 0),
		gsl_vector_get (s->x, 1),
		gsl_vector_get (s->x, 2),
		gsl_vector_get (s->x, 3),
		gsl_blas_dnrm2 (s->f));
}