#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <string.h>
#include <pslib.h>

// All these includes required to get stamps out of an array of pmSources
#include "pmMoments.h"
#include "pmPeaks.h"
#include "pmResiduals.h"
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmGrowthCurve.h"
#include "pmTrend2D.h"
#include "pmPSF.h"
#include "pmModel.h"
#include "pmSource.h"


#include "pmSubtraction.h"
#include "pmSubtractionStamps.h"

#define STAMP_LIST_BUFFER 20            // Number of stamps to add to list at a time

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Private (file-static) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Free function for pmSubtractionStampList
static void subtractionStampListFree(pmSubtractionStampList *list // Stamp list to free
                                     )
{
    psFree(list->stamps);
    psFree(list->regions);
    psFree(list->x);
    psFree(list->y);
    psFree(list->flux);
}

// Free function for pmSubtractionStamp
static void subtractionStampFree(pmSubtractionStamp *stamp // Stamp to free
                                 )
{
    psFree(stamp->image1);
    psFree(stamp->image2);
    psFree(stamp->weight);
    psFree(stamp->convolutions1);
    psFree(stamp->convolutions2);

    psFree(stamp->matrix1);
    psFree(stamp->matrix2);
    psFree(stamp->matrixX);
    psFree(stamp->vector1);
    psFree(stamp->vector2);

}

// Is this region OK?
static bool checkStampRegion(int x, int y, // Coordinates of stamp
                             const psRegion *region // Region of interest
                             )
{
    if (!region) {
        return true;
    }
    return (x < region->x0 || x > region->x1 || y < region->y0 || y > region->y1) ?
        false : true;
}

