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

/* 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 "pmSubtractionStamps.h"

#include "pmVisual.h"

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

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

//variables to determine when things are plotted
static bool isVisual             = false;
static bool plotConvKernels      = true;
static bool plotStamps           = true;
static bool plotLeastSquares     = true;
static bool plotImage            = true;
// variables to store plotting window indices
static int kapa  = -1;
static int kapa2 = -1;

/** function prototypes*/
static bool plotStampLocations(pmSubtractionStampList *stamps, pmReadout *ro);

// Initialization Routines

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


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

// Plotting Routines

/** Display images of the convolution kernels
 *  @param convKernels the kernels to plot
 *    @return true for success */
bool pmSubtractionVisualPlotConvKernels(psImage *convKernels) {
    if (!isVisual || !plotConvKernels) return true;
    if (!pmVisualInitWindow(&kapa, "ppSub:Images")) {
        isVisual = false;
        return false;
    }
    pmVisualScaleImage(kapa, convKernels, "Convolution_Kernels", 0, true);
    pmVisualAskUser(&plotConvKernels, &isVisual);
    return true;
}


/** Display the postage stamps used to determine the convolution kernels
    @param stamps to display
    @return true for success */
bool pmSubtractionVisualPlotStamps(pmSubtractionStampList *stamps, pmReadout *ro) {
    if (!isVisual || !plotStamps) return true;
    if (!pmVisualInitWindow (&kapa, "ppSub:Images")) {
        isVisual = false;
        return false;
    }
    if (!pmVisualInitWindow (&kapa2, "ppSub:StampMasterImage")) {
        isVisual = false;
        return false;
    }

    //Plot Location of stamps
    plotStampLocations(stamps, ro);

    //Find the stamp size
    int imageMax = -1;
    int numStamps = 0;
    psElemType type = PS_TYPE_F32; // will be overwritten
    for(int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i];
        if (stamp == NULL) continue;

        psKernel *k1 = stamp->image1;
        if (k1 == NULL) continue;

        psImage *i1 = k1->image;
        if (i1 == NULL) continue;

        imageMax = PS_MAX(imageMax, i1->numCols);
        imageMax = PS_MAX(imageMax, i1->numRows);
        numStamps++;
        type = i1->type.type;
    }
    if (imageMax == -1) return false;

    int border = 30;
    int tileRowCount = (int) ceil(sqrt(numStamps));
    int canvasX = tileRowCount * (imageMax * 2 + border);
    int canvasY = tileRowCount * (imageMax + border);
    psImage *canvas = psImageAlloc (canvasX, canvasY, type);
    psImageInit (canvas, NAN);

    //overlay the images
    int stampNum = 0;
    int stampListNum = 0;
    while (stampNum < numStamps) {
        int x0 = (2 * imageMax + border) * (stampNum % tileRowCount);
        int y0 = (imageMax + border) * (stampNum / tileRowCount);

        pmSubtractionStamp *stamp = stamps->stamps->data[stampListNum++];
        if (stamp == NULL) continue;

        psKernel *k1 = stamp->image1;
        psKernel *k2 = stamp->image2;

        if (k1 == NULL) continue;
        psImage *i1 = k1->image;
        psImage *i2 = k2->image;

        if (i1 == NULL) continue;
        psImageOverlaySection(canvas, i1, x0 + 0 * imageMax, y0, "=");
        psImageOverlaySection(canvas, i2, x0 + 1 * imageMax + border / 4, y0, "=");
        stampNum++;
    }
    pmVisualScaleImage(kapa, canvas, "Subtraction_Stamps", 0, true);

    pmVisualAskUser(&plotStamps, &isVisual);
    return true;
}

