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

/* need to pass in a function of the form:
   funcs (x, t, a, Npar, dy/da) 
   returns y (x,t) for Npar parameters a along with dy/da at (x,t)
   dy carries 1/sig^2 
*/

# define VERY_VERBOSE 0

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

static float *parmin = NULL;
static float *parmax = NULL;

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;

  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];
    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];
      
# if (VERY_VERBOSE)
  for (j = 0; j < Npar; j++) {
    for (k = 0; k < Npar; k++) {
      gprint (GP_ERR, "%9.3e  ", ta[j][k]);
    }
    gprint (GP_ERR, "    :   %9.3e  ", tb[j][0]);
    gprint (GP_ERR, "\n");
  }
# endif

  return (chisq);

}

float mrq2dchi (float *x, float *t, float *y, float *dy, int Npts, 
		float *par, int Npar, 
		float (funcs)(float, float, float *, int, float *)) {

  int i;
  float ydiff, chisq;

  chisq = 0.0;
  for (i = 0; i < Npts; i++) {
    ydiff = funcs (x[i], t[i], par, Npar, dyda) - y[i];
    chisq += SQ(ydiff) * dy[i];
  }
  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 VERBOSE) {

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

  /* keep this test in here? */
  if (!fgaussj (talpha, Npar, tbeta, 1)) {
    lambda *= 10.0;
    return (ochisq);
  }

  for (j = 0; j < Npar; j++) {
    partry[j] = par[j] - tbeta[j][0];
    /*
    if (parmin != NULL) partry[j] = MAX (parmin[j], partry[j]);
    if (parmax != NULL) partry[j] = MIN (parmax[j], partry[j]);
    */
  }

  chisq = mrq2dcof (x, t, 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");
  }

  /* 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);

}

int mrq2dlimits (float *pmin, float *pmax, int Npar) {

  int i;

  ALLOCATE (parmin, float, Npar);
  ALLOCATE (parmax, float, Npar);
  for (i = 0; i < Npar; i++) {
    parmin[i] = pmin[i];
    parmax[i] = pmax[i];
  }
  return (TRUE);
}

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

  if (parmin != NULL) free (parmin);
  if (parmax != NULL) free (parmax);

}