// Is this position unmasked?
static bool checkStampMask(int x, int y, // Coordinates of stamp
                           const psImage *mask, // Mask
                           pmSubtractionMode mode // Mode for subtraction
                           )
{
    if (!mask) {
        return true;
    }
    if (x < 0 || x >= mask->numCols || y < 0 || y >= mask->numRows) {
        return false;
    }

    psMaskType maskVal = PM_SUBTRACTION_MASK_BORDER | PM_SUBTRACTION_MASK_REJ; // Mask value
    switch (mode) {
      case PM_SUBTRACTION_MODE_1:
        maskVal |= PM_SUBTRACTION_MASK_FOOTPRINT_1;
        break;
      case PM_SUBTRACTION_MODE_2:
        maskVal |= PM_SUBTRACTION_MASK_FOOTPRINT_2;
        break;
      case PM_SUBTRACTION_MODE_UNSURE:
      case PM_SUBTRACTION_MODE_DUAL:
        maskVal |= PM_SUBTRACTION_MASK_FOOTPRINT_1 | PM_SUBTRACTION_MASK_FOOTPRINT_2;
        break;
      default:
        psAbort("Unsupported subtraction mode: %x", mode);
    }

    return (mask->data.PS_TYPE_MASK_DATA[y][x] & maskVal) ? false : true;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

pmSubtractionStampList *pmSubtractionStampListAlloc(int numCols, int numRows, const psRegion *region,
                                                    int footprint, float spacing)
{
    pmSubtractionStampList *list = psAlloc(sizeof(pmSubtractionStampList)); // Stamp list to return
    psMemSetDeallocator(list, (psFreeFunc)subtractionStampListFree);

    // Get region in which to find stamps: [xMin:xMax,yMin:yMax]
    int xMin = 0, xMax = numCols, yMin = 0, yMax = numRows;
    if (region) {
        xMin = PS_MAX(region->x0, xMin);
        xMax = PS_MIN(region->x1, xMax);
        yMin = PS_MAX(region->y0, yMin);
        yMax = PS_MIN(region->y1, yMax);
    }
    int xSize = xMax - xMin, ySize = yMax - yMin; // Size of region of interest
    int xStamps = (float)xSize / spacing + 1, yStamps = (float)ySize / spacing + 1; // Number of stamps

    list->num = xStamps * yStamps;
    list->stamps = psArrayAlloc(list->num);
    list->regions = psArrayAlloc(list->num);

    for (int y = 0, index = 0; y < yStamps; y++) {
        int yStart = yMin + y * ((float)ySize / (float)(yStamps)); // Subregion starts here
        int yStop = yMin + (y + 1) * ((float)ySize / (float)(yStamps)) - 1; // Subregion stops here
        assert(yStart >= yMin && yStop < yMax);

        for (int x = 0; x < xStamps; x++, index++) {
            int xStart = xMin + x * ((float)xSize / (float)(xStamps)); // Subregion starts here
            int xStop = xMin + (x + 1) * ((float)xSize / (float)(xStamps)) - 1; // Subregion stops here
            assert(xStart >= xMin && xStop < xMax);

            list->stamps->data[index] = pmSubtractionStampAlloc();
            psTrace("psModules.imcombine", 6, "Stamp region %d: [%d:%d,%d:%d]\n",
                    index, xStart, xStop, yStart, yStop);
            list->regions->data[index] = psRegionAlloc(xStart, xStop, yStart, yStop);
        }
    }

    list->x = NULL;
    list->y = NULL;
    list->flux = NULL;
    list->footprint = footprint;

    return list;
}

pmSubtractionStamp *pmSubtractionStampAlloc(void)
{
    pmSubtractionStamp *stamp = psAlloc(sizeof(pmSubtractionStamp)); // Stamp to return
    psMemSetDeallocator(stamp, (psFreeFunc)subtractionStampFree);

    stamp->x = NAN;
    stamp->y = NAN;
    stamp->flux = NAN;
    stamp->xNorm = NAN;
    stamp->yNorm = NAN;
    stamp->status = PM_SUBTRACTION_STAMP_INIT;

    stamp->image1 = NULL;
    stamp->image2 = NULL;
    stamp->weight = NULL;
    stamp->convolutions1 = NULL;
    stamp->convolutions2 = NULL;

    stamp->matrix1 = NULL;
    stamp->matrix2 = NULL;
    stamp->matrixX = NULL;
    stamp->vector1 = NULL;
    stamp->vector2 = NULL;

    return stamp;
}


pmSubtractionStampList *pmSubtractionStampsFind(pmSubtractionStampList *stamps, const psImage *image,
                                                const psImage *subMask, const psRegion *region,
                                                float threshold, int footprint, float spacing,
                                                pmSubtractionMode mode)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, NULL);
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(image, subMask, NULL);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_MASK, NULL);
    }
    PS_ASSERT_INT_NONNEGATIVE(footprint, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(spacing, 0.0, NULL);
    if (region) {
        if (psRegionIsNaN(*region)) {
            psString string = psRegionToString(*region);
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Input region (%s) contains NAN values", string);
            psFree(string);
            return false;
        }
        if (region->x0 < 0 || region->x1 > image->numCols ||
            region->y0 < 0 || region->y1 > image->numRows) {
            psString string = psRegionToString(*region);
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Input region (%s) does not fit in image (%dx%d)",
                    string, image->numCols, image->numRows);
            psFree(string);
            return false;
        }
    }

    int numRows = image->numRows, numCols = image->numCols; // Size of image

    if (!stamps) {
        stamps = pmSubtractionStampListAlloc(numCols, numRows, region, footprint, spacing);
    }

    int numStamps = stamps->num;        // Number of stamp regions
    int numFound = 0;                   // Number of stamps found
    int numValid = 0;                   // Number of valid regions
    for (int i = 0; i < numStamps; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest

        // Only find a new stamp if we need to
        if (stamp->status != PM_SUBTRACTION_STAMP_REJECTED && stamp->status != PM_SUBTRACTION_STAMP_INIT) {
            continue;
        }
        numValid++;

        float xStamp = 0, yStamp = 0;   // Coordinates of stamp
        float fluxStamp = NAN;          // Flux of stamp
        bool goodStamp = false;         // Found a good stamp?

        // A couple different ways of finding stamps:
        if (stamps->x && stamps->y) {
            // Get the next stamp from the list
            psVector *xList = stamps->x->data[i], *yList = stamps->y->data[i]; // Coordinate lists
            psVector *fluxList = stamps->flux->data[i]; // List of stamp fluxes

            // Take stamp off the top of the (sorted) list
            if (xList->n > 0) {
                int index = xList->n - 1; // Index of new stamp
                xStamp = xList->data.F32[index];
                yStamp = yList->data.F32[index];
                fluxStamp = fluxList->data.F32[index];

                // Chop off the top of the list
                xList->n = index;
                yList->n = index;
                fluxList->n = index;

                goodStamp = true;
            }
        } else {
            // Use a simple method of automatically finding stamps --- take the highest pixel in the subregion
            fluxStamp = threshold;
            psRegion *subRegion = stamps->regions->data[i]; // Sub-region of interest
            for (int y = subRegion->y0; y <= subRegion->y1; y++) {
                for (int x = subRegion->x0; x <= subRegion->x1; x++) {
                    if (checkStampMask(x, y, subMask, mode) && image->data.F32[y][x] > fluxStamp) {
                        fluxStamp = image->data.F32[y][x];
                        xStamp = x;
                        yStamp = y;
                        goodStamp = true;
                    }
                }
            }
        }

        if (goodStamp) {
            stamp->x = xStamp;
            stamp->y = yStamp;
            stamp->flux = fluxStamp;

            // Reset the postage stamps since we're making a new stamp
            psFree(stamp->image1);
            psFree(stamp->image2);
            psFree(stamp->weight);
            psFree(stamp->convolutions1);
            psFree(stamp->convolutions2);
            stamp->image1 = stamp->image2 = stamp->weight = NULL;
            stamp->convolutions1 = stamp->convolutions2 = NULL;

            stamp->status = PM_SUBTRACTION_STAMP_FOUND;
            numFound++;
            psTrace("psModules.imcombine", 5, "Found stamp in subregion %d: %d,%d\n",
                    i, (int)stamp->x, (int)stamp->y);
        } else {
            stamp->status = PM_SUBTRACTION_STAMP_NONE;
        }
    }

    if (numValid > 0) {
        psLogMsg("psModules.imcombine", PS_LOG_INFO, "Found %d stamps", numFound);
    }

    return stamps;
}


