/** Diagnostic plots for relastro
 * @author Chris Beaumont, IfA
 */

#include "kapa.h"
#include "relastro.h"

#define TESTING

#define KAPAX 700
#define KAPAY 700

static int kapa = -1;
static int kapa2 = -1;
static int kapa3 = -1;

static int isVisual = FALSE;
static int plotRawRef = FALSE;
static int plotScatter = FALSE;
static int plotResid = FALSE;
static int plotVector = FALSE;
static int plotOutliers = TRUE;

/** Spawn a kapa window */
static int initWindow(int *kapid) {
  if (*kapid == -1) {
    *kapid = KapaOpenNamedSocket("kapa", "relastro_plots");
    if (*kapid == -1) {
      fprintf(stderr, "Failure to open kapa.\n");
      isVisual = 0;
      return 0;
    }
    //    KapaResize (*kapid, KAPAX, KAPAY);
  }
  return 1;
}

/** Ask the user how to proceed */
static int askUser(int *flag) {
  char key[10];
  fprintf (stdout, "[c]ontinue? [s]kip the rest of these plots? [a]bort all visual plots?");
  if (!fgets(key, 8, stdin)) {
    fprintf(stderr, "warning: Unable to read option");
  }
  if (key[0] == 's') {
    *flag = 0;
  }
  if (key[0] == 'a') {
    isVisual = 0;
  }
  return 1;
}

/** Size graphdata to encompass all points */
static int scaleGraphdata(float x[], float y[], Graphdata *graphdata, int n) {
    float xlo = FLT_MAX, xhi = -FLT_MAX, ylo = FLT_MAX, yhi = -FLT_MAX;
    int i;
    int goodData = 0;
    for(i = 0; i < n; i++) {
        goodData++;
        if(!finite(x[i]) || !finite(y[i])) continue;
        if(x[i] < xlo) xlo = x[i];
        if(x[i] > xhi) xhi = x[i];
        if(y[i] < ylo) ylo = y[i];
        if(y[i] > yhi) yhi = y[i];
    }
    if (goodData < 2) return 0;
    graphdata->xmin = xlo;
    graphdata->ymin = ylo;
    graphdata->xmax = xhi;
    graphdata->ymax = yhi;
    return 1;
}

static int residPlot(float x[], float y[],
                     float xVec[], float yVec[],
                     int npts, int *kapaID) {
    if (!isVisual || !plotResid) return TRUE;
    Graphdata graphdata;
    KapaSection section;

    if (!initWindow(kapaID)) return FALSE;

    KapaInitGraph(&graphdata);
    KapaClearPlots(*kapaID);
    KapaSetFont(*kapaID, "helvetica", 14);

    section.name = "0";
    section.x = 0.0; section.y = 0.0;
    section.dx = .45, section.dy = .45;
    graphdata.ptype = 7;
    graphdata.style = 2;

    KapaSetSection(*kapaID, &section);
    if(!scaleGraphdata(x, xVec, &graphdata, npts)) return 0;
    KapaSetLimits(*kapaID, &graphdata);
    KapaBox(*kapaID, &graphdata);
    KapaSendLabel(*kapaID, "x", KAPA_LABEL_XM);
    KapaSendLabel(*kapaID, "dx", KAPA_LABEL_YM);
    KapaPrepPlot(*kapaID, npts, &graphdata);
    KapaPlotVector(*kapaID, npts, x, "x");
    KapaPlotVector(*kapaID, npts, xVec, "y");

    section.x = .5; section.y = 0; section.name="1";
    KapaSetSection(*kapaID, &section);
    if(!scaleGraphdata(x, yVec, &graphdata, npts)) return 0;
    KapaSetLimits(*kapaID, &graphdata);
    KapaBox(*kapaID, &graphdata);
    KapaSendLabel(*kapaID, "x", KAPA_LABEL_XM);
    KapaSendLabel(*kapaID, "dy", KAPA_LABEL_YP);
    KapaPrepPlot(*kapaID, npts, &graphdata);
    KapaPlotVector(*kapaID, npts, x, "x");
    KapaPlotVector(*kapaID, npts, yVec, "y");

    section.x = .0; section.y = .5; section.name="2";
    KapaSetSection(*kapaID, &section);
    if(!scaleGraphdata(y, xVec, &graphdata, npts)) return 0;;
    KapaSetLimits(*kapaID, &graphdata);
    KapaBox(*kapaID, &graphdata);
    KapaSendLabel(*kapaID, "y", KAPA_LABEL_XM);
    KapaSendLabel(*kapaID, "dx", KAPA_LABEL_YM);
    KapaPrepPlot(*kapaID, npts, &graphdata);
    KapaPlotVector(*kapaID, npts, y, "x");
    KapaPlotVector(*kapaID, npts, xVec, "y");

    section.x = .5; section.y = .5; section.name="3";
    KapaSetSection(*kapaID, &section);
    if(!scaleGraphdata(y, yVec, &graphdata, npts)) return 0;
    KapaSetLimits(*kapaID, &graphdata);
    KapaBox(*kapaID, &graphdata);
    KapaSendLabel(*kapaID, "y", KAPA_LABEL_XM);
    KapaSendLabel(*kapaID, "dy", KAPA_LABEL_YP);
    KapaPrepPlot(*kapaID, npts, &graphdata);
    KapaPlotVector(*kapaID, npts, y, "x");
    KapaPlotVector(*kapaID, npts, yVec, "y");

    return 1;
}