/** Plot the least-squares matrix of each stamp */
bool pmSubtractionVisualPlotLeastSquares (pmSubtractionStampList *stamps, bool dual) {

    if (!isVisual || !plotLeastSquares) return true;
    if (!pmVisualInitWindow (&kapa, "PPSub:Images")) {
        isVisual = false;
        return false;
    }

    //Find the stamp size
    int imageMax = -1;
    int numStamps = 0;
    psElemType type = PS_TYPE_F64;

    for(int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i];
        if (stamp == NULL) continue;

        psImage *im = stamp->matrix1;
        if (im == NULL) im = stamp->matrix2;
        if (im == NULL) im = stamp->matrixX;
        if (im == NULL) continue;

        imageMax = PS_MAX(imageMax, im->numCols);
        imageMax = PS_MAX(imageMax, im->numRows);
        numStamps++;
        type = im->type.type;
    }
    if (imageMax == -1) return false;

    int border = 15;
    imageMax += border;
    int tileRowCount = (int) ceil(sqrt(numStamps));
    int canvasX = tileRowCount * (imageMax);
    int canvasY = tileRowCount * (imageMax);
    psImage *canvas = psImageAlloc (canvasX, canvasY, type);
    psImageInit (canvas, NAN);

    //overlay the images
    int stampNum = 0;
    int stampListNum = 0;
    while (stampNum < numStamps) {
        int x0 = (imageMax) * (stampNum % tileRowCount);
        int y0 = (imageMax) * (stampNum / tileRowCount);

        pmSubtractionStamp *stamp = stamps->stamps->data[stampListNum++];
        if (stamp == NULL) continue;

        psImage *im = stamp->matrix1;
        if (im == NULL) im = stamp->matrix2;
        if (im == NULL) im = stamp->matrixX;
        if (im == NULL) continue;

        psImageOverlaySection(canvas, im, x0, y0, "=");
        stampNum++;
    }

    psImage *canvas32 = pmVisualImageToFloat(canvas);
    pmVisualScaleImage(kapa, canvas32, "Least_Squares", 0, true);

    pmVisualAskUser(&plotLeastSquares, &isVisual);
    psFree(canvas);
    psFree(canvas32);
    return true;
}

bool pmSubtractionVisualShowSubtraction(psImage *image, psImage *ref, psImage *sub) {
    if (!isVisual || !plotImage) return true;
    if (!pmVisualInitWindow (&kapa, "PPSub:Images")) {
        isVisual = false;
        return false;
    }

    pmVisualScaleImage(kapa, image, "Image", 0, true);
    pmVisualScaleImage(kapa, ref, "Reference", 1, true);
    pmVisualScaleImage(kapa, sub, "Subtraction", 2, true);
    pmVisualAskUser(&plotImage, &isVisual);
    return true;
}

static bool plotStampLocations(pmSubtractionStampList *stamps, pmReadout *ro) {

    if (!pmVisualScaleImage(kapa2, ro->image, "Stamp_master_image", 0, true)) {
        fprintf(stderr, "Cannot display postage stamp master image. Skipping \n");
        return false;
    }

    int Noverlay;
    KiiOverlay *overlay;

    // note: this uses the Ohana allocation tools:
    // ALLOCATE (overlay, KiiOverlay, 3*peaks->n + 1);
    ALLOCATE (overlay, KiiOverlay, stamps->num);

    Noverlay = 0;
    char boxID[10];
    for (int i = 0; i < stamps->num; i++) {

        pmSubtractionStamp *stamp = stamps->stamps->data[i];
        if (stamp == NULL) continue;

        overlay[Noverlay].type = KII_OVERLAY_BOX;
        overlay[Noverlay].x = stamp->x;
        overlay[Noverlay].y = stamp->y;
        overlay[Noverlay].dx = 40.0;
        overlay[Noverlay].dy = 40.0;
        overlay[Noverlay].angle = 0.0;
        sprintf(boxID, "%d", i);
        overlay[Noverlay].text = boxID;
        Noverlay ++;
    }

    KiiLoadOverlay (kapa2, overlay, Noverlay, "red");
    FREE (overlay);
    return true;
}

#else
bool pmSubtractionSetVisual (bool mode) {return true;}
bool pmSubtractionVisualClose() {return true;}
bool pmSubtractionVisualPlotConvKernels(psImage *convKernels) {return true;}
bool pmSubtractionVisualPlotStamps(pmSubtractionStampList *stamps, pmReadout *ro) {return true;}
bool pmSubtractionVisualPlotLeastSquares(pmSubtractionStampList *stamps) {return true;}
bool pmSubtractionVisualShowSubtraction(psImage *image, psImage *ref, psImage *sub) {return true;}
#endif