pmSubtractionStampList *pmSubtractionStampsSet(const psVector *x, const psVector *y, const psVector *flux,
                                               const psImage *image, const psImage *subMask,
                                               const psRegion *region, int footprint, float spacing,
                                               pmSubtractionMode mode)

{
    PS_ASSERT_VECTOR_NON_NULL(x, NULL);
    PS_ASSERT_VECTOR_TYPE(x, PS_TYPE_F32, NULL);
    PS_ASSERT_VECTOR_NON_NULL(y, NULL);
    PS_ASSERT_VECTOR_TYPE(y, PS_TYPE_F32, NULL);
    PS_ASSERT_VECTORS_SIZE_EQUAL(y, x, NULL);
    if (flux) {
        PS_ASSERT_VECTOR_NON_NULL(flux, NULL);
        PS_ASSERT_VECTOR_TYPE(flux, PS_TYPE_F32, NULL);
        PS_ASSERT_VECTORS_SIZE_EQUAL(flux, x, NULL);
    } else {
        PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    }
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, NULL);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_MASK, NULL);
        if (image) {
            PS_ASSERT_IMAGE_NON_NULL(image, NULL);
            PS_ASSERT_IMAGES_SIZE_EQUAL(image, subMask, NULL);
        }
    }
    PS_ASSERT_FLOAT_LARGER_THAN(spacing, 0.0, NULL);

    int numStars = x->n;                // Number of stars
    pmSubtractionStampList *stamps = pmSubtractionStampListAlloc(subMask->numCols, subMask->numRows,
                                                                 region, footprint, spacing); // Stamp list
    int numStamps = stamps->num;        // Number of stamps

    // Initialise the lists
    stamps->x = psArrayAlloc(numStamps);
    stamps->y = psArrayAlloc(numStamps);
    stamps->flux = psArrayAlloc(numStamps);
    for (int i = 0; i < numStamps; i++) {
        stamps->x->data[i] = psVectorAllocEmpty(STAMP_LIST_BUFFER, PS_TYPE_F32);
        stamps->y->data[i] = psVectorAllocEmpty(STAMP_LIST_BUFFER, PS_TYPE_F32);
        stamps->flux->data[i] = psVectorAllocEmpty(STAMP_LIST_BUFFER, PS_TYPE_F32);
    }

    // Put the stars into their appropriate subregions
    for (int i = 0; i < numStars; i++) {
        float xStamp = x->data.F32[i], yStamp = y->data.F32[i]; // Coordinates of stamp
        int xPix = xStamp + 0.5, yPix = yStamp + 0.5; // Pixel coordinate of stamp
        if (!checkStampRegion(xPix, yPix, region)) {
            // It's not in the big region
            psTrace("psModules.imcombine", 9, "Rejecting input stamp (%d,%d) because outside region",
                    xPix, yPix);
            continue;
        }
        if (!checkStampMask(xPix, yPix, subMask, mode)) {
            // Not a good stamp
            psTrace("psModules.imcombine", 9, "Rejecting input stamp (%d,%d) because bad mask",
                    xPix, yPix);
            continue;
        }

        bool found = false;
        for (int j = 0; j < numStamps && !found; j++) {
            psRegion *subRegion = stamps->regions->data[j]; // Subregion of interest
            if (checkStampRegion(xPix, yPix, subRegion)) {
                psVector *xList = stamps->x->data[j], *yList = stamps->y->data[j]; // Pixel lists
                psVector *fluxList = stamps->flux->data[j]; // Flux list

                int index = xList->n;   // Index of new stamp candidate

                psVectorExtend(xList, STAMP_LIST_BUFFER, 1);
                psVectorExtend(yList, STAMP_LIST_BUFFER, 1);
                psVectorExtend(fluxList, STAMP_LIST_BUFFER, 1);

                xList->data.F32[index] = xStamp;
                yList->data.F32[index] = yStamp;

                if (flux) {
                    fluxList->data.F32[index] = flux->data.F32[i];
                } else {
                    fluxList->data.F32[index] = image->data.F32[yPix][xPix];
                }

                found = true;
                psTrace("psModules.imcombine", 9, "Putting input stamp (%d,%d) into subregion %d",
                        xPix, yPix, j);
            }
        }

        if (!found) {
            psTrace("psModules.imcombine", 9, "Unable to find subregion for stamp (%d,%d)",
                    xPix, yPix);
        }
    }

    // Sort the list by flux, with the brightest last
    for (int i = 0; i < numStamps; i++) {
        psVector *xList = stamps->x->data[i], *yList = stamps->y->data[i]; // Pixel lists
        psVector *fluxList = stamps->flux->data[i]; // Flux list

        psVector *indexes = psVectorSortIndex(NULL, fluxList); // Indices to sort flux
        int num = indexes->n;           // Number of candidate stamps in this subregion

        psVector *xSorted = psVectorAlloc(num, PS_TYPE_F32); // Sorted version of x list
        psVector *ySorted = psVectorAlloc(num, PS_TYPE_F32); // Sorted version of y list
        psVector *fluxSorted = psVectorAlloc(num, PS_TYPE_F32); // Sorted version of flux list
        for (int j = 0; j < num; j++) {
            int k = indexes->data.S32[j]; // Sorted index
            xSorted->data.F32[j] = xList->data.F32[k];
            ySorted->data.F32[j] = yList->data.F32[k];
            fluxSorted->data.F32[j] = fluxList->data.F32[k];
        }
        psFree(indexes);

        psFree(stamps->x->data[i]);
        psFree(stamps->y->data[i]);
        psFree(stamps->flux->data[i]);

        stamps->x->data[i] = xSorted;
        stamps->y->data[i] = ySorted;
        stamps->flux->data[i] = fluxSorted;
    }


    return stamps;
}