/**Plot a vector field*/
static int plotVectorField(float x[], float y[],
                           float xVec[], float yVec[],
                           int npts, int *kapaID, double maxVecLength) {

    if(!plotVector) return 1;

    Graphdata graphdata;
    float singleX[2], singleY[2];
    float vecScaleFactor;
    float graphSize;
    int i;
    char plotTitle[50];
    sprintf(plotTitle, "Maximum Vector Size = %5.1e", maxVecLength);

    if (!initWindow(kapaID)) return 0;

    KapaInitGraph(&graphdata);
    KapaClearPlots(*kapaID);
    if(!scaleGraphdata(x, y, &graphdata, npts)) return 0;

    graphSize = graphdata.xmax - graphdata.xmin;
    if ((graphdata.ymax - graphdata.ymin) > graphSize) {
        graphSize = graphdata.ymax - graphdata.ymin;
    }

    vecScaleFactor = graphSize * 0.02 / (float)maxVecLength;
#ifdef TESTING
    fprintf(stderr, "GraphSize: %e\n maxVecLength: %e\n vecScaleFactor: %e\n",
            graphSize, maxVecLength, vecScaleFactor);
#endif


    KapaSetFont (*kapaID, "helvetica", 14);
    KapaSetLimits(*kapaID, &graphdata);
    KapaBox(*kapaID, &graphdata);
    KapaSendLabel(*kapaID, plotTitle, KAPA_LABEL_XP);

    graphdata.ptype = 7;
    graphdata.style = 2;
    KapaPrepPlot(*kapaID, npts, &graphdata);
    KapaPlotVector(*kapaID, npts, x, "x");
    KapaPlotVector(*kapaID, npts, y, "y");

    //plot each vector individually
    graphdata.ptype = 0;
    graphdata.style = 0;
    graphdata.color = KapaColorByName("blue");
    for(i = 0; i < npts; i++) {
        singleX[0] = x[i];
        singleY[0] = y[i];
        singleX[1] = x[i] + xVec[i] * vecScaleFactor;
        singleY[1] = y[i] + yVec[i] * vecScaleFactor;
        KapaPrepPlot(*kapaID, 2, &graphdata);
        KapaPlotVector(*kapaID, 2, singleX, "x");
        KapaPlotVector(*kapaID, 2, singleY, "y");
    }
    return 1;
}

