# include "data.h"  /* only needed for the ALLOCATE def */

/* need to pass in a function of the form:
   funcs (x, a, Npar, dyda) 
   returns f (x) for Npar parameters a, also df/da at x 
   dy carries 1/sig^2 
*/

static float **alpha, **talpha;
static float **beta, **tbeta;
static float *partry, *dyda;
static float ochisq, lambda;

float mrqcof (float *x, float *y, float *dy, int Npts, 
	      float *par, int Npar, float **ta, float **tb, 
	      float (funcs)(float, float *, int, float *)) {

  int k, j, i;
  float ydiff, wt, chisq;

  for (j = 0; j < Npar; j++) {
    for (k = 0; k <= j; k++) ta[j][k] = 0.0;
    tb[j][0] = 0.0;
  }

  chisq = 0.0;
  for (i = 0; i < Npts; i++) {

    ydiff = funcs (x[i], par, Npar, dyda) - y[i];
    chisq += SQ(ydiff) * dy[i];

    for (j = 0; j < Npar; j++) {
      wt = dyda[j] * dy[i];
      for (k = 0; k <= j; k++) ta[j][k] += wt * dyda[k];
      tb[j][0] += wt * ydiff;
    }
  }

  for (j = 1; j < Npar; j++)
    for (k = 0; k < j; k++) 
      ta[k][j] = ta[j][k];

  return (chisq);

}

float mrqmin (float *x, float *y, float *dy, int Npts, 
	      float *par, int Npar, 
	      float (funcs)(float, float *, int, float *), int VERBOSE) {

  int j, k;
  float chisq;
  float rho, dX, dL;

  /* set up test matrixes for this run */
  for (j = 0; j < Npar; j++) {
    for (k = 0; k < Npar; k++) talpha[j][k] = alpha[j][k];
    talpha[j][j] = alpha[j][j] * (1.0 + lambda);
    tbeta[j][0] = beta[j][0];
  }

  fgaussj (talpha, Npar, tbeta, 1);

  for (j = 0; j < Npar; j++) partry[j] = par[j] - tbeta[j][0];

  /* get linear model prediction */
  dL = 0;
  for (j = 0; j < Npar; j++) {
      dL += 0.5*lambda*SQ(tbeta[j][0]) + tbeta[j][0]*beta[j][0];
  }

  chisq = mrqcof (x, y, dy, Npts, partry, Npar, talpha, tbeta, funcs);
  if (VERBOSE) { 
    gprint (GP_ERR, "chisq: %f  ", chisq);
    gprint (GP_ERR, "lambda: %f  ", lambda);
    for (j = 0; j < Npar; j++) {
      gprint (GP_ERR, "%f ", partry[j]);
    }
    gprint (GP_ERR, "\n");
  }

  /* compare linear model with actual */
  dX = ochisq - chisq;
  rho = dX / dL;

  /* if good, save temp values */
  if (rho > 0) {
    lambda *= 0.1;
    ochisq = chisq;
    for (j = 0; j < Npar; j++) {
      for (k = 0; k < Npar; k++) alpha[j][k] = talpha[j][k];
      beta[j][0] = tbeta[j][0];
      par[j] = partry[j];
    }
  } else {
    lambda *= 10.0;
    chisq = ochisq;
  }

  return (chisq);

}

float mrqinit (float *x, float *y, float *dy, int Npts, 
	      float *par, int Npar, 
	      float (funcs)(float, float *, int, float *), int VERBOSE) {

  int i;

  ALLOCATE (dyda, float, Npar);
  ALLOCATE (partry, float, Npar);
  ALLOCATE (alpha, float *, Npar);
  ALLOCATE (beta, float *, Npar);
  ALLOCATE (talpha, float *, Npar);
  ALLOCATE (tbeta, float *, Npar);
  for (i = 0; i < Npar; i++) {
    ALLOCATE (alpha[i], float, Npar);
    ALLOCATE (beta[i], float, Npar);
    ALLOCATE (talpha[i], float, Npar);
    ALLOCATE (tbeta[i], float, Npar);
  }
  
  lambda = 0.01;
  
  ochisq = mrqcof (x, y, dy, Npts, par, Npar, alpha, beta, funcs);
  if (VERBOSE) {
    gprint (GP_ERR, "chisq: %f  ", ochisq);
    gprint (GP_ERR, "lambda: %f  ", lambda);
    for (i = 0; i < Npar; i++) {
      gprint (GP_ERR, "%f ", par[i]);
    }
    gprint (GP_ERR, "\n");
  }
  return (ochisq);
}

/* don't invoke this in the middle of a run, only near the end */ 
float **mrqcovar (int Npar) {
  fgaussj (alpha, Npar, beta, 1);
  return (alpha);
} 

void mrqfree (int Npar) {

  int i;

  for (i = 0; i < Npar; i++) {
    free (alpha[i]);
    free (talpha[i]);
    free (beta[i]);
    free (tbeta[i]);
  }
  free (alpha);
  free (talpha);
  free (beta);
  free (tbeta);
  free (partry);
  free (dyda);
}