bool pmSubtractionStampsExtract(pmSubtractionStampList *stamps, psImage *image1, psImage *image2,
                                psImage *weight, int kernelSize)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PS_ASSERT_IMAGE_NON_NULL(image1, false);
    PS_ASSERT_IMAGE_TYPE(image1, PS_TYPE_F32, false);
    if (image2) {
        PS_ASSERT_IMAGE_NON_NULL(image2, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(image2, image1, false);
        PS_ASSERT_IMAGE_TYPE(image2, PS_TYPE_F32, false);
    }
    PS_ASSERT_IMAGE_NON_NULL(weight, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(weight, image1, false);
    PS_ASSERT_IMAGE_TYPE(weight, PS_TYPE_F32, false);
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);

    int numCols = image1->numCols, numRows = image1->numRows; // Size of images
    int size = kernelSize + stamps->footprint; // Size of postage stamps

    for (int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
        if (!stamp || stamp->status != PM_SUBTRACTION_STAMP_FOUND) {
            continue;
        }

        if (isnan(stamp->xNorm)) {
            stamp->xNorm = 2.0 * (stamp->x - (float)numCols/2.0) / (float)numCols;
        }
        if (isnan(stamp->yNorm)) {
            stamp->yNorm = 2.0 * (stamp->y - (float)numRows/2.0) / (float)numRows;
        }

        int x = stamp->x + 0.5, y = stamp->y + 0.5; // Stamp coordinates
        if (x < size || x > numCols - size || y < size || y > numRows - size) {
            psError(PS_ERR_UNKNOWN, false, "Stamp %d (%d,%d) is within the image border.\n", i, x, y);
            return false;
        }

        // Catch memory leaks --- these should have been freed and NULLed before
        assert(stamp->image1 == NULL);
        assert(stamp->image2 == NULL);
        assert(stamp->weight == NULL);

        psRegion region = psRegionSet(x - size, x + size + 1, y - size, y + size + 1); // Region of interest

        psImage *sub1 = psImageSubset(image1, region); // Subimage with stamp
        stamp->image1 = psKernelAllocFromImage(sub1, size, size);
        psFree(sub1);                   // Drop reference

        if (image2) {
            psImage *sub2 = psImageSubset(image2, region); // Subimage with stamp
            stamp->image2 = psKernelAllocFromImage(sub2, size, size);
            psFree(sub2);               // Drop reference
        }

        psImage *wtSub = psImageSubset(weight, region); // Subimage with stamp
        stamp->weight = psKernelAllocFromImage(wtSub, size, size);
        psFree(wtSub);                  // Drop reference

        stamp->status = PM_SUBTRACTION_STAMP_CALCULATE;
    }

    return true;
}

