/** Diagnostic plots called from within pmAstrometry routines.
 * @author      Chris Beaumont
 * @date        October 2008
 */

/* Include Files   */
#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <strings.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include <pslib.h>

#include "pmKapaPlots.h"

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmAstrometryObjects.h"

# if (HAVE_KAPA)
# include <kapa.h>

# define KAPAX 700
# define KAPAY 700


//variables to determine when things are plotted
static bool isVisual             = false;
static bool plotGridMatch        = true;
static bool plotTweak            = true;

// variables to store plotting window indices
static int kapa = -1;


/* Utility Routine Prototypes */
bool pmAstromVisualScaleGraphdata(Graphdata *graphdata, psVector *xVec, psVector *yVec, bool clip);


/* Initialization Routines  */


/** start or stop plotting */
bool pmAstromSetVisual (bool mode) {
    isVisual = mode;
    return true;
}


/** open name, and resize a window if necessary */
bool pmAstromVisualInitWindow (int *kapid, char *name) {
    if (*kapid == -1) {
        *kapid = KapaOpenNamedSocket("kapa", name);
        if (*kapid == -1) {
            fprintf (stderr, "failure to open kapa; visual mode disabled for pmAstrom\n");
            isVisual = false;
            return false;
        }
        KapaResize (*kapid, KAPAX, KAPAY);
    }
    return true;
}


/** ask the user how to proceed */
bool pmAstromVisualAskUser(bool *plotflag)
{
    char key[10];
    fprintf (stdout, "[c]ontinue? [s]kip the rest of these plots? [a]bort all visual plots?");
    if (!fgets(key, 8, stdin)) {
        psWarning("Unable to read option");
    }
    if (key[0] == 's') {
        *plotflag = false;
    }
    if (key[0] == 'a') {
        pmAstromSetVisual(false);
    }
    return true;
}


/** destroy windows at the end of a run*/
bool pmAstromVisualClose()
{
    if(kapa != -1)
        KiiClose(kapa);
    return true;
}


/* plotting routines */

/**
 * Plot the offset between every pair of reference and raw source locations. The peak of this
 * distribution nominally gives the offset, scale difference, and rotation of the two catalogs.
 * Overplots the location of this peak as determined by pmAstromGridAngle, as well as some profiles
 * along horizontal and vertical cuts through this peak.
 */
