/* compile with 
gcc -lgsl -lgslcblas testpsf.c 
*/
#include <stdio.h>
#include <math.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>
struct data {
   size_t n;
   double *x;
   double *y;
   double *e;
};

#define SIGMA 2.
const int ntrial=100,nsample=100,npar=3;
const double sky=100.,amp=5000.,xc=30.,sigma=SIGMA;
const double norm=1./2.50662827463;
const double alpha=0.2,deltax=1.,sigma2=2.*SIGMA;

int gaussplussky_f(const gsl_vector *par, void *data, gsl_vector *f)
{
  int n = ((struct data *)data)->n;
  double *x = ((struct data *)data)->x;
  double *y = ((struct data *)data)->y;
  double *e = ((struct data *) data)->e;

  double sky = gsl_vector_get(par, 0);
  double amp = gsl_vector_get(par, 1);
  double xc = gsl_vector_get(par, 2);

  int i;
  for (i=0;i<n;i++){
    /* Model Yi = sky+amp*norm/sigma*exp(-0.5((x-xc)/sigma)^2) */
    double Yi = sky+amp*norm/sigma*exp(-0.5*pow((x[i]-xc)/sigma,2));
    gsl_vector_set(f, i, (Yi-y[i])/e[i]);
  }
  return GSL_SUCCESS;
}

int gaussplussky_df(const gsl_vector *par, void *data, gsl_matrix *J)
{
  size_t n = ((struct data *)data)->n;
  double *x = ((struct data *) data)->x;
  double *e = ((struct data *) data)->e;

  double sky = gsl_vector_get(par, 0);
  double amp = gsl_vector_get(par, 1);
  double xc = gsl_vector_get(par, 2);

  int i;
  double s,dx,expg;
  for (i=0;i<n;i++){
    /* Jacobian matrix J(i,j) = dfi / dxj, */
    /* where fi = (Yi - yi)/e[i],      */
    /*       Yi = sky+amp*norm/sigma*exp(-0.5((x-xc)/sigma)^2 */
    /* and the xj are the parameters (A,lambda,b) */
    s = e[i];
    dx = (x[i]-xc)/sigma;
    expg = norm/sigma*exp(-0.5*pow(dx,2));
    gsl_matrix_set(J, i, 0, 1./s); 
    gsl_matrix_set(J, i, 1, expg/s);
    gsl_matrix_set(J, i, 2, expg*amp*dx/sigma/s);
  }
  return GSL_SUCCESS;
}

int gaussplussky_fdf (const gsl_vector *par, void *data, 
                      gsl_vector *f, gsl_matrix *J)
{
  gaussplussky_f (par, data, f);
  gaussplussky_df (par, data, J);
  return GSL_SUCCESS;
}

