#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 "pmSpan.h"
#include "pmFootprint.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 "pmErrorCodes.h"

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

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

#define SOURCE_MASK (PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_SATSTAR | PM_SOURCE_MODE_BLEND | \
                     PM_SOURCE_MODE_BADPSF | PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | \
                     PM_SOURCE_MODE_CR_LIMIT) // Mask for bad sources
#define SOURCE_FAINTEST 50.0            // Faintest magnitude to consider


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// 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);
    psFree(list->window);
}

// 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->matrix);
    psFree(stamp->vector);
}

// 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;
}


// Search a region for a suitable stamp
bool stampSearch(float *xStamp, float *yStamp, // Coordinates of stamp, to return
                 float *fluxStamp, // Flux of stamp, to return
                 const psImage *image1, const psImage *image2, // Images to search
                 float thresh1, float thresh2, // Thresholds for images
                 const psImage *subMask, // Subtraction mask
                 int xMin, int xMax, int yMin, int yMax, // Bounds of search
                 int numCols, int numRows, // Size of images
                 int border             // Border around image
    )
{
    bool found = false;                 // Found a suitable stamp?
    *fluxStamp = -INFINITY;             // Flux of best stamp

    // fprintf (stderr, "xMin, xMax: %d, %d -> ", xMin, xMax);

    // Ensure we're not going to go outside the bounds of the image
    xMin = PS_MAX(border, xMin);
    xMax = PS_MIN(numCols - border - 1, xMax);
    yMin = PS_MAX(border, yMin);
    yMax = PS_MIN(numRows - border - 1, yMax);

    if (xMax < xMin) return false;
    if (yMax < yMin) return false;

    psAssert (xMin <= xMax, "x mismatch?");
    psAssert (yMin <= yMax, "y mismatch?");

    for (int y = yMin; y <= yMax; y++) {
        for (int x = xMin; x <= xMax; x++) {
            if ((image1 && image1->data.F32[y][x] < thresh1) ||
                (image2 && image2->data.F32[y][x] < thresh2)) {
                continue;
            }

            if (subMask && subMask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] &
                (PM_SUBTRACTION_MASK_BORDER | PM_SUBTRACTION_MASK_REJ)) {
                return false;
            }

            // We take the MIN to attempt to avoid transients in both images
            float flux = (image1 && image2) ? PS_MIN(image1->data.F32[y][x], image2->data.F32[y][x]) :
                ((image1) ? image1->data.F32[y][x] : image2->data.F32[y][x]); // Flux at pixel
            if (flux > *fluxStamp) {
                *xStamp = x + 0.5;
                *yStamp = y + 0.5;
                *fluxStamp = flux;
                found = true;
            }
        }
    }

    return found;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Functions for generating ds9 region files
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

static bool ds9regions = false;         // Save ds9 region files?

void pmSubtractionRegions(bool state)
{
    ds9regions = state;
}

FILE *pmSubtractionStampsFile(const pmSubtractionStampList *stamps, const char *filename,
                              const char *description)
{
    if (!ds9regions || !stamps || !filename) {
        return NULL;
    }

    psLogMsg("psModules.imcombine", PS_LOG_INFO, "Writing %s to ds9 region file: %s",
             description, filename);

    FILE *file = fopen(filename, "w");

    // Outline the stamps
    for (int i = 0; i < stamps->num; i++) {
        psRegion *region = stamps->regions->data[i]; // Region of interest
        float xCentre = 0.5 * (region->x0 + region->x1), yCentre = 0.5 * (region->y0 + region->y1);
        fprintf(file, "image;box(%f,%f,%f,%f,0) # color=blue\nimage;text(%f,%f,{%d}) # color=blue\n",
                xCentre, yCentre, region->x1 - region->x0, region->y1 - region->y0,
                xCentre, yCentre, i);
    }

    return file;
}

