# include "astro.h"
# define J2000 51544.5       /* Modified Julian date at standard epoch J2000 */

# define ESCAPE(MSG,...) {			\
    gprint (GP_ERR, MSG, __VA_ARGS__);		\
    return FALSE; }

typedef struct {
  double Ro, dRo;
  double Do, dDo;

  double uR, duR;
  double uD, duD;

  double p, dp;

  double chisq;
  int Nfit;
} PlxFit;

int FitPMandPar (PlxFit *fit, double *X, double *dX, double *Y, double *dY, double *T, double *pR, double *pD, int Npts, int VERBOSE);
int sun_ecliptic (double mjd, double *lambda, double *beta, double *epsilon, double *Radius);
int ParFactor (double *pR, double *pD, double RA, double DEC, double Time);

int fitplx (int argc, char **argv) {
  
  int i, N;

  Vector *rvec, *dvec, *tvec, *dRvec, *dDvec;

  Vector *mvec = NULL; // mask vector
  if ((N = get_argument (argc, argv, "-mask"))) {
    remove_argument (N, &argc, argv);
    if ((mvec = SelectVector (argv[N], ANYVECTOR, TRUE)) == NULL) return (FALSE);
    remove_argument (N, &argc, argv);
    CastVector (mvec, OPIHI_INT);
  }

  int VERBOSE = FALSE;
  if ((N = get_argument (argc, argv, "-v"))) {
    remove_argument (N, &argc, argv);
    VERBOSE = TRUE;
  }
  if ((N = get_argument (argc, argv, "-vv"))) {
    remove_argument (N, &argc, argv);
    VERBOSE = 2;
  }

  if (argc != 6) {
    gprint (GP_ERR, "USAGE: fitplx (ra) (dR) (dec) (dD) (mjd) [-mask mask]\n");
    // what about the errors?
    return (FALSE);
  }

  /* select input / output buffers */
  // if ((buf = SelectBuffer (argv[1], OLDBUFFER, TRUE)) == NULL) return (FALSE);
  // Nx = buf[0].header.Naxis[0];
  // Ny = buf[0].header.Naxis[1];
  
  if ((rvec = SelectVector (argv[1], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[1]);
  if ((dvec = SelectVector (argv[3], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[3]);
  if ((tvec = SelectVector (argv[5], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[5]);

  if ((dRvec = SelectVector (argv[2], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[2]);
  if ((dDvec = SelectVector (argv[4], ANYVECTOR, TRUE)) == NULL) ESCAPE ("missing vector %s\n", argv[4]);

  double *R = rvec->elements.Flt;
  double *D = dvec->elements.Flt;
  double *T = tvec->elements.Flt;
  
  double *dR = dRvec->elements.Flt;
  double *dD = dDvec->elements.Flt;

  int *mask = NULL;
  if (mvec) {
    mask = mvec->elements.Int;
  }

  N = tvec->Nelements; // XXX check other lengths

  // find mean values to remove
  double Npts = 0;
  double Tmean = 0;
  double Rmean = 0;
  double Dmean = 0;
  double Tmin = +1000000;
  double Tmax = -1000000;
  for (i = 0; i < N; i++) {
    if (mask && !mask[i]) continue;
    Rmean += R[i];
    Dmean += D[i];
    Tmean += T[i];
    Tmin = MIN(Tmin, T[i]);
    Tmax = MAX(Tmax, T[i]);
    Npts += 1.0;
  }
  Rmean /= Npts;
  Dmean /= Npts;
  Tmean /= Npts;

  float Trange = Tmax - Tmin;
  // fprintf (stderr, "R,D : %f,%f, T: %f, Trange: %f, Tmin: %f, Tmax: %f\n", Rmean, Dmean, Tmean, Trange, Tmin, Tmax);

  /* project coordinates to a plane centered on the object with units of arcsec */
  Coords coords;
  InitCoords (&coords, "DEC--SIN");
  coords.crval1 = Rmean;
  coords.crval2 = Dmean;
  coords.cdelt1 = coords.cdelt2 = 1.0 / 3600.0;

  double *X, *Y, *t, *pX, *pY, *dX, *dY;
  ALLOCATE (X, double, N);
  ALLOCATE (Y, double, N);
  ALLOCATE (dX, double, N);
  ALLOCATE (dY, double, N);
  ALLOCATE (t, double, N);
  ALLOCATE (pX, double, N);
  ALLOCATE (pY, double, N);

  float pXmin = +2.0;
  float pXmax = -2.0;
  float pYmin = +2.0;
  float pYmax = -2.0;

  int n = 0;
  for (i = 0; i < N; i++) {
    if (mask && !mask[i]) continue;
    RD_to_XY (&X[n], &Y[n], R[i], D[i], &coords);
    dX[n] = dR[i];
    dY[n] = dD[i];
    t[n] = (T[i] - Tmean) / 365.25;
    ParFactor (&pX[n], &pY[n], R[i], D[i], T[i]);
    pXmin = MIN (pXmin, pX[n]);
    pXmax = MAX (pXmax, pX[n]);
    pYmin = MIN (pYmin, pY[n]);
    pYmax = MAX (pYmax, pY[n]);
    n++;
  }
  float dXRange = pXmax - pXmin;
  float dYRange = pYmax - pYmin;
  float parRange = hypot (dXRange, dYRange);
	
  // fprintf (stderr, "par factor range: %f\n", parRange);

  PlxFit fit;
  if (!FitPMandPar (&fit, X, dX, Y, dY, t, pX, pY, n, VERBOSE)) {
    return FALSE;
  }

  // fprintf (stderr, "Roff, Doff: %f, %f; dRo, dDo: %f, %f\n", fit.Ro, fit.Do, fit.dRo, fit.dDo);
  
  XY_to_RD (&Rmean, &Dmean, fit.Ro, fit.Do, &coords);
  if (VERBOSE) {
    fprintf (stderr, "Ro, Do: %f, %f +/- %f, %f\n", Rmean, Dmean, fit.dRo, fit.dDo);
    fprintf (stderr, "uR, uD: %f, %f; duR, duD: %f, %f\n", fit.uR, fit.uD, fit.duR, fit.duD);
    fprintf (stderr, "par: %f +/- %f\n", fit.p, fit.dp);
    fprintf (stderr, "chisq: %f Nfit %d\n", fit.chisq, fit.Nfit);
  }

  set_variable ("RA",   Rmean);
  set_variable ("DEC",  Dmean);
  set_variable ("dR",   fit.dRo);
  set_variable ("dD",   fit.dDo);
  set_variable ("uR",   fit.uR);
  set_variable ("uD",   fit.uD);
  set_variable ("duR",   fit.duR);
  set_variable ("duD",   fit.duD);
  set_variable ("plx",  fit.p);
  set_variable ("dplx", fit.dp);
  
  set_variable ("Tmean",  Tmean);
  set_variable ("Trange", Trange);
  set_variable ("Prange", parRange);

  set_variable ("chisq", fit.chisq);
  set_variable ("Nfit",  fit.Nfit);

  return (TRUE);
}

/* do we want an init function which does the alloc and a clear function to free? */
int FitPMandPar (PlxFit *fit, double *X, double *dX, double *Y, double *dY, double *T, double *pR, double *pD, int Npts, int VERBOSE) {

  int i;

  static double **A = NULL;
  static double **B = NULL;
  double wx, wy, Wx, Wy, Tx, Ty, Tx2, Ty2, Xs, Ys, XT, YT;
  double PR, PD, PRT, PDT, PRX, PDY, PR2, PD2;
  double chisq, Xf, Yf;

  if (A == NULL) {
    ALLOCATE (A, double *, 5);
    ALLOCATE (B, double *, 5);
    for (i = 0; i < 5; i++) {
      ALLOCATE (A[i], double, 5);
      ALLOCATE (B[i], double, 1);
      memset (A[i], 0, 5*sizeof(double));
      memset (B[i], 0, 1*sizeof(double));
    }
  }

  PR = PD = PRT = PDT = PRX = PDY = PR2 = PD2 = 0.0;
  Wx = Wy = Tx = Ty = Tx2 = Ty2 = Xs = Ys = XT = YT = 0.0;
  for (i = 0; i < Npts; i++) {

    if (VERBOSE == 2) fprintf (stderr, "%f %f : %f %f : %f : %f %f\n", X[i], dX[i], Y[i], dY[i], T[i], pR[i], pD[i]);

    /* handle case where dX or dY = 0.0 */
    wx = (fabs(dX[i]) < 0.0001) ? 1.0 : 1.0 / SQ(dX[i]);
    wy = (fabs(dY[i]) < 0.0001) ? 1.0 : 1.0 / SQ(dY[i]);

    Wx += wx;
    Wy += wy;

    Tx += T[i]*wx;
    Ty += T[i]*wy;
    
    Tx2 += SQ(T[i])*wx;
    Ty2 += SQ(T[i])*wy;
    
    PR += pR[i]*wx;
    PD += pD[i]*wy;
    
    PRT += pR[i]*T[i]*wx;
    PDT += pD[i]*T[i]*wy;
    
    PRX += pR[i]*X[i]*wx;
    PDY += pD[i]*Y[i]*wy;
    
    PR2 += SQ(pR[i])*wx;
    PD2 += SQ(pD[i])*wy;

    Xs += X[i]*wx;
    Ys += Y[i]*wy;

    XT += X[i]*T[i]*wx;
    YT += Y[i]*T[i]*wy;
  }

  A[0][0] = Wx;
  A[0][1] = Tx;
  A[0][4] = PR;

  A[1][0] = Tx;
  A[1][1] = Tx2;
  A[1][4] = PRT;

  A[2][2] = Wy;
  A[2][3] = Ty;
  A[2][4] = PD;

  A[3][2] = Ty;
  A[3][3] = Ty2;
  A[3][4] = PDT;

  A[4][0] = PR;
  A[4][1] = PRT;
  A[4][2] = PD;
  A[4][3] = PDT;
  A[4][4] = PR2 + PD2;

  B[0][0] = Xs;
  B[1][0] = XT;
  B[2][0] = Ys;
  B[3][0] = YT;
  B[4][0] = PRX + PDY;

  if (!dgaussjordan ((double **)A, (double **)B, 5, 1)) {
    if (VERBOSE) fprintf (stderr, "error in fit\n");
    if (VERBOSE == 2) {
      int j;
      for (i = 0; i < 5; i++) {
	for (j = 0; j < 5; j++) {
	  fprintf (stderr, "%e ", A[i][j]);
	}
	fprintf (stderr, " : %e\n", A[i][0]);
      }
    }
    return FALSE;
  }

  fit[0].Ro = B[0][0];
  fit[0].uR = B[1][0];
  fit[0].Do = B[2][0];
  fit[0].uD = B[3][0];
  fit[0].p  = B[4][0];
  
  fit[0].dRo = sqrt(A[0][0]);
  fit[0].duR = sqrt(A[1][1]);
  fit[0].dDo = sqrt(A[2][2]);
  fit[0].duD = sqrt(A[3][3]);
  fit[0].dp  = sqrt(A[4][4]);
  
  // add up the chi square for the fit
  chisq = 0.0;
  for (i = 0; i < Npts; i++) {
    Xf = fit[0].Ro + fit[0].uR*T[i] + fit[0].p*pR[i];
    Yf = fit[0].Do + fit[0].uD*T[i] + fit[0].p*pD[i];
    wx = (fabs(dX[i]) < 0.0001) ? 1.0 : 1.0 / SQ(dX[i]);
    wy = (fabs(dY[i]) < 0.0001) ? 1.0 : 1.0 / SQ(dY[i]);
    chisq += SQ(X[i] - Xf) * wx;
    chisq += SQ(Y[i] - Yf) * wy;
    // if (VERBOSE) fprintf (stderr, "chisq contrib : %f %f : %f %f : %f %f : %f %f : %f\n", Xf, Yf, X[i] - Xf, Y[i] - Yf, dX[i], dY[i], (X[i] - Xf) / dX[i], (Y[i] - Yf) / dY[i], chisq);
  }
  fit[0].Nfit = Npts;

  // the reduced chisq is divided by (Ndof = 2*Npts - 5)
  fit[0].chisq = chisq / (2.0*Npts - 5.0);
  return (TRUE);
}

/* Low precision formulae for the sun, from Astro. Almanac p. C5 (2012) */
int sun_ecliptic (double mjd, double *lambda, double *beta, double *epsilon, double *Radius) {

  double n = mjd - J2000;                          // day number relative to standard epoch
  double L = 280.460 + 0.9856474 * n;	           // mean solar longitute (corr. for aberration)
  double g = (357.528 + 0.9856003 * n)*RAD_DEG;    // Mean anomaly

  *lambda = L + 1.915 * sin(g) + 0.020 * sin(2*g); // solar longitude in degrees
  *beta = 0.0;					   // approx latitude
  *epsilon = (23.439 - 0.0000004 * n);		   // obliquity of ecliptic in degrees
  *Radius = 1.00014 - 0.01671*cos(g) - 0.00014*cos(2*g); // earth-to-sun dist in AU
  return TRUE;
}

/* given RA, DEC, Time, calculate the parallax factor */
// RA,DEC are decimal degrees
// Time is MJD
int ParFactor (double *pR, double *pD, double RA, double DEC, double Time) {

  double lambda, beta, epsilon, Radius;

  // Time must be mjd
  sun_ecliptic (Time, &lambda, &beta, &epsilon, &Radius);

  double lambda_rad = lambda*RAD_DEG;
  double epsilon_rad = epsilon*RAD_DEG;
  double RA_rad = RA*RAD_DEG;
  double DEC_rad = DEC*RAD_DEG;

  double x = Radius*cos(lambda_rad);
  double y = Radius*cos(epsilon_rad)*sin(lambda_rad);
  double z = Radius*sin(epsilon_rad)*sin(lambda_rad);

  *pR = +(y*cos(RA_rad) - x*sin(RA_rad));
  *pD = -(y*sin(RA_rad) + x*cos(RA_rad))*sin(DEC_rad) + z*cos(DEC_rad);

  return TRUE;
}