int relastroVisualPlotScatter(double values[], double thresh, int npts) {
    float *x, *data;
    int i;
    float xline[2], yline[2];
    if (!isVisual || !plotScatter) return 1;
    if (!initWindow(&kapa2)) return 0;

    ALLOCATE(x, float, npts);
    ALLOCATE(data, float, npts);

    for(i = 0; i < npts; i++) {
        x[i] = i;
        data[i] = (float) values[i];
    }

    Graphdata graphdata;
    KapaSection section;
    section.x = 0; section.y = 0; section.dx = 1; section.dy = 1;
    section.name = "junk";

    KapaInitGraph(&graphdata);
    KapaClearPlots(kapa2);
    KapaSetSection(kapa2, &section);

    graphdata.ptype = 0;
    graphdata.style = 0;
    graphdata.xmin = 0;
    graphdata.xmax = npts;
    graphdata.ymin = 0;
    graphdata.ymax = data[npts-1];

    KapaSetFont(kapa2, "helvetica", 14);
    KapaSetLimits(kapa2, &graphdata);
    KapaBox(kapa2, &graphdata);
    KapaSendLabel( kapa2, "Object", KAPA_LABEL_XM);
    KapaSendLabel( kapa2, "Offset(pixels)", KAPA_LABEL_YM);
    KapaSendLabel( kapa2, "Astrometric Offset with fit cutoff",
                   KAPA_LABEL_XP);
    KapaPrepPlot(kapa2, npts, &graphdata);
    KapaPlotVector(kapa2, npts, x, "x");
    KapaPlotVector(kapa2, npts, data, "y");

    graphdata.color = KapaColorByName("red");
    KapaPrepPlot(kapa2, 2, &graphdata);
    yline[0] = (float) thresh;
    yline[1] = (float) thresh;
    xline[0] = graphdata.xmin;
    xline[1] = graphdata.xmax;
    KapaPlotVector(kapa2, 2, xline, "x");
    KapaPlotVector(kapa2, 2, yline, "y");

    askUser(&plotScatter);
    return 1;
}