void pmSubtractionStampPrint(FILE *ds9, float x, float y, float size, const char *color)
{
    if (!ds9regions || !ds9) {
        return;
    }
    fprintf(ds9, "image;circle(%f,%f,%f)", x, y, size);
    if (color && strlen(color) > 0) {
        fprintf(ds9, " # color=%s", color);
    }
    fprintf(ds9, "\n");
    return;
}


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

pmSubtractionStampList *pmSubtractionStampListAlloc(int numCols, int numRows, const psRegion *region,
                                                    int footprint, float spacing, float normFrac,
                                                    float sysErr, float skyErr)
{
    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->window = NULL;
    list->normFrac = normFrac;
    list->normWindow = 0;
    list->footprint = footprint;
    list->sysErr = sysErr;
    list->skyErr = skyErr;

    return list;
}

pmSubtractionStampList *pmSubtractionStampListCopy(const pmSubtractionStampList *in)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(in, NULL);

    pmSubtractionStampList *out = psAlloc(sizeof(pmSubtractionStampList)); // Copied stamp list to return
    psMemSetDeallocator(out, (psFreeFunc)subtractionStampListFree);

    int num = out->num = in->num;       // Number of stamps
    out->stamps = psArrayAlloc(num);
    out->regions = psArrayAlloc(num);
    out->x = NULL;
    out->y = NULL;
    out->flux = NULL;
    out->window = psMemIncrRefCounter(in->window);
    out->footprint = in->footprint;
    out->normWindow = in->normWindow;

    for (int i = 0; i < num; i++) {
        psRegion *inRegion = in->regions->data[i]; // Input region
        out->regions->data[i] = psRegionAlloc(inRegion->x0, inRegion->x1, inRegion->y0, inRegion->y1);

        pmSubtractionStamp *inStamp = in->stamps->data[i]; // Input stamp
        pmSubtractionStamp *outStamp = psAlloc(sizeof(pmSubtractionStamp));
        psMemSetDeallocator(outStamp, (psFreeFunc)subtractionStampFree);
        outStamp->x = inStamp->x;
        outStamp->y = inStamp->y;
        outStamp->flux = inStamp->flux;
        outStamp->xNorm = inStamp->xNorm;
        outStamp->yNorm = inStamp->yNorm;
        outStamp->status = inStamp->status;

        outStamp->image1 = inStamp->image1 ? psKernelCopy(inStamp->image1) : NULL;
        outStamp->image2 = inStamp->image2 ? psKernelCopy(inStamp->image2) : NULL;
        outStamp->weight = inStamp->weight ? psKernelCopy(inStamp->weight) : NULL;

        if (inStamp->convolutions1) {
            int size = inStamp->convolutions1->n; // Size of array
            outStamp->convolutions1 = psArrayAlloc(size);
            for (int j = 0; j < size; j++) {
                psKernel *conv = inStamp->convolutions1->data[j]; // Convolution
                outStamp->convolutions1->data[j] = conv ? psKernelCopy(conv) : NULL;
            }
        } else {
            outStamp->convolutions1 = NULL;
        }
        if (inStamp->convolutions2) {
            int size = inStamp->convolutions2->n; // Size of array
            outStamp->convolutions2 = psArrayAlloc(size);
            for (int j = 0; j < size; j++) {
                psKernel *conv = inStamp->convolutions2->data[j]; // Convolution
                outStamp->convolutions2->data[j] = conv ? psKernelCopy(conv) : NULL;
            }
        } else {
            outStamp->convolutions2 = NULL;
        }

        outStamp->matrix = inStamp->matrix ? psImageCopy(NULL, inStamp->matrix, PS_TYPE_F64) : NULL;
        outStamp->vector = inStamp->vector ? psVectorCopy(NULL, inStamp->vector, PS_TYPE_F64) : NULL;

        out->stamps->data[i] = outStamp;
    }

    return out;
}

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->matrix = NULL;
    stamp->vector = NULL;
    stamp->norm = NAN;

    return stamp;
}


