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

#include <stdio.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->matrix);
    psFree(stamp->vector);
    psFree(stamp->reference);
    psFree(stamp->input);
    psFree(stamp->weight);
    psFree(stamp->convolutions);
}

// 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
                           )
{
    if (!mask) {
        return true;
    }
    return (mask->data.PS_TYPE_MASK_DATA[y][x] & (PM_SUBTRACTION_MASK_BORDER |
                                                  PM_SUBTRACTION_MASK_FOOTPRINT | PM_SUBTRACTION_MASK_REJ)) ?
        false : true;
}

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

pmSubtractionStampList *pmSubtractionStampListAlloc(int numCols, int numRows, const psRegion *region,
                                                    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 = xSize / spacing + 1, yStamps = ySize / spacing + 1; // Number of stamps in x and y

    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 * ySize / (yStamps + 1); // Subregion starts here
        int yStop = yMin + (y + 1) * ySize / (yStamps + 1) - 1; // Subregion stops here
        assert(yStart >= yMin && yStop < yMax);

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

            list->stamps->data[index] = pmSubtractionStampAlloc();
            list->regions->data[index] = psRegionAlloc(xStart, xStop, yStart, yStop);
        }
    }

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

    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->matrix = NULL;
    stamp->vector = NULL;
    stamp->status = PM_SUBTRACTION_STAMP_INIT;

    stamp->reference = NULL;
    stamp->input = NULL;
    stamp->weight = NULL;
    stamp->convolutions = NULL;

    return stamp;
}


pmSubtractionStampList *pmSubtractionStampsFind(pmSubtractionStampList *stamps, const psImage *image,
                                                const psImage *subMask, const psRegion *region,
                                                float threshold, float spacing)
{
    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_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, 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->y1 ; x++) {
                    if (checkStampMask(x, y, subMask) && 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->reference);
            psFree(stamp->input);
            psFree(stamp->weight);
            stamp->reference = stamp->input = stamp->weight = NULL;

            stamp->status = PM_SUBTRACTION_STAMP_CALCULATE;
            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, float spacing, int exclusionZone)
{
    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);
    PS_ASSERT_INT_NONNEGATIVE(exclusionZone, NULL);

    int numStars = x->n;                // Number of stars
    psVector *exclude = psVectorAlloc(numStars, PS_TYPE_U8); // Exclude a star?
    psVectorInit(exclude, 0);

    // Apply exclusion zone around each star --- important when we're convolving to a specified PSF; less so
    // when we're trying to get a good subtraction.
    if (exclusionZone > 0) {
        psTrace("psModules.imcombine", 2, "Applying exclusion zone of %d pixels for stamps\n", exclusionZone);
        // We use something resembling a binary search --- coordinates are sorted in the x dimension, and then
        // to exclude stars within a nominated distance, we need only examine (i.e., calculate the
        // 2-dimensional distance to) other stars in the list that are within that distance of the x
        // coordinate of the star of interest.  By marking both stars when we find a violation of the
        // exclusion zone, we need only search upwards from the x coordinate of the star of interest.

        int minDistance2 = PS_SQR(exclusionZone); // Minimum square distance for other stars
        psVector *indexes = psVectorSortIndex(NULL, x); // Indices for sorting in x
        for (int i = 0; i < numStars - 1; i++) {
            int iIndex = indexes->data.S32[i]; // Index for star i
            if (exclude->data.U8[iIndex]) {
                continue;
            }
            float ix = x->data.F32[iIndex], iy = y->data.F32[iIndex]; // Coordinates for star i
            float jx, jy;                   // Coordinates for star j
            for (int j = i + 1, jIndex = indexes->data.S32[j];
                 j < numStars && (jx = x->data.F32[jIndex]) < ix + exclusionZone;
                 j++, jIndex = indexes->data.S32[j]) {
                jy = y->data.F32[jIndex];
                if (PS_SQR(ix - jx) + PS_SQR(iy - jy) < minDistance2) {
                    exclude->data.U8[iIndex] = 0xff;
                    exclude->data.U8[jIndex] = 0xff;
                    psTrace("psModules.imcombine", 5, "Excluding stamps %d,%d and %d,%d\n",
                            (int)x->data.F32[iIndex], (int)y->data.F32[iIndex],
                            (int)x->data.F32[jIndex], (int)y->data.F32[jIndex]);
                    // Would 'break' here, but there might be more stamps within the exclusion zone.
                }
            }
        }
        psFree(indexes);
    }

    pmSubtractionStampList *stamps = pmSubtractionStampListAlloc(subMask->numCols, subMask->numRows,
                                                                 region, 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++) {
        if (exclude->data.U8[i]) {
            // Star has been excluded
            continue;
        }
        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
            continue;
        }
        if (!checkStampMask(xPix, yPix, subMask)) {
            // Not a good stamp
            continue;
        }

        for (int j = 0; j < numStamps; 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
                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];
                }

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

                break;
            }
        }
    }
    psFree(exclude);

    // 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 *reference, psImage *input,
                                psImage *weight, int footprint, int kernelSize)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PS_ASSERT_IMAGE_NON_NULL(reference, false);
    PS_ASSERT_IMAGE_TYPE(reference, PS_TYPE_F32, false);
    if (input) {
        PS_ASSERT_IMAGE_NON_NULL(input, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(input, reference, false);
        PS_ASSERT_IMAGE_TYPE(input, PS_TYPE_F32, false);
    }
    if (weight) {
        PS_ASSERT_IMAGE_NON_NULL(weight, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(weight, reference, false);
        PS_ASSERT_IMAGE_TYPE(weight, PS_TYPE_F32, false);
    }
    PS_ASSERT_INT_NONNEGATIVE(footprint, false);
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);

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

    if (!weight) {
        // Use the input (or if I must, the reference) as a rough approximation to the variance map, and HOPE
        // that it's positive.
        weight = input ? input : reference;
    }

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

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

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

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

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

    return true;
}


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

    int size = kernelSize + 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);
            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->input);
        stamp->input = psKernelAlloc(-size, size, -size, size);
        psKernel *input = stamp->input; // Target stamp

        // Put in a Gaussian, just for fun!
        for (int v = -size; v <= size; v++) {
            for (int u = -size; u <= size; u++) {
                input->kernel[v][u] = flux / sigma * 0.5 * M_2_SQRTPI * M_SQRT1_2 *
                    expf(-0.5 * (PS_SQR(u + xStamp) + PS_SQR(v + yStamp)) / PS_SQR(sigma));
            }
        }

    }

    return true;
}


pmSubtractionStampList *pmSubtractionStampsSetFromSources(const psArray *sources, const psImage *subMask,
                                                          const psRegion *region, float spacing,
                                                          int exclusionZone)
{
    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,
                                                            spacing, exclusionZone); // Stamps to return
    psFree(x);
    psFree(y);
    psFree(flux);

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

    return stamps;
}
