# include "astro.h"

# define SUBTRACT 0

static float (*fitfunc)(float, float, float *, int, float *);

float fgaussTD (float, float, float *, int, float *);
float fgalaxyTD (float, float, float *, int, float *);
float fbarTD (float, float, float *, int, float *);
void fix_ellipsegauss_pars (float *par, int Npar);

/*** for some reason, fgaussTD fails occasionally if we use these two vars... */ 
static float f1, f2;

int imfit (int argc, char **argv) {

  int i, j, N, done, Npts, Subtract;
  double Npar, DZ, DZ2;
  float dx, dy, chisq, nchisq;
  float **covar;
  float *V, *x, *y, *z, *dz, *par, tmpsky, DFact;
  int sx, sy, nx, ny;
  char *c, line[64];
  Buffer *buf;

  Subtract = FALSE;
  if (N = get_argument (argc, argv, "-sub")) {
    remove_argument (N, &argc, argv);
    Subtract  = TRUE;
  }

  DFact = 1;
  if (N = get_argument (argc, argv, "-D")) {
    remove_argument (N, &argc, argv);
    DFact  = atof(argv[N]);
    remove_argument (N, &argc, argv);
  }

  fitfunc = fgaussTD;
  if (N = get_argument (argc, argv, "-gal")) {
    remove_argument (N, &argc, argv);
    fitfunc = fgalaxyTD; 
  }
  if (N = get_argument (argc, argv, "-bar")) {
    remove_argument (N, &argc, argv);
    fitfunc = fbarTD; 
  }

  if (argc != 6) {
    fprintf (stderr, "USAGE: imfit <buffer> sx sy nx ny\n");
    return (FALSE);
  }

  if ((buf = SelectBuffer (argv[1], OLDBUFFER, TRUE)) == NULL) return (FALSE);

  f1 = 1;
  if ((c = get_variable ("BETA1")) != (char *) NULL) f1 = atof (c);

  f2 = 1;
  if ((c = get_variable ("BETA2")) != (char *) NULL) f2 = atof (c);

  sx = atof (argv[2]);
  sy = atof (argv[3]);
  nx = atof (argv[4]);
  ny = atof (argv[5]);

  if ((sx < 0) || (sy < 0) || 
      (sx+nx > buf[0].matrix.Naxis[0]) || 
      (sy+ny > buf[0].matrix.Naxis[1])) {
    fprintf (stderr, "region out of range\n");
    return (FALSE);
  }

  Npar = 7;
  ALLOCATE (par, float, 2*Npar);
  bzero (par, Npar*sizeof(float));

  /* convert array z[x,y] to x[i], y[i], z[i] */
  Npts = nx*ny;
  ALLOCATE (x, float, 2*Npts);
  ALLOCATE (y, float, 2*Npts);
  ALLOCATE (z, float, 2*Npts);
  ALLOCATE (dz, float, 2*Npts);

  for (N = j = 0; j < ny; j++) {
    V = (float *)(buf[0].matrix.buffer) + (j+sy)*buf[0].matrix.Naxis[0] + sx; 
    for (i = 0; i < nx; i++, V++, N++) {
      dx = i + sx;
      dy = j + sy;
      z[N] = *V;
      dz[N] = 1.0 / *V;  /* temporarily no gain (or, gain = 1) */
      x[N] = dx;
      y[N] = dy;
    }
  }

  /* make initial guesses on parameters */
  /* fitting X^2 / (2*Sx), Y^2 / (2*Sy) */
  get_aperture_stats (&buf[0].matrix, (int)(sx + 0.5*nx), (int)(sy + 0.5*ny), 15, 2, 65000);
  par[0] = atof (get_variable ("Xg"));
  par[1] = atof (get_variable ("Yg"));
  par[2] = SQ (atof (get_variable ("SXg")) / 2.355);
  par[3] = SQ (atof (get_variable ("SYg")) / 2.355);
  par[4] = 0;
  par[5] = atof (get_variable ("Zcg")) / 15.0;
  /* par[5] is peak intensity, Zcg is total flux, scale down some */
  par[6] = atof (get_variable ("Sg"));

  /* soften the errors */
  for (i = 0; i < Npts; i++) {
    double tmp1;
    tmp1 = fitfunc (x[i], y[i], par, Npar, (float *) NULL);
    if (!finite (tmp1)) { 
      fprintf (stderr, "nan for tmp @ %d (%f %f %f)\n", i, x[i], y[i], z[i]); 
    } else {
      DZ = DFact * fabs (z[i] - tmp1);
      DZ2 = DFact * fabs (z[i] - tmp1);
      dz[i] *= 1 / (dz[i]*DZ + 1);  /* dz[i] = 1 / (DZ + 1/dz[i]) */
    }
  }
  /* run fit routine */
  chisq = mrq2dinit (x, y, z, dz, Npts, par, Npar, fitfunc);
  done = FALSE;
  for (i = 0; (i < 30) && !done; i++) {
    nchisq = mrq2dmin (x, y, z, dz, Npts, par, Npar, fitfunc);
    if (chisq - nchisq < 0.1) done = TRUE;
    if (nchisq > 0) chisq = nchisq;
  }  
  covar = mrq2dcovar (Npar);
  mrq2dfree (Npar);

  if (Subtract) {
    tmpsky = par[6];
    par[6] = 0;
    for (N = j = 0; j < ny; j++) {
      V = (float *)(buf[0].matrix.buffer) + (j+sy)*buf[0].matrix.Naxis[0] + sx; 
      for (i = 0; i < nx; i++, V++, N++) {
	dx = i + sx;
	dy = j + sy;
	*V -= fitfunc (dx, dy, par, Npar, (float *) NULL);
      }
    }
    par[6] = tmpsky;
  }

  fix_ellipsegauss_pars (par, Npar); 

  for (i = 0; i < Npar; i++) {
    sprintf (line, "PAR%d", i);
    set_variable (line, (double) par[i]);
    fprintf (stderr, "%f ", par[i]);
  }
  fprintf (stderr, "  %f\n", nchisq/(1.0*Npts));

  free (par);
  free (x);
  free (y);
  free (z);
  free (dz);

  return (TRUE);

}