pmSubtractionStampList *pmSubtractionStampsFind(pmSubtractionStampList *stamps, const psImage *image1,
                                                const psImage *image2, const psImage *subMask,
                                                const psRegion *region, float thresh1, float thresh2,
                                                int size, int footprint, float spacing, float normFrac,
                                                float sysErr, float skyErr, pmSubtractionMode mode)
{
    if (!image1 && !image2) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Must specify an image");
        return NULL;
    }
    int numCols = 0, numRows = 0;       // Size of images
    if (image1) {
        PS_ASSERT_IMAGE_NON_NULL(image1, NULL);
        PS_ASSERT_IMAGE_TYPE(image1, PS_TYPE_F32, NULL);
        if (subMask) {
            PS_ASSERT_IMAGES_SIZE_EQUAL(image1, subMask, NULL);
        }
        numCols = image1->numCols;
        numRows = image1->numRows;
    }
    if (image2) {
        PS_ASSERT_IMAGE_NON_NULL(image2, NULL);
        PS_ASSERT_IMAGE_TYPE(image2, PS_TYPE_F32, NULL);
        if (subMask) {
            PS_ASSERT_IMAGES_SIZE_EQUAL(image2, subMask, NULL);
        }
        numCols = image2->numCols;
        numRows = image2->numRows;
    }
    if (image1 && image2) {
        PS_ASSERT_IMAGES_SIZE_EQUAL(image1, image2, NULL);
    }
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, NULL);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_IMAGE_MASK, NULL);
        PS_ASSERT_IMAGE_SIZE(subMask, numCols, numRows, NULL);
    }
    PS_ASSERT_INT_POSITIVE(size, 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(PM_ERR_PROG, true, "Input region (%s) contains NAN values", string);
            psFree(string);
            return false;
        }
        if (region->x0 < 0 || region->x1 > numCols ||
            region->y0 < 0 || region->y1 > numRows) {
            psString string = psRegionToString(*region);
            psError(PM_ERR_PROG, true, "Input region (%s) does not fit in image (%dx%d)",
                    string, numCols, numRows);
            psFree(string);
            return false;
        }
    }

    int border = size + footprint;      // Border size

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

    int numStamps = stamps->num;        // Number of stamp regions
    int numFound = 0;                   // Number of stamps found
    int numSearch = 0;                  // Number of regions searched for new stamp
    int numGood = 0;                    // Number of good stamps in total
    for (int i = 0; i < numStamps; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest

        switch (stamp->status) {
          case PM_SUBTRACTION_STAMP_NONE:
            continue;
          case PM_SUBTRACTION_STAMP_FOUND:
          case PM_SUBTRACTION_STAMP_CALCULATE:
          case PM_SUBTRACTION_STAMP_USED:
            numGood++;
            continue;
          case PM_SUBTRACTION_STAMP_INIT:
          case PM_SUBTRACTION_STAMP_REJECTED:
            numSearch++;

            float xStamp = 0.0, yStamp = 0.0; // Coordinates of stamp
            float fluxStamp = -INFINITY; // 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 stamps off the top of the (sorted) list
                for (int j = xList->n - 1; j >= 0 && !goodStamp; j--) {
                    // Chop off the top of the list
                    xList->n = j;
                    yList->n = j;
                    fluxList->n = j;

#if 0
                    // Fish around a bit to see if we can find a pixel that isn't masked
                    // This is not a good idea if we're using the window feature
                    psTrace("psModules.imcombine", 7, "Searching for stamp %d around %d,%d\n",
                            i, xCentre, yCentre);

                    // Search bounds
                    int xCentre = xList->data.F32[j] - 0.5, yCentre = yList->data.F32[j] - 0.5;// Stamp centre
                    int search = footprint - size; // Search radius
                    int xMin = PS_MAX(border, xCentre - search);
                    int xMax = PS_MIN(numCols - border -1, xCentre + search);
                    int yMin = PS_MAX(border, yCentre - search);
                    int yMax = PS_MIN(numRows - border - 1, yCentre + search);

                    goodStamp = stampSearch(&xStamp, &yStamp, &fluxStamp, image1, image2, thresh1, thresh2,
                                            subMask, xMin, xMax, yMin, yMax, numCols, numRows, border);
                    // fprintf (stderr, "find: %d %d ==> %5.1f %5.1f (\n", xCentre, yCentre, xStamp, yStamp);
#else
                    // Only search the exact centre pixel
                    goodStamp = stampSearch(&xStamp, &yStamp, &fluxStamp, image1, image2, thresh1, thresh2,
                                            subMask, xList->data.F32[j], xList->data.F32[j],
                                            yList->data.F32[j], yList->data.F32[j], numCols, numRows, border);
#endif
                    if (0 && goodStamp) {
                        fprintf (stderr, "find: %6.1f < %6.1f < %6.1f, %6.1f < %6.1f %6.1f\n",
                                 region->x0 + size, xStamp, region->x1 - size,
                                 region->y0 + size, yStamp, region->y1 - size);
                    }
                }
            } else {
                // Use a simple method of automatically finding stamps --- take the highest pixel in the
                // subregion
                psRegion *subRegion = stamps->regions->data[i]; // Sub-region of interest

                goodStamp = stampSearch(&xStamp, &yStamp, &fluxStamp, image1, image2, thresh1, thresh2,
                                        subMask, subRegion->x0, subRegion->x1, subRegion->y0, subRegion->y1,
                                        numCols, numRows, border);
            }

            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 (numSearch > 0) {
        psLogMsg("psModules.imcombine", PS_LOG_INFO, "Found %d stamps", numFound);
    }

    if (numGood == 0 && numFound == 0) {
        // No good stamps
        psError(PM_ERR_STAMPS, true, "Unable to find suitable stamps");
        psFree(stamps);
        return NULL;
    }

    return stamps;
}