/** plot raw vs ref (L, M). Only those whose distance is < drMax are used in fit*/
int relastroVisualPlotRawRef(StarData *raw, StarData *ref, double dRmax, int numObj) {

  if( !isVisual || !plotRawRef) return 1;
  if( !initWindow(&kapa)) return 0;

  float *rawX, *rawY,  *refX, *refY;
  float *rawXfit, *rawYfit, *refXfit, *refYfit;
  float *magRaw, *magRef, *magRawfit, *magReffit;
  float *xVec, *yVec;
  int numFit = 0, numNoFit = 0;
  double dL, dM, dR;

  ALLOCATE(rawX,      float, numObj);
  ALLOCATE(rawY,      float, numObj);
  ALLOCATE(refX,      float, numObj);
  ALLOCATE(refY,      float, numObj);
  ALLOCATE(rawXfit,   float, numObj);
  ALLOCATE(rawYfit,   float, numObj);
  ALLOCATE(refXfit,   float, numObj);
  ALLOCATE(refYfit,   float, numObj);
  ALLOCATE(magRaw,    float, numObj);
  ALLOCATE(magRef,    float, numObj);
  ALLOCATE(magRawfit, float, numObj);
  ALLOCATE(magReffit, float, numObj);

  int i;
  for(i = 0; i < numObj; i++) {
    if  (raw[i].mask) continue;

    dL = raw[i].L - ref[i].L;
    dM = raw[i].M - ref[i].M;
    dR = hypot (dL, dM);
    if (dR > dRmax) {
      rawX[numNoFit] = raw[i].X;
      rawY[numNoFit] = raw[i].Y;
      refX[numNoFit] = ref[i].X;
      refY[numNoFit] = ref[i].Y;
      magRaw[numNoFit] = raw[i].Mag;
      magRef[numNoFit] = ref[i].Mag;
      numNoFit++;
    } else {
      rawXfit[numFit] = raw[i].X;
      rawYfit[numFit] = raw[i].Y;
      refXfit[numFit] = ref[i].X;
      refYfit[numFit] = ref[i].Y;
      magRaw[numFit] = raw[i].Mag;
      magRef[numFit] = ref[i].Mag;
      numFit++;
    }
  }

  if (numFit == 0) return 0;

  Graphdata graphdata;

  KapaInitGraph(&graphdata);
  KapaClearPlots(kapa);

  graphdata.ptype = 7;
  graphdata.style = 2;

  if(!scaleGraphdata(rawXfit, rawYfit, &graphdata, numFit)) {
      fprintf(stderr, "Not enough finite points for plotting");
      return 0;
  }

  KapaSetFont(kapa, "helvetica", 14);
  KapaSetLimits(kapa, &graphdata);
  KapaBox(kapa, &graphdata);
  KapaSendLabel( kapa, "X", KAPA_LABEL_XM);
  KapaSendLabel( kapa, "Y", KAPA_LABEL_YM);
  KapaSendLabel( kapa, "orange, red, green, blue: (raw, ref), (nofit, fit)",
                 KAPA_LABEL_XP);

  graphdata.color = KapaColorByName("orange");
  graphdata.size = 1;
  KapaPrepPlot(kapa, numNoFit, &graphdata);
  KapaPlotVector(kapa, numNoFit, rawX, "x");
  KapaPlotVector(kapa, numNoFit, rawY, "y");

  graphdata.color = KapaColorByName("red");
  graphdata.size = 2;
  KapaPrepPlot(kapa, numNoFit, &graphdata);
  KapaPlotVector(kapa, numNoFit, refX, "x");
  KapaPlotVector(kapa, numNoFit, refY, "y");

  graphdata.color = KapaColorByName("green");
  graphdata.size = 1;
  KapaPrepPlot(kapa, numFit,  &graphdata);
  KapaPlotVector(kapa, numFit, rawXfit, "x");
  KapaPlotVector(kapa, numFit, rawYfit, "y");

  graphdata.color = KapaColorByName("blue");
  graphdata.size = 2;
  KapaPrepPlot(kapa, numFit, &graphdata);
  KapaPlotVector(kapa, numFit, refXfit, "x");
  KapaPlotVector(kapa, numFit, refYfit, "y");

  ALLOCATE(xVec, float, numFit);
  ALLOCATE(yVec, float, numFit);

  //plot the fitted objects as vectors
  for(i = 0; i < numFit; i++) {
      xVec[i] = rawXfit[i] - refXfit[i];
      yVec[i] = rawYfit[i] - refYfit[i];
  }

  plotVectorField(rawXfit, rawYfit, xVec, yVec, numFit, &kapa3, dRmax);
  if(!residPlot(rawXfit, rawYfit, xVec, yVec, numFit, &kapa2)) {
      fprintf(stderr, "Unable to plot residuals");
      return 0;
  }

  FREE(xVec);
  FREE(yVec);

  askUser(&plotRawRef);

  FREE(rawX);
  FREE(rawY);
  FREE(refX);
  FREE(refY);
  FREE(rawXfit);
  FREE(rawYfit);
  FREE(refXfit);
  FREE(refYfit);
  FREE(magRaw);
  FREE(magRef);
  FREE(magRawfit);
  FREE(magReffit);

  return 1;
}