/* pars: x, y, sx, sy, sxy, I, sky */
float fgaussTD (float x, float y, float *par, int Npar, float *dpar) {

  float X, Y, Z, E, F, q, R, f, p2, p3;

  X = x - par[0];
  Y = y - par[1];
  
  p2 = X / par[2];
  p3 = Y / par[3];

  Z = 0.5*p2*X + X*Y*par[4] + 0.5*p3*Y;                 /* R^2 */
  /* E = 1.0 / (1 + f1*Z + f2*(0.5*Z*Z + 0.1666666*Z*Z*Z)); */  /* ~ exp (-Z) */
  E = 1.0 / (1 + Z + (0.5*Z*Z + 0.1666666*Z*Z*Z));   /* ~ exp (-Z) */
  F = (f1 + f2*(Z + 0.5*Z*Z));

  q = par[5] * E;
  R = q*E*F;
  
  f = q + par[6];

  if (dpar != (float *) NULL) {
    dpar[0] = R*(p2 + par[4]*Y);
    dpar[1] = R*(p3 + par[4]*X);
    dpar[2] = 0.5*R*p2*p2;
    dpar[3] = 0.5*R*p3*p3;
    dpar[4] = -R*X*Y;
    
    dpar[5] = E;
    dpar[6] = 1;
  }

  /*
  if ((fabs(X) < 2) && (fabs(Y) < 2)) {
    int i;

    for (i = 0; i < Npar; i++) {
      fprintf (stderr, "%d  %f  %f\n", i, par[i], dpar[i]);
    }
    fprintf (stderr, "func: (%f, %f)  %f\n", X, Y, f);
  }
  */

  if (!finite (f)) {
    fprintf (stderr, "nan for fgaussTD\n");
  }

  return (f);

}

/* pars: x, y, sx, sy, sxy, I, sky */
float fgalaxyTD (float x, float y, float *par, int Npar, float *dpar) {

  float X, Y, Z, E, F, q, R, f, p2, p3;

  X = x - par[0];
  Y = y - par[1];
  
  p2 = X / par[2];
  p3 = Y / par[3];

  Z = sqrt (0.5*p2*X + X*Y*par[4] + 0.5*p3*Y);                 /* R */
  E = 1.0 / (1 + Z);   

  q = par[5] * E;
  R = q*E;
  F = 0.5 / Z;
  
  f = q + par[6];

  if (dpar != (float *) NULL) {
    dpar[0] = F*R*(p2 + par[4]*Y);
    dpar[1] = F*R*(p3 + par[4]*X);
    dpar[2] = F*0.5*R*p2*p2;
    dpar[3] = F*0.5*R*p3*p3;
    dpar[4] = -R*X*Y*F;
    
    dpar[5] = E;
    dpar[6] = 1;
  }

  /*
  if ((fabs(X) < 2) && (fabs(Y) < 2)) {
    int i;

    for (i = 0; i < Npar; i++) {
      fprintf (stderr, "%d  %f  %f\n", i, par[i], dpar[i]);
    }
    fprintf (stderr, "func: (%f, %f)  %f\n", X, Y, f);
  }
  */
  return (f);

}


/* pars: x, y, sx, sy, sxy, I, sky */
float fbarTD (float x, float y, float *par, int Npar, float *dpar) {

  float X, Y, Z, E, F, q, R, f, p2, p3;

  X = x - par[0];
  Y = y - par[1];
  
  p2 = X / par[2];
  p3 = Y / par[3];

  Z = 0.5*p2*X + X*Y*par[4] + 0.5*p3*Y;                 /* R */
  E = 1.0 / (1 + Z*Z*Z);   

  q = par[5] * E;
  F = 3*Z*Z;
  R = q*E*F;
  
  f = q + par[6];

  if (dpar != (float *) NULL) {
    dpar[0] = R*(p2 + par[4]*Y);
    dpar[1] = R*(p3 + par[4]*X);
    dpar[2] = 0.5*R*p2*p2;
    dpar[3] = 0.5*R*p3*p3;
    dpar[4] = -R*X*Y;
    
    dpar[5] = E;
    dpar[6] = 1;
  }

  /*
  if ((fabs(X) < 2) && (fabs(Y) < 2)) {
    int i;

    for (i = 0; i < Npar; i++) {
      fprintf (stderr, "%d  %f  %f\n", i, par[i], dpar[i]);
    }
    fprintf (stderr, "func: (%f, %f)  %f\n", X, Y, f);
  }
  */
  return (f);

}


void fix_ellipsegauss_pars (float *par, int Npar) {

  float p2, p4, angle, t1, t2, tmp, area;

  /* par[0], par[1] = Xo, Yo - stay the same */

  p2 = 1/par[2];
  p4 = 1/par[3];

  angle = 0.5 * atan2 (-2*par[4], p4 - p2); 

  tmp = sqrt (SQ(p2 - p4) + 4*SQ(par[4]));
  t1 = (p2 + p4 + tmp) / 2;
  t2 = t1 - tmp;

  par[2] = 2.35482*sqrt(1/t2);
  par[3] = 2.35482*sqrt(1/t1);
  par[4] = DEG_RAD * angle;

  area = 2*M_PI/sqrt(t1*t2);

  par[5] *= area;

}