pmSubtractionStampList *pmSubtractionStampsSet(const psVector *x, const psVector *y,
                                               const psImage *image, const psImage *subMask,
                                               const psRegion *region, int size, int footprint,
                                               float spacing, float normFrac, float sysErr, float skyErr,
                                               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);
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, NULL);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_IMAGE_MASK, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(image, subMask, NULL);
    }
    PS_ASSERT_INT_POSITIVE(size, NULL);
    PS_ASSERT_INT_POSITIVE(footprint, 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,
                                                                 normFrac, sysErr, skyErr); // Stamp list
    int numStamps = stamps->num;        // Number of stamps

    psString ds9name = NULL;            // Filename for ds9 region file
    static int ds9num = 0;              // File number for ds9 region file
    psStringAppend(&ds9name, "stamps_all_%d.ds9", ds9num);
    FILE *ds9 = pmSubtractionStampsFile(stamps, ds9name, "all stamps"); // ds9 region file
    psFree(ds9name);
    ds9num++;

    // 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);
            pmSubtractionStampPrint(ds9, xPix, yPix, footprint, "red");
            continue;
        }

        // fprintf (stderr, "stamp: %5.1f %5.1f == %d %d\n", xStamp, yStamp, xPix, yPix);

        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;
                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);
                pmSubtractionStampPrint(ds9, xPix, yPix, footprint, "green");
            }
        }

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

    if (ds9) {
        fclose(ds9);
    }

    // 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 pmSubtractionStampsGetWindow(pmSubtractionStampList *stamps, int kernelSize)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);

    int size = stamps->footprint; // Size of postage stamps

    psFree (stamps->window);
    stamps->window = psKernelAlloc(-size, size, -size, size);
    psImageInit(stamps->window->image, 0.0);

    // generate normalizations for each stamp
    psVector *norm1 = psVectorAlloc(stamps->num, PS_TYPE_F32);
    psVector *norm2 = psVectorAlloc(stamps->num, PS_TYPE_F32);
    for (int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
        if (!stamp) continue;
        if (!stamp->image1) continue;
        if (!stamp->image2) continue;

        float sum1 = 0.0;
        float sum2 = 0.0;
        for (int y = -size; y <= size; y++) {
            for (int x = -size; x <= size; x++) {
                sum1 += stamp->image1->kernel[y][x];
                sum2 += stamp->image2->kernel[y][x];
            }
        }
        norm1->data.F32[i] = sum1;
        norm2->data.F32[i] = sum2;
    }

    // storage vector for flux data
    psStats *stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    psVector *flux1 = psVectorAllocEmpty(2*stamps->num, PS_TYPE_F32);
    psVector *flux2 = psVectorAllocEmpty(2*stamps->num, PS_TYPE_F32);

    // generate the window pixels
    double sum = 0.0;                   // Sum inside the window
    float maxValue = 0.0;               // Maximum value, for normalisation
    for (int y = -size; y <= size; y++) {
        for (int x = -size; x <= size; x++) {

            flux1->n = 0;
            flux2->n = 0;
            for (int i = 0; i < stamps->num; i++) {
                pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
                if (!stamp) continue;
                if (!stamp->image1) continue;
                if (!stamp->image2) continue;

                psVectorAppend(flux1, stamp->image1->kernel[y][x] / norm1->data.F32[i]);
                psVectorAppend(flux2, stamp->image2->kernel[y][x] / norm2->data.F32[i]);
            }

            psStatsInit (stats);
            if (!psVectorStats (stats, flux1, NULL, NULL, 0)) {
                psAbort ("failed to generate stats");
            }
            float f1 = stats->robustMedian;
            psStatsInit (stats);
            if (!psVectorStats (stats, flux2, NULL, NULL, 0)) {
                psAbort ("failed to generate stats");
            }
            float f2 = stats->robustMedian;

            stamps->window->kernel[y][x] = f1 + f2;
            if (PS_SQR(x) + PS_SQR(y) <= PS_SQR(size)) {
                sum += stamps->window->kernel[y][x];
            }
            maxValue = PS_MAX(maxValue, stamps->window->kernel[y][x]);
        }
    }

    psTrace("psModules.imcombine", 3, "Window total: %f, threshold: %f\n",
            sum, (1.0 - stamps->normFrac) * sum);
    bool done = false;
    for (int radius = 1; radius <= size && !done; radius++) {
        double within = 0.0;
        for (int y = -radius; y <= radius; y++) {
            for (int x = -radius; x <= radius; x++) {
                if (PS_SQR(x) + PS_SQR(y) <= PS_SQR(radius)) {
                    within += stamps->window->kernel[y][x];
                }
            }
        }
        psTrace("psModules.imcombine", 5, "Radius %d: %f\n", radius, within);
        if (within > (1.0 - stamps->normFrac) * sum) {
            stamps->normWindow = radius;
            done = true;
        }
    }

    psTrace("psModules.imcombine", 3, "Normalisation window radius set to %d\n", stamps->normWindow);
    if (stamps->normWindow == 0 || stamps->normWindow >= size) {
        psError(PM_ERR_STAMPS, true, "Unable to determine normalisation window size.");
        return false;
    }

    // re-normalize so chisquare values are sensible
    for (int y = -size; y <= size; y++) {
        for (int x = -size; x <= size; x++) {
            stamps->window->kernel[y][x] /= maxValue;
        }
    }