int relastroVisualPlotOutliers(Catalog *catalog, int offset, int Nmeasure,
			       StatType statsR, StatType statsD, double thresh) {
    
  float *Din, *Rin, *Dout, *Rout;
  double xmin, xmax, ymin, ymax, range;
  float xCirc[100], yCirc[100];
  int m, i;
  int Nin, Nout;
  Measure meas;
  Graphdata graphdata;
  KapaSection section;
  
  if (!isVisual || !plotOutliers) return 1;
  if (!initWindow(&kapa)) return 0;
  
  // populate vectors
  ALLOCATE(Din, float, Nmeasure);
  ALLOCATE(Rin, float, Nmeasure);
  ALLOCATE(Dout, float, Nmeasure);
  ALLOCATE(Rout, float, Nmeasure);
  
  //create the threshhold ellipse
  for(i = 0; i < 100; i++) {
    xCirc[i] = statsR.median + thresh * statsR.sigma * cos(2 * 3.14 / 99. * i);
    yCirc[i] = statsD.median + thresh * statsD.sigma * sin(2 * 3.14 / 99. * i);  
  }

  m = offset;
  Nin = Nout = 0;
  xmin = +FLT_MAX;
  xmax = -FLT_MAX;
  ymax = -FLT_MAX;
  ymin = +FLT_MAX;
  for(i = 0; i < Nmeasure; i++, m++) {
    meas = catalog[0].measure[m];
    if (!MeasFilterTest(&meas)) continue;
    xmin = MIN(xmin, meas.dR);
    xmax = MAX(xmax, meas.dR);
    ymin = MIN(ymin, meas.dD);
    ymax = MAX(ymax, meas.dD);
    
    if (meas.dbFlags & ID_MEAS_POOR_ASTROM) {
      Rout[Nout] = (float)(meas.dR);
      Dout[Nout] = (float)(meas.dD);
      fprintf(stderr, "r: %f\td: %f\t outlier: 1\n", Rout[Nout], Dout[Nout]);
      Nout++;
    } else {
      Rin[Nin] = (float)(meas.dR);
      Din[Nin] = (float)(meas.dD);
      fprintf(stderr, "r: %f\td: %f\t outlier: 0\n", Rin[Nin], Din[Nin]);
      Nin++;
    }
  }
 
  range = (xmax - xmin);
  xmin -= .1 * range;
  xmax += .1 * range;
  range = (ymax - ymin);
  ymax += .1 * range;
  ymin -= .1 * range;

  //temporary fix
  xmin = -1; xmax = 1; ymin = -1; ymax = 1;


  //initialize graph info
  section.x = 0; section.y = 0; section.dx = 1; section.dy = 1;
  section.name = "junk";
  
  KapaInitGraph(&graphdata);
  KapaClearPlots(kapa);
  KapaSetFont(kapa, "helvetica", 14);
 
  graphdata.ptype = 7;
  graphdata.style = 2;
  graphdata.size = 3;
  graphdata.xmin = xmin;
  graphdata.xmax = xmax;
  graphdata.ymin = ymin;
  graphdata.ymax = ymax;

  KapaSetSection(kapa, &section);
  KapaSetLimits(kapa, &graphdata);
  KapaBox(kapa, &graphdata);

  KapaSendLabel( kapa, "RA (arcsec)", KAPA_LABEL_XM);
  KapaSendLabel( kapa, "Dec (arcsec)", KAPA_LABEL_YM);
  KapaSendLabel( kapa, "Points flagged as outliers (red)",
		 KAPA_LABEL_XP);

  graphdata.color = KapaColorByName("green");
  KapaPrepPlot(kapa, Nin, &graphdata);
  KapaPlotVector(kapa, Nin, Rin, "x");
  KapaPlotVector(kapa, Nin, Din, "y");

  graphdata.color = KapaColorByName("red");
  KapaPrepPlot(kapa, Nout, &graphdata);
  KapaPlotVector(kapa, Nout, Rout, "x");
  KapaPlotVector(kapa, Nout, Dout, "y");

  graphdata.color = KapaColorByName("black");
  graphdata.ptype = 0;
  graphdata.style = 0;
  KapaPrepPlot(kapa, 100, &graphdata);
  KapaPlotVector(kapa, 100, xCirc, "x");
  KapaPlotVector(kapa, 100, yCirc, "y");


  
  askUser(&plotOutliers);

  FREE(Rout);
  FREE(Dout);
  FREE(Rin);
  FREE(Din);

  return TRUE;
}
  
  