#if 0
bool pmSubtractionStampsGenerate(pmSubtractionStampList *stamps, float fwhm, int kernelSize)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PS_ASSERT_FLOAT_LARGER_THAN(fwhm, 0.0, false);
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);

    int size = kernelSize + stamps->footprint; // Size of postage stamps
    int num = stamps->num;              // Number of stamps
    float sigma = fwhm / (2.0 * sqrtf(2.0 * log(2.0))); // Gaussian sigma

    for (int i = 0; i < num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
        if (!(stamp->status & PM_SUBTRACTION_STAMP_CALCULATE)) {
            continue;
        }

        float x = stamp->x, y = stamp->y; // Coordinates of stamp
        float flux = stamp->flux; // Flux of star
        if (!isfinite(flux)) {
            psWarning("Unable to generate PSF for stamp %d --- bad flux.", i);
            stamp->status = PM_SUBTRACTION_STAMP_REJECTED;
            continue;
        }

        float xStamp = x - (int)(x + 0.5); // x coordinate of star in stamp frame
        float yStamp = y - (int)(y + 0.5); // y coordinate of star in stamp frame

        psFree(stamp->image2);
        stamp->image2 = psKernelAlloc(-size, size, -size, size);
        psKernel *target = stamp->image2; // Target stamp

        // Put in a Waussian, just for fun!
        for (int v = -size; v <= size; v++) {
            for (int u = -size; u <= size; u++) {
                float z = (PS_SQR(u + xStamp) + PS_SQR(v + yStamp)) / (2.0 * PS_SQR(sigma));
                target->kernel[v][u] = flux / sigma * 0.5 * M_2_SQRTPI * M_SQRT1_2 / (1.0 + z + PS_SQR(z));
            }
        }

    }

    return true;
}
#endif