bool pmAstromVisualPlotGridMatch (const psArray *raw, ///< raw stars
                                  const psArray *ref, ///< reference stars
                                  psImage *gridNP,    ///< a 2D histogram of raw-ref star distances
                                  double offsetX,     ///< The X location (FP coordinates) of the peak of gridNP
                                  double offsetY,     ///< the Y location (FP coordinates) of the peak of gridNP
                                  double maxOffpix,   ///< The half-width of gridNP in FP coordinates
                                  double Scale,       ///< The pixel size of gridNP in histogram-bin-coordinates
                                  double Offset       ///< The (x,y) location (histogram-bin coordinates) of the FP point (0,0) in gridNP
                                  )
{
    //make sure we want to plot this
    if (!isVisual || !plotGridMatch) return true;
    if (!pmAstromVisualInitWindow(&kapa, "pmAstrom:plots"))
        return false;

    KapaSection section = {"s1", 0.05, 0.05, .75, .75};
    KapaSection sectionY = {"s2", 0.8, 0.05, .15, .75};
    KapaSection sectionX = {"s3", .05, .8, .75, .15};

    Graphdata graphdata;
    int nplot = raw->n * ref->n;                      // number of points to plot
    float dXplot[nplot];                              // x data points
    float dYplot[nplot];                              // y data points
    pmAstromObj *ob1; pmAstromObj *ob2;               // shortcuts to the data in raw and ref
    psU32 **NP = gridNP->data.U32;                    // shortcut to the gridNP data
    float vertHistSlice[gridNP->numRows];             // vertical histogram slice through peak
    float horizHistSlice[gridNP->numCols];            // horizontal histogram slice through peak
    float horizontalIndices[gridNP->numCols];         // the horizontal offset corresponding to each bin of horizHistSlice
    float verticalIndices[gridNP->numRows];           // the vertical offset corresponding to each bin of vertHistSlice
    int maxHorizontalSlice = 0;                       // peak value of horizHistSlice
    int maxVerticalSlice = 0;                         // peak value of vertHistSlice
    int peakXbin = (int) (offsetX / Scale + Offset);  // X bin index of peak
    int peakYbin = (int) (offsetY / Scale + Offset);  // Y bin index of peak


    // set up plot information
    KapaClearPlots(kapa);
    KapaInitGraph(&graphdata);
    KapaSetSection(kapa, &section);

    graphdata.xmin = -1.0 * maxOffpix;
    graphdata.xmax =  1.0 * maxOffpix;
    graphdata.ymin = -1.0 * maxOffpix;
    graphdata.ymax =  1.0 * maxOffpix;
    KapaSetLimits(kapa, &graphdata);

    KapaSetFont(kapa, "helvetica", 14);
    KapaBox(kapa, &graphdata);
    KapaSendLabel (kapa, "X offset (FP)", KAPA_LABEL_XM);
    KapaSendLabel (kapa, "Y offset (FP)", KAPA_LABEL_YM);
    KapaSendLabel (kapa, "pmAstromGridAngle residuals. Box: Correlation Peak.",
                   KAPA_LABEL_XP);
    graphdata.style = 2;
    graphdata.ptype = 0;
    graphdata.size = 0.4;
    graphdata.color = KapaColorByName ("black");

    // calculate the plot points
    float dX, dY;
    for (int i = 0; i < raw->n; i++) {
        ob1 = (pmAstromObj *)raw->data[i];
        for (int j = 0; j < ref->n; j++) {
            ob2 = (pmAstromObj *)ref->data[j];
            dX = ob1->FP->x - ob2->FP->x;
            dY = ob1->FP->y - ob2->FP->y;
            dXplot[(i * ref->n) + j] = dX;
            dYplot[(i * ref->n) + j] = dY;
        }
    }

    // calculate the points for the profiles
    for (int i = 0; i < gridNP->numRows; i++) {
        vertHistSlice[i] = NP[i][peakXbin];
        verticalIndices[i] = (i - Offset) * Scale;
        if (vertHistSlice[i] > maxVerticalSlice) {
            maxVerticalSlice = vertHistSlice[i];
        }
    }
    for (int i = 0; i < gridNP->numCols; i++) {
        horizHistSlice[i] = NP[peakYbin][i];
        horizontalIndices[i] = (i - Offset) * Scale;
        if (horizHistSlice[i] > maxHorizontalSlice) {
            maxHorizontalSlice = horizHistSlice[i];
        }
    }

    //Plot the offsets
    KapaPrepPlot(kapa, nplot, &graphdata);
    KapaPlotVector (kapa, nplot, dXplot, "x");
    KapaPlotVector (kapa, nplot, dYplot, "y");

    //Overplot bounding box, peak of distribution
    float xbound[5] = { -maxOffpix, maxOffpix, maxOffpix, -maxOffpix, -maxOffpix};
    float ybound[5] = { -maxOffpix, -maxOffpix, maxOffpix, maxOffpix, -maxOffpix};
    float xbin[5] = {offsetX - 0.5 * Scale, offsetX + 0.5 * Scale,
                     offsetX + 0.5 * Scale, offsetX - 0.5 * Scale,
                     offsetX - 0.5 * Scale};
    float ybin[5] = {offsetY - 0.5 * Scale, offsetY - 0.5 * Scale,
                     offsetY + 0.5 * Scale, offsetY + 0.5 * Scale,
                     offsetY - 0.5 * Scale};
    graphdata.color = KapaColorByName("red");
    graphdata.style = 0;
    graphdata.size = 1.0;
    KapaPrepPlot(kapa, 5, &graphdata);
    KapaPlotVector (kapa, 5, xbound, "x");
    KapaPlotVector (kapa, 5, ybound, "y");
    KapaPrepPlot(kapa, 5, &graphdata);
    KapaPlotVector (kapa, 5, xbin, "x");
    KapaPlotVector (kapa, 5, ybin, "y");

    //plot X profile
    KapaSetSection(kapa, &sectionX);
    graphdata.color = KapaColorByName("black");
    graphdata.ptype = 1;
    graphdata.style = 1;
    graphdata.ymin = 0;
    graphdata.ymax = maxHorizontalSlice + 0.5;
    KapaSetLimits(kapa, &graphdata);

    KapaBox(kapa, &graphdata);
    KapaPrepPlot(kapa, gridNP->numCols, &graphdata);
    KapaPlotVector (kapa, gridNP->numCols, horizontalIndices, "x");
    KapaPlotVector (kapa, gridNP->numCols, horizHistSlice, "y");
    float xslice[2] = {offsetX - Scale / 2., offsetX - Scale / 2.};
    float yslice[2] = {-5, 100};
    graphdata.color = KapaColorByName("red");
    KapaPrepPlot(kapa, 2, &graphdata);
    KapaPlotVector (kapa, 2, xslice, "x");
    KapaPlotVector (kapa, 2, yslice, "y");

    //plot Y profile
    KapaSetSection(kapa, &sectionY);
    graphdata.color = KapaColorByName("black");
    graphdata.ptype = 1;
    graphdata.style = 1;
    graphdata.ymin = -maxOffpix;
    graphdata.ymax = maxOffpix;
    graphdata.xmin = -1.0 ;
    graphdata.xmax = maxVerticalSlice + 0.5;
    KapaSetLimits(kapa, &graphdata);

    KapaBox(kapa, &graphdata);
    KapaPrepPlot(kapa, gridNP->numRows, &graphdata);
    KapaPlotVector (kapa, gridNP->numRows, vertHistSlice, "x");
    KapaPlotVector (kapa, gridNP->numRows, verticalIndices, "y");
    yslice[0] = yslice[1] = offsetY - Scale / 2.;
    xslice[0] = -5; xslice[1] = 100;
    graphdata.color = KapaColorByName("red");
    KapaPrepPlot(kapa, 2, &graphdata);
    KapaPlotVector (kapa, 2, xslice, "x");
    KapaPlotVector (kapa, 2, yslice, "y");

    pmAstromVisualAskUser(&plotGridMatch);
    return true;
} // end of pmAstromVisualPlotGridMatch