void print_state (size_t iter, gsl_multifit_fdfsolver * s)
{
  printf ("iter: %3u x = % 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_blas_dnrm2 (s->f));
}

int main (void)
{
  int i,j,ps;
  double x[nsample],m[nsample],f[nsample],e[nsample];
  double sumchi2,sumdevsky,sumdevamp,sumdevxc,sumdev2sky,sumdev2amp,sumdev2xc;
  double alph,chi2;
  /* Random number generator setup; see example 
http://www.gnu.org/software/gsl/manual/html_node/Random-Number-Distribution-Examples.html#Random-Number-Distribution-Examples
  */
  const gsl_rng_type *T;
  gsl_rng *r;
  gsl_rng_env_setup();
  T = gsl_rng_default;
  r = gsl_rng_alloc (T);
  /* NLSQ setup; see example
http://www.gnu.org/software/gsl/manual/html_node/Example-programs-for-Nonlinear-Least_002dSquares-Fitting.html#Example-programs-for-Nonlinear-Least_002dSquares-Fitting
  */
  const gsl_multifit_fdfsolver_type *F;
  gsl_multifit_fdfsolver *s;
  gsl_matrix *covar = gsl_matrix_alloc (npar,npar);
  struct data d = {nsample, x, f, e};
  gsl_multifit_function_fdf gaussplussky;
  double par_init[3];
  par_init[0] = sky;
  par_init[1] = amp;
  par_init[2] = xc;
  gsl_vector_view par = gsl_vector_view_array(par_init, npar);
  gaussplussky.f = &gaussplussky_f;
  gaussplussky.df = &gaussplussky_df;
  gaussplussky.fdf = &gaussplussky_fdf;
  gaussplussky.n = nsample;
  gaussplussky.p = npar;
  gaussplussky.params = &d;
  /* loop over pure Gaussian and one with distortion */
  for(ps=2;ps<=3;ps++){
    alph = alpha*(double)(ps-2); /* no dist. for ps=2, full dist. for ps=3 */
    printf("***** Run with alpha=%10.6f\n",alph);
  /* ntrial times, ... */
    sumchi2=sumdevsky=sumdevamp=sumdevxc=sumdev2sky=sumdev2amp=sumdev2xc=0;
    printf("It. #Chi2     DOF   Sky     +/-   Amp       +/-     Xc       +/-\n");
    for(i=0;i<ntrial;i++){
      /* contruct a star */
      for(j=0;j<nsample;j++){
	x[j] = (double)j;
	m[j] = sky+amp*(1.-alph)*norm/sigma*exp(-0.5*pow((x[j]-xc)/sigma,2))
	          +amp*alph*norm/sigma2*exp(-0.5*pow((x[j]-xc-deltax)/sigma2,2));
      }
      /* add noise */
      for(j=0;j<nsample;j++){
	f[j] = (double)gsl_ran_poisson(r,m[j]);
      }
      /* calculate errors from "observed" fluxes */
      for(j=0;j<nsample;j++){
	e[j] = sqrt(f[j]);
      /*  printf("%5.1f %10.2f %10.2f\n",x[j],f[j],e[j]); */
      }
      F = gsl_multifit_fdfsolver_lmsder;
      s = gsl_multifit_fdfsolver_alloc(F, nsample, npar);
      gsl_multifit_fdfsolver_set(s, &gaussplussky, &par.vector);

      int iter=0,status;
      /* 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);
      /* get covariances */
      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))
      chi2 = pow(gsl_blas_dnrm2(s->f),2);
      printf("%3d %8.2f %4d  %7.3f %6.3f %8.3f %7.3f %9.5f %7.5f\n",
	     i,chi2,nsample-npar,FIT(0),ERR(0),FIT(1),ERR(1),FIT(2),ERR(2));
      sumchi2    += chi2;
      sumdevsky  += FIT(0)-sky;
      sumdevamp  += FIT(1)-amp;
      sumdevxc   += FIT(2)-xc;
      sumdev2sky += pow(FIT(0)-sky,2);
      sumdev2amp += pow(FIT(1)-amp,2);
      sumdev2xc  += pow(FIT(2)-xc,2);
    }
    printf("                  MD(Sky)   SD    MD(Amp)   SD      MD(Xc)   SD\n");
    printf("    %8.2f %4d  %7.3f %6.3f %8.3f %7.3f %9.5f %7.5f\n",
           sumchi2/(double)ntrial,nsample-npar,
	   sumdevsky/(double)ntrial,sqrt(sumdev2sky/(double)ntrial),
	   sumdevamp/(double)ntrial,sqrt(sumdev2amp/(double)ntrial),
	   sumdevxc/(double)ntrial, sqrt(sumdev2xc/(double)ntrial));
    printf("Scatter around MD          %6.3f          %7.3f           %7.5f\n",
	   sqrt((sumdev2sky-pow(sumdevsky,2)/(double)ntrial)/(double)(ntrial-1)),
	   sqrt((sumdev2amp-pow(sumdevamp,2)/(double)ntrial)/(double)(ntrial-1)),
	   sqrt((sumdev2xc -pow(sumdevxc,2) /(double)ntrial)/(double)(ntrial-1)));
  }
  gsl_multifit_fdfsolver_free(s);
  return 0;
}