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

/* need to pass in a function of the form:

   funcs (x, t, a, Npar, dyda) 

   returns f (x, t) for Npar parameters a, also df/da at x,t 

   dy carries 1/sig^2 

*/

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

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

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

  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], t[i], par, Npar, dyda) - y[i];
    value = SQ(ydiff) * fabs(dy[i]);
    if (value < 0.0) fprintf (stderr, "neg chi: %f %f\n", value, chisq);
    chisq += value; 
    
    if (!finite (dyda[i])) { dyda[i] = 1e10; }
    if (!finite (dy[i])) { dyda[i] = 1e10; }

    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 mrq2dmin (float *x, float *t, float *y, float *dy, int Npts, 
	      float *par, int Npar, 
	      float (funcs)(float, float, float *, int, float *)) {

  int j, k;
  float chisq;

  /* 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];
  }

  if (!fgaussj (talpha, Npar, tbeta, 1)) {
    lambda *= 10.0;
    return (ochisq);
  }

  /* if we keep some fixed, set partry to just them 
     for (j = 0; j < Npar; j++) partry[j] = par[j]; */

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

  chisq = mrq2dcof (x, t, y, dy, Npts, partry, Npar, talpha, tbeta, funcs);

  /* if good, save temp values */
  if (chisq < ochisq) {
    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 mrq2dinit (float *x, float *t, float *y, float *dy, int Npts, 
	      float *par, int Npar, 
	      float (funcs)(float, float, float *, int, float *)) {

  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.001;
  
  ochisq = mrq2dcof (x, t, y, dy, Npts, par, Npar, alpha, beta, funcs);

  return (ochisq);

}

/* don't invoke this in the middle of a run, only near the end */ 
float **mrq2dcovar (int Npar) {

  fgaussj (alpha, Npar, beta, 1);
  return (alpha);

} 

void mrq2dfree (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);

}