/** Plot the refinements made within pmAstromGridTweak.
 * After pmAstromGridMatch finds the best rotaion/scale/offset between raw and reference stars
 * within a coarse grid of rotations/scales, pmAstromGridTweak computes a higher precision
 * estimate of the offset. It computes the 2 point correlation function between raw and ref
 * stars along horizontal and vertical cuts through the first-guess offset. It finds the peak
 * of these two profiles and adjusts the offset accordingly. This procedure plots the profiles.
 */
bool pmAstromVisualPlotTweak (psVector *xHist, ///< Smoothed Horizontal cut through the histogram
                              psVector *yHist, ///< Smoothed Vertical cut throug the histogram
                              int xBin,        ///< X Bin index of the histogram peak
                              int yBin         ///< Y bin index of the histogram peak
    )
{
    //make sure we want to plot this
    if (!isVisual || !plotTweak) return true;
    if (!pmAstromVisualInitWindow(&kapa, "pmAstrom:plots")) {
        return false;
    }

    Graphdata graphdata;
    KapaSection section1 = {"s1", 0.05, 0.05, 0.90, 0.4};
    KapaSection section2 = {"s2", 0.05, 0.5, 0.90, 0.4};

    psVector *xIndices = psVectorAlloc (xHist->n, PS_TYPE_F32);
    psVector *yIndices = psVectorAlloc (yHist->n, PS_TYPE_F32);

    //populate the Indices vectors
    for(int i = 0; i < xHist->n; i++) {
        xIndices->data.F32[i] = i;
    }
    for(int i = 0; i < yHist->n; i++) {
        yIndices->data.F32[i] = i;
    }

    // set up plot information
    KapaClearPlots(kapa);
    KapaInitGraph(&graphdata);

    // plot the X histogram
    pmAstromVisualScaleGraphdata(&graphdata, xIndices, xHist, false);
    KapaSetSection(kapa, &section1);
    KapaSetLimits (kapa, &graphdata);
    KapaSetFont(kapa, "helvetica", 14);
    KapaBox(kapa, &graphdata);
    KapaSendLabel (kapa, "X offset Bin", KAPA_LABEL_XM);
    KapaSendLabel (kapa, "Number of Sources", KAPA_LABEL_YM);
    KapaSendLabel (kapa, "Horizontal Profile",
                   KAPA_LABEL_XP);
    graphdata.style = 1;
    graphdata.ptype = 0;
    graphdata.size = 0.4;
    graphdata.color = KapaColorByName ("black");

    KapaPrepPlot (kapa, xHist->n, &graphdata);
    KapaPlotVector (kapa, xHist->n, xIndices->data.F32, "x");
    KapaPlotVector (kapa, xHist->n, xHist->data.F32, "y");

    //overplot the peak
    float x[2] = {xBin, xBin};
    float y[2] = {-500, 500};
    graphdata.color = KapaColorByName ("red");
    KapaPrepPlot (kapa, 2, &graphdata);
    KapaPlotVector (kapa, 2, x, "x");
    KapaPlotVector (kapa, 2, y, "y");

    //plot the Y histogram
    pmAstromVisualScaleGraphdata(&graphdata, yIndices, yHist, false);
    KapaSetSection(kapa, &section2);
    KapaSetLimits (kapa, &graphdata);
    KapaSetFont(kapa, "helvetica", 14);
    graphdata.color = KapaColorByName ("black");
    KapaBox(kapa, &graphdata);
    KapaSendLabel (kapa, "Y offset Bin", KAPA_LABEL_XM);
    KapaSendLabel (kapa, "Number of Sources", KAPA_LABEL_YM);
    KapaSendLabel (kapa, "Vertical Profile",
                   KAPA_LABEL_XP);
    graphdata.style = 1;
    graphdata.ptype = 0;
    graphdata.size = 0.4;

    KapaPrepPlot (kapa, yHist->n, &graphdata);
    KapaPlotVector (kapa, yHist->n, yIndices->data.F32, "x");
    KapaPlotVector (kapa, yHist->n, yHist->data.F32, "y");

    //overplot the peak
    x[0] = x[1] = yBin;
    graphdata.color = KapaColorByName ("red");
    KapaPrepPlot (kapa, 2, &graphdata);
    KapaPlotVector (kapa, 2, x, "x");
    KapaPlotVector (kapa, 2, y, "y");

    //plot title
    graphdata.color = KapaColorByName("black");
    KapaSection section3 = {"s3", 0, 0, 1, 1};
    KapaSetSection( kapa, &section3);
    KapaSendLabel (kapa, "Tweaking the Astrometry Grid Solution. Smoothed profiles + peak location",
                   KAPA_LABEL_XP);

    pmAstromVisualAskUser(&plotTweak);

    psFree(xIndices);
    psFree(yIndices);
    return true;
} //end of pmAstromPlotTweak