pmSubtractionStampList *pmSubtractionStampsSetFromSources(const psArray *sources, const psImage *subMask,
                                                          const psRegion *region, int footprint,
                                                          float spacing, pmSubtractionMode mode)
{
    PS_ASSERT_ARRAY_NON_NULL(sources, NULL);
    // Let pmSubtractionStampsSet take care of the rest of the assertions

    int numSources = sources->n;          // Number of stars

    psVector *x = psVectorAlloc(numSources, PS_TYPE_F32); // x coordinates
    psVector *y = psVectorAlloc(numSources, PS_TYPE_F32); // y coordinates
    psVector *flux = psVectorAlloc(numSources, PS_TYPE_F32); // Fluxes

    for (int i = 0; i < numSources; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (source->modelPSF) {
            x->data.F32[i] = source->modelPSF->params->data.F32[PM_PAR_XPOS];
            y->data.F32[i] = source->modelPSF->params->data.F32[PM_PAR_YPOS];
        } else {
            x->data.F32[i] = source->peak->xf;
            y->data.F32[i] = source->peak->yf;
        }
        flux->data.F32[i] = powf(10.0, -0.4 * source->psfMag);
    }

    pmSubtractionStampList *stamps = pmSubtractionStampsSet(x, y, flux, NULL, subMask, region,
                                                            footprint, spacing, mode); // Stamps to return
    psFree(x);
    psFree(y);
    psFree(flux);

    if (!stamps) {
        psError(PS_ERR_UNKNOWN, false, "Unable to set stamps from sources.");
    }

    return stamps;
}


pmSubtractionStampList *pmSubtractionStampsSetFromFile(const char *filename, const psImage *image,
                                                       const psImage *subMask, const psRegion *region,
                                                       int footprint, float spacing, pmSubtractionMode mode)
{
    PS_ASSERT_STRING_NON_EMPTY(filename, NULL);
    // Let pmSubtractionStampsSet take care of the rest of the assertions

    psArray *data = psVectorsReadFromFile(filename, "%f %f");
    if (!data) {
        psError(PS_ERR_IO, false, "Unable to read stamps file %s", filename);
        return NULL;
    }
    psVector *x = data->data[0], *y = data->data[1]; // Stamp positions

    // Correct for IRAF/FITS (unit-offset) positions to C (zero-offset) positions
    psBinaryOp(x, x, "-", psScalarAlloc(1.0, PS_TYPE_F32));
    psBinaryOp(y, y, "-", psScalarAlloc(1.0, PS_TYPE_F32));

    pmSubtractionStampList *stamps = pmSubtractionStampsSet(x, y, NULL, image, subMask, region, footprint,
                                                            spacing, mode);
    psFree(data);

    return stamps;

}