#if 0
    {
        psFits *fits = psFitsOpen ("window.fits", "w");
        psFitsWriteImage (fits, NULL, stamps->window->image, 0, NULL);
        psFitsClose (fits);
    }
#endif

    psFree (stats);
    psFree (flux1);
    psFree(flux2);
    psFree (norm1);
    psFree (norm2);
    return true;
}

bool pmSubtractionStampsExtract(pmSubtractionStampList *stamps, psImage *image1, psImage *image2,
                                psImage *variance, int kernelSize, psRegion bounds)
{
    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);
    }
    if (variance) {
        PS_ASSERT_IMAGE_NON_NULL(variance, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(variance, image1, false);
        PS_ASSERT_IMAGE_TYPE(variance, PS_TYPE_F32, false);
        PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);
    }

    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;
        }

        p_pmSubtractionPolynomialNormCoords(&stamp->xNorm, &stamp->yNorm, stamp->x, stamp->y,
                                            bounds.x0, bounds.x1, bounds.y0, bounds.y1);

        int x = stamp->x - 0.5, y = stamp->y - 0.5; // Stamp coordinates
        // fprintf (stderr, "stamp: %5.1f %5.1f == %d %d (size: %d)\n", stamp->x, stamp->y, x, y, size);

        if (x < bounds.x0 + size || x > bounds.x1 - size || y < bounds.y0 + size || y > bounds.y1 - size) {
            psError(PM_ERR_PROG, false, "Stamp %d (%d,%d) is within the region 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
        }

        psKernel *weight = stamp->weight = psKernelAlloc(-size, size, -size, size); // Weight = 1/variance
        if (variance) {
            psImage *varSub = psImageSubset(variance, region); // Subimage with stamp
            psKernel *var = psKernelAllocFromImage(varSub, size, size); // Variance postage stamp

            if ((isfinite(stamps->skyErr) && (stamps->skyErr > 0)) ||
                (isfinite(stamps->sysErr) && (stamps->sysErr > 0))) {
                float sysErr = 0.25 * PS_SQR(stamps->sysErr); // Systematic error
                float skyErr = stamps->skyErr;
                psKernel *image1 = stamp->image1, *image2 = stamp->image2; // Input images
                for (int y = -size; y <= size; y++) {
                    for (int x = -size; x <= size; x++) {
                        float additional = image1->kernel[y][x] + image2->kernel[y][x];
                        weight->kernel[y][x] = 1.0 / (skyErr + var->kernel[y][x] + sysErr * PS_SQR(additional));
                    }
                }
            } else {
                for (int y = -size; y <= size; y++) {
                    for (int x = -size; x <= size; x++) {
                        weight->kernel[y][x] = 1.0 / var->kernel[y][x];
                    }
                }
            }
            psFree(var);
            psFree(varSub);
        } else {
            psImageInit(weight->image, 1.0);
        }

        stamp->status = PM_SUBTRACTION_STAMP_CALCULATE;
    }

    return true;
}

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

    int numIn = sources->n;             // Number of sources in input list

    psVector *x = psVectorAllocEmpty(numIn, PS_TYPE_F32); // x coordinates
    psVector *y = psVectorAllocEmpty(numIn, PS_TYPE_F32); // y coordinates

    int numOut = 0;                     // Number of sources in output list
    for (int i = 0; i < numIn; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (!source || (source->mode & SOURCE_MASK) || source->psfMag > SOURCE_FAINTEST) {
            continue;
        }
        if (source->modelPSF) {
            x->data.F32[numOut] = source->modelPSF->params->data.F32[PM_PAR_XPOS];
            y->data.F32[numOut] = source->modelPSF->params->data.F32[PM_PAR_YPOS];
        } else {
            x->data.F32[numOut] = source->peak->xf;
            y->data.F32[numOut] = source->peak->yf;
        }
        // fprintf (stderr, "stamp: %5.1f %5.1f\n", x->data.F32[numOut], y->data.F32[numOut]);
        numOut++;
    }
    x->n = numOut;
    y->n = numOut;

    pmSubtractionStampList *stamps = pmSubtractionStampsSet(x, y, image, subMask, region, size,
                                                            footprint, spacing, normFrac,
                                                            sysErr, skyErr, mode); // Stamps
    psFree(x);
    psFree(y);

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

    return stamps;
}


pmSubtractionStampList *pmSubtractionStampsSetFromFile(const char *filename, const psImage *image,
                                                       const psImage *subMask, const psRegion *region,
                                                       int size, int footprint, float spacing, float normFrac,
                                                       float sysErr, float skyErr, 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(psErrorCodeLast(), 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, image, subMask, region, size, footprint,
                                                            spacing, normFrac, sysErr, skyErr, mode);
    psFree(data);

    return stamps;

}


bool pmSubtractionStampsResetStatus (pmSubtractionStampList *stamps) {

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