/** pmAstromScaleGraphdata
 * Scale the graphdata structure based on x and y coordinates. Use sigma clipping to
 * prevent outliers from making te plot region too big.
 */
bool pmAstromVisualScaleGraphdata(Graphdata *graphdata, psVector *xVec, psVector *yVec, bool clip) {

    graphdata->xmin = +FLT_MAX;
    graphdata->xmax = -FLT_MAX;
    graphdata->ymin = +FLT_MAX;
    graphdata->ymax = -FLT_MAX;

    //determine standard deviation of xVec and yVec
    psStats *statsX = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);
    psStats *statsY = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);
    psVectorStats (statsX, xVec, NULL, NULL, 0);
    psVectorStats (statsY, yVec, NULL, NULL, 0);

    float xhi  = statsX->sampleMedian + 3 *statsX->sampleStdev;
    float xlo = statsX->sampleMedian - 3 *statsX->sampleStdev;
    float yhi = statsY->sampleMedian + 3 *statsY->sampleStdev;
    float ylo = statsY->sampleMedian - 3 *statsY->sampleStdev;

    // don't sigma clip
    if (!clip) {
        xhi = +FLT_MAX;
        xlo = -FLT_MAX;
        yhi = +FLT_MAX;
        ylo = -FLT_MAX;
    }

    // abort if there is no good data
    if (!isfinite(xhi) || !isfinite(xlo) || !isfinite(yhi) || !isfinite(ylo)) {
        graphdata->xmin = -1;
        graphdata->ymin  = -1;
        graphdata->xmax = 1;
        graphdata->ymax = 1;
        psFree(statsX);
        psFree(statsY);
        return false;
    }

    for(int i = 0; i < xVec->n; i++) {
        if (!isfinite(xVec->data.F32[i])) continue;
        if (xVec->data.F32[i] > xhi || xVec->data.F32[i] < xlo) continue;
        graphdata->xmin = PS_MIN (graphdata->xmin, xVec->data.F32[i]);
        graphdata->xmax = PS_MAX (graphdata->xmax, xVec->data.F32[i]);
    }

    for (int i = 0; i < yVec->n; i++) {
        if (!isfinite(xVec->data.F32[i])) continue;
        if (yVec->data.F32[i] > yhi || yVec->data.F32[i] < ylo) continue;
        graphdata->ymin = PS_MIN (graphdata->ymin, yVec->data.F32[i]);
        graphdata->ymax = PS_MAX (graphdata->ymax, yVec->data.F32[i]);
    }

    // add a whitespace border
    float range = graphdata->xmax - graphdata->xmin;
    if (range == 0) range = 1;
    graphdata->xmin -= .05 * range;
    graphdata->xmax += .05 * range;

    range = graphdata->ymax - graphdata->ymin;
    if (range == 0) range = 1;
    graphdata->ymin -= .05 * range;
    graphdata->ymax += .05 * range;

    psFree (statsX);
    psFree (statsY);
    return true;
}


# else

bool pmAstromSetVisual(bool mode) { return true; }
bool pmAstromVisualInitWindow (int *kapid, char *name) { return true; }
bool pmAstromVisualAskUser (bool *plotflag) { return true; }
bool pmAstromVisualClose() { return true; }
bool pmAstromVisualPlotGridMatch (const psArray *raw, const psArray *ref, psImage *gridNP, double offsetX, double offsetY, double maxOffpix, double Scale, double Offset) { return true; }
bool pmAstromVisualPlotTweak (psVector *xHist, psVector *yHist, int xBin, int yBin) {return true;}

# endif
