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

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

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmSubtractionParams.h"
#include "pmSubtractionKernels.h"
#include "pmSubtractionStamps.h"
#include "pmSubtractionEquation.h"
#include "pmSubtraction.h"
#include "pmSubtractionAnalysis.h"
#include "pmSubtractionMask.h"
#include "pmSubtractionThreads.h"
#include "pmSubtractionMatch.h"


#define BG_STAT PS_STAT_ROBUST_MEDIAN   // Statistic to use for background

static bool useFFT = true;              // Do convolutions using FFT


//#define TESTING
//#define TESTING_MEMORY

// Output memory usage information
static void memCheck(const char *where)
{
#ifdef TESTING_MEMORY
    psMemBlock **leaks = NULL;
    int numLeaks = psMemCheckLeaks(0, &leaks, NULL, true);
    size_t largestSize = 0;
    psMemId largest = 0;
    size_t totalSize = 0;
    for (int i = 0; i < numLeaks; i++) {
        psMemBlock *mb = leaks[i];
        totalSize += mb->userMemorySize;
        if (mb->userMemorySize > largestSize) {
            largestSize = mb->userMemorySize;
            largest = mb->id;
        }
    }
    psFree(leaks);
    fprintf(stderr, "%s:\n", where);
    fprintf(stderr, "    Memory in use: %zd\n", totalSize);
    fprintf(stderr, "    Largest block: %ld\n", largest);
    fprintf(stderr, "    sbrk(): %zd\n", (size_t)sbrk(0));
#endif
    return;
}


static bool getStamps(pmSubtractionStampList **stamps, // Stamps to read
                      const pmReadout *ro1, // Readout 1
                      const pmReadout *ro2, // Readout 2
                      const psImage *subMask, // Mask for subtraction, or NULL
                      psImage *weight,  // Weight map
                      const psRegion *region, // Region of interest, or NULL
                      float threshold,  // Threshold for stamp finding
                      float stampSpacing, // Spacing between stamps
                      int size,         // Kernel half-size
                      int footprint,     // Convolution footprint for stamps
                      pmSubtractionMode mode // Mode for subtraction
    )
{
    psTrace("psModules.imcombine", 3, "Finding stamps...\n");
    *stamps = pmSubtractionStampsFind(*stamps, ro1->image, subMask, region, threshold, footprint,
                                      stampSpacing, mode);
    if (!*stamps) {
        psError(PS_ERR_UNKNOWN, false, "Unable to find stamps.");
        return false;
    }

    memCheck("  find stamps");

    psTrace("psModules.imcombine", 3, "Extracting stamps...\n");
    if (!pmSubtractionStampsExtract(*stamps, ro1->image, ro2 ? ro2->image : NULL, weight, size)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to extract stamps.");
        return false;
    }

    memCheck("   extract stamps");

    return true;
}


bool pmSubtractionMatch(pmReadout *conv1, pmReadout *conv2, const pmReadout *ro1, const pmReadout *ro2,
                        int footprint, int stride, float regionSize, float stampSpacing, float threshold,
                        const psArray *sources, const char *stampsName,
                        pmSubtractionKernelsType type, int size, int spatialOrder,
                        const psVector *isisWidths, const psVector *isisOrders,
                        int inner, int ringsOrder, int binning, float penalty,
                        bool optimum, const psVector *optFWHMs, int optOrder, float optThreshold,
                        int iter, float rej, float sysError, psImageMaskType maskVal, psImageMaskType maskBad,
                        psImageMaskType maskPoor, float poorFrac, float badFrac, pmSubtractionMode subMode)
{
    if (subMode != PM_SUBTRACTION_MODE_2) {
        PM_ASSERT_READOUT_NON_NULL(conv1, false);
        if (conv1->image) {
            psFree(conv1->image);
            conv1->image = NULL;
        }
        if (conv1->mask) {
            psFree(conv1->mask);
            conv1->mask = NULL;
        }
        if (conv1->weight) {
            psFree(conv1->weight);
            conv1->weight = NULL;
        }
    }
    if (subMode != PM_SUBTRACTION_MODE_1) {
        PM_ASSERT_READOUT_NON_NULL(conv2, false);
        if (conv2->image) {
            psFree(conv2->image);
            conv2->image = NULL;
        }
        if (conv2->mask) {
            psFree(conv2->mask);
            conv2->mask = NULL;
        }
        if (conv2->weight) {
            psFree(conv2->weight);
            conv2->weight = NULL;
        }
    }

    PM_ASSERT_READOUT_NON_NULL(ro1, false);
    PM_ASSERT_READOUT_NON_NULL(ro2, false);
    PM_ASSERT_READOUT_IMAGE(ro1, false);
    PM_ASSERT_READOUT_IMAGE(ro2, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(ro1->image, ro2->image, false);

    PS_ASSERT_INT_NONNEGATIVE(footprint, false);
    PS_ASSERT_INT_NONNEGATIVE(stride, false);
    // regionSize can be just about anything (except maybe negative, but it can be NAN)
    PS_ASSERT_FLOAT_LARGER_THAN(stampSpacing, 0.0, false);
    // Don't care what threshold is
    if (sources) {
        PS_ASSERT_ARRAY_NON_NULL(sources, false);
    }
    // stampsName may be anything
    // We'll check kernel type when we allocate the kernels
    PS_ASSERT_INT_POSITIVE(size, false);
    PS_ASSERT_INT_NONNEGATIVE(spatialOrder, false);
    if (isisWidths || isisOrders) {
        PS_ASSERT_VECTOR_NON_NULL(isisWidths, false);
        PS_ASSERT_VECTOR_TYPE(isisWidths, PS_TYPE_F32, false);
        PS_ASSERT_VECTOR_NON_NULL(isisOrders, false);
        PS_ASSERT_VECTOR_TYPE(isisOrders, PS_TYPE_S32, false);
        PS_ASSERT_VECTORS_SIZE_EQUAL(isisWidths, isisOrders, false);
    }
    PS_ASSERT_INT_NONNEGATIVE(inner, false);
    PS_ASSERT_INT_NONNEGATIVE(ringsOrder, false);
    PS_ASSERT_INT_POSITIVE(binning, false);
    if (optimum) {
        PS_ASSERT_VECTOR_NON_NULL(optFWHMs, false);
        PS_ASSERT_INT_NONNEGATIVE(optOrder, false);
        PS_ASSERT_FLOAT_LARGER_THAN(optThreshold, 0.0, false);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(optThreshold, 1.0, false);
    }
    PS_ASSERT_INT_POSITIVE(iter, false);
    PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    if (isfinite(sysError)) {
        PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(sysError, 0.0, NULL);
        PS_ASSERT_FLOAT_LESS_THAN(sysError, 1.0, NULL);
    }
    // Don't care about maskVal
    // Don't care about maskBad
    // Don't care about maskPoor
    PS_ASSERT_FLOAT_LARGER_THAN(poorFrac, 0.0, NULL);
    PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(poorFrac, 1.0, NULL);
    if (isfinite(badFrac)) {
        PS_ASSERT_FLOAT_LARGER_THAN(badFrac, 0.0, NULL);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(badFrac, 1.0, NULL);
    }

    // If the stamp footprint is smaller than the kernel size, then we won't get much signal in the outer
    // parts of the kernel, which can result in bad matching artifacts.
    if (footprint < size) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "Stamp footprint (%d) should be larger than or equal to the kernel size (%d)",
                footprint, size);
        return false;
    }

    // Where does our weight map come from?
    // Getting the weight exactly right is not necessary --- it's just used for weighting.
    psImage *weight = NULL;             // Weight image to use
    if (ro1->weight && ro2->weight) {
        weight = (psImage*)psBinaryOp(NULL, ro1->weight, "+", ro2->weight);
    } else if (ro1->weight) {
        weight = psMemIncrRefCounter(ro1->weight);
    } else if (ro2->weight) {
        weight = psMemIncrRefCounter(ro2->weight);
    } else {
        weight = (psImage*)psBinaryOp(NULL, ro1->image, "+", ro2->image);
    }

    // Putting important variable declarations here, since they are freed after a "goto" if there is an error.
    psImage *subMask = NULL;            // Mask for subtraction
    psRegion *region = NULL;            // Iso-kernel region
    psString regionString = NULL;       // String for region
    pmSubtractionStampList *stamps = NULL; // Stamps for matching PSF
    pmSubtractionKernels *kernels = NULL; // Kernel basis functions
    psMetadata *analysis = psMetadataAlloc(); // QA data

    int numCols = ro1->image->numCols, numRows = ro1->image->numRows; // Image dimensions

    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0); // Random number generator

    memCheck("start");

    subMask = pmSubtractionMask(ro1->mask, ro2 ? ro2->mask : NULL, maskVal, size, footprint,
                                badFrac, useFFT);
    if (!subMask) {
        psError(PS_ERR_UNKNOWN, false, "Unable to generate subtraction mask.");
        goto MATCH_ERROR;
    }

    memCheck("mask");

    // Get region of interest
    int xRegions = 1, yRegions = 1;     // Number of iso-kernel regions
    float xRegionSize = 0, yRegionSize = 0; // Size of iso-kernel regions
    if (isfinite(regionSize) && regionSize != 0.0) {
        xRegions = numCols / regionSize + 1;
        yRegions = numRows / regionSize + 1;
        xRegionSize = (float)numCols / (float)xRegions;
        yRegionSize = (float)numRows / (float)yRegions;
        region = psRegionAlloc(NAN, NAN, NAN, NAN);
    }

    {
        psStats *bg = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV); // Statistics for backgroun
        if (!psImageBackground(bg, NULL, ro1->image, ro1->mask, maskVal, rng)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to measure background statistics.");
            psFree(bg);
            psFree(rng);
            goto MATCH_ERROR;
        }
        threshold = bg->robustMedian + threshold * bg->robustStdev;
        psFree(bg);
    }

    // Iterate over iso-kernel regions
    for (int j = 0; j < yRegions; j++) {
        for (int i = 0; i < xRegions; i++) {
            psTrace("psModules.imcombine", 1, "Subtracting region %d of %d...\n",
                    j * xRegions + i + 1, xRegions * yRegions);
            if (region) {
                *region = psRegionSet((int)(i * xRegionSize), (int)((i + 1) * xRegionSize),
                                      (int)(j * yRegionSize), (int)((j + 1) * yRegionSize));
                psFree(regionString);
                regionString = psRegionToString(*region);
                psTrace("psModules.imcombine", 3, "Iso-kernel region: %s out of %d,%d\n",
                        regionString, numCols, numRows);
            }

            if (sources) {
                stamps = pmSubtractionStampsSetFromSources(sources, ro1->image, subMask, region, footprint,
                                                           stampSpacing, subMode);
            } else if (stampsName && strlen(stampsName) > 0) {
                stamps = pmSubtractionStampsSetFromFile(stampsName, ro1->image, subMask, region, footprint,
                                                        stampSpacing, subMode);
            }

            // We get the stamps here; we will also attempt to get stamps at the first iteration, but it
            // doesn't matter.
            if (!getStamps(&stamps, ro1, ro2, subMask, weight, NULL, threshold, stampSpacing,
                           size, footprint, subMode)) {
                goto MATCH_ERROR;
            }

            if (subMode == PM_SUBTRACTION_MODE_UNSURE) {
                // Get backgrounds
                psStats *bgStats = psStatsAlloc(BG_STAT); // Statistics for background
                psVector *buffer = NULL;// Buffer for stats
                if (!psImageBackground(bgStats, &buffer, ro1->image, ro1->mask, maskVal, rng)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to measure background of image 1.");
                    psFree(bgStats);
                    psFree(buffer);
                    goto MATCH_ERROR;
                }
                float bg1 = psStatsGetValue(bgStats, BG_STAT); // Background for image 1
                if (!psImageBackground(bgStats, &buffer, ro2->image, ro2->mask, maskVal, rng)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to measure background of image 2.");
                    psFree(bgStats);
                    psFree(buffer);
                    goto MATCH_ERROR;
                }
                float bg2 = psStatsGetValue(bgStats, BG_STAT); // Background for image 2
                psFree(bgStats);
                psFree(buffer);

                pmSubtractionMode newMode = pmSubtractionOrder(stamps, bg1, bg2); // Subtraction mode to use
                switch (newMode) {
                  case PM_SUBTRACTION_MODE_1:
                    psLogMsg("psModules.imcombine", PS_LOG_INFO, "Convolving image 1 to match image 2.");
                    break;
                  case PM_SUBTRACTION_MODE_2:
                    psLogMsg("psModules.imcombine", PS_LOG_INFO, "Convolving image 2 to match image 1.");
                    break;
                  default:
                    psError(PS_ERR_UNKNOWN, false, "Unable to determine subtraction order.");
                    goto MATCH_ERROR;
                }
                subMode = newMode;
            }

            // Define kernel basis functions
            if (optimum && (type == PM_SUBTRACTION_KERNEL_ISIS || type == PM_SUBTRACTION_KERNEL_GUNK)) {
                kernels = pmSubtractionKernelsOptimumISIS(type, size, inner, spatialOrder, optFWHMs, optOrder,
                                                          stamps, footprint, optThreshold, penalty, subMode);
                if (!kernels) {
                    psErrorClear();
                    psWarning("Unable to derive optimum ISIS kernel --- switching to default.");
                }
            }
            if (kernels == NULL) {
                // Not an ISIS/GUNK kernel, or the optimum kernel search failed
                kernels = pmSubtractionKernelsGenerate(type, size, spatialOrder, isisWidths, isisOrders,
                                                       inner, binning, ringsOrder, penalty, subMode);
            }

            memCheck("kernels");

            int numRejected = -1;       // Number of rejected stamps in each iteration
            for (int k = 0; k < iter && numRejected != 0; k++) {
                psLogMsg("psModules.imcombine", PS_LOG_INFO, "Iteration %d.", k);

                if (!getStamps(&stamps, ro1, ro2, subMask, weight, region, threshold, stampSpacing,
                               size, footprint, subMode)) {
                    goto MATCH_ERROR;
                }

                psTrace("psModules.imcombine", 3, "Calculating equation...\n");
                if (!pmSubtractionCalculateEquation(stamps, kernels)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
                    goto MATCH_ERROR;
                }

                memCheck("  calculate equation");

                psTrace("psModules.imcombine", 3, "Solving equation...\n");

                if (!pmSubtractionSolveEquation(kernels, stamps)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
                    goto MATCH_ERROR;
                }

                memCheck("  solve equation");

                psVector *deviations = pmSubtractionCalculateDeviations(stamps, kernels); // Stamp deviations
                if (!deviations) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate deviations.");
                    goto MATCH_ERROR;
                }

                memCheck("   calculate deviations");

                psTrace("psModules.imcombine", 3, "Rejecting stamps...\n");
                numRejected = pmSubtractionRejectStamps(kernels, stamps, deviations, subMask, rej, footprint);
                if (numRejected < 0) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to reject stamps.");
                    psFree(deviations);
                    goto MATCH_ERROR;
                }
                psFree(deviations);

                memCheck("  reject stamps");
            }

            if (numRejected > 0) {
                psTrace("psModules.imcombine", 3, "Solving equation...\n");
                if (!pmSubtractionSolveEquation(kernels, stamps)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
                    goto MATCH_ERROR;
                }
                psVector *deviations = pmSubtractionCalculateDeviations(stamps, kernels); // Stamp deviations
                if (!deviations) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate deviations.");
                    goto MATCH_ERROR;
                }
                pmSubtractionRejectStamps(kernels, stamps, deviations, subMask, NAN, footprint);
                psFree(deviations);
            }
            psFree(stamps);
            stamps = NULL;

            memCheck("solution");

            if (!pmSubtractionAnalysis(analysis, kernels, region, numCols, numRows)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to generate QA data");
                goto MATCH_ERROR;
            }

            memCheck("diag outputs");

            psTrace("psModules.imcombine", 2, "Convolving...\n");
            if (!pmSubtractionConvolve(conv1, conv2, ro1, ro2, subMask, stride, maskBad, maskPoor, poorFrac,
                                       sysError, region, kernels, true, useFFT)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to convolve image.");
                goto MATCH_ERROR;
            }

            psFree(kernels);
            kernels = NULL;

            // There is data in the readout now
            if (subMode == PM_SUBTRACTION_MODE_1 || subMode == PM_SUBTRACTION_MODE_DUAL) {
                conv1->data_exists = true;
                if (conv1->parent) {
                    conv1->parent->data_exists = true;
                    conv1->parent->parent->data_exists = true;
                }
            }
            if (subMode == PM_SUBTRACTION_MODE_2 || subMode == PM_SUBTRACTION_MODE_DUAL) {
                conv2->data_exists = true;
                if (conv2->parent) {
                    conv2->parent->data_exists = true;
                    conv2->parent->parent->data_exists = true;
                }
            }
        }
    }
    psFree(rng);
    rng = NULL;
    psFree(region);
    region = NULL;
    psFree(regionString);
    regionString = NULL;
    psFree(subMask);
    subMask = NULL;
    psFree(weight);
    weight = NULL;

    if (!pmSubtractionBorder(conv1->image, conv1->weight, conv1->mask, size, maskBad)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to set border of convolved image.");
        goto MATCH_ERROR;
    }

    memCheck("convolution");

    if (conv1) {
        psMetadataCopy(conv1->analysis, analysis);
    }
    if (conv2) {
        psMetadataCopy(conv2->analysis, analysis);
    }
    psFree(analysis);

#ifdef TESTING
    {
        if (subMode == PM_SUBTRACTION_MODE_1 || subMode == PM_SUBTRACTION_MODE_DUAL) {
            psFits *fits = psFitsOpen("convolved1.fits", "w");
            psFitsWriteImage(fits, NULL, conv1->image, 0, NULL);
            psFitsClose(fits);
        }

        if (subMode == PM_SUBTRACTION_MODE_2 || subMode == PM_SUBTRACTION_MODE_DUAL) {
            psFits *fits = psFitsOpen("convolved2.fits", "w");
            psFitsWriteImage(fits, NULL, conv2->image, 0, NULL);
            psFitsClose(fits);
        }
    }
#endif

    return true;

MATCH_ERROR:
    psFree(analysis);
    psFree(region);
    psFree(regionString);
    psFree(subMask);
    psFree(kernels);
    psFree(stamps);
    psFree(weight);
    psFree(rng);
    return false;
}


// Determine a rough width (integer value) of the star in the image
// XXX Could improve this by using a user-provided list of floating-point widths (or an end point and
// increment).
static int subtractionOrderWidth(const psKernel *kernel, // Image
                                 float bg, // Background in image
                                 int size, // Maximum size
                                 const psArray *models, // Buffer of models
                                 const psVector *modelSums // Buffer of model sums
    )
{
    assert(kernel);
    assert(models);
    assert(modelSums);

    int xMin = -size, xMax = size; // Bounds in x
    int yMin = -size, yMax = size; // Bounds in y

    // Fit gaussians of varying widths to the image, record the chi^2
    psVector *chi2 = psVectorAlloc(size, PS_TYPE_F32); // chi^2 as a function of radius
    for (int sigma = 0; sigma < size; sigma++) {
        double sumFG = 0.0; // Sum for calculating the normalisation of the Gaussian
        psKernel *model = models->data[sigma]; // Model of interest
        for (int y = yMin; y <= yMax; y++) {
            for (int x = xMin; x <= xMax; x++) {
                sumFG += model->kernel[y][x] * (kernel->kernel[y][x] - bg);
            }
        }
        float norm = sumFG * modelSums->data.F64[sigma]; // Normalisation for Gaussian
        double sumDev2 = 0.0;           // Sum of square deviations
        for (int y = yMin; y <= yMax; y++) {
            for (int x = xMin; x <= xMax; x++) {
                float dev = kernel->kernel[y][x] - bg - norm * model->kernel[y][x]; // Deviation
                sumDev2 += PS_SQR(dev);
            }
        }
        chi2->data.F32[sigma] = sumDev2;
    }

    // Find the minimum chi^2
    int bestIndex = -1;                 // Index of best chi^2
    float bestChi2 = INFINITY;          // Best chi^2
    for (int i = 0; i < size; i++) {
        if (chi2->data.F32[i] < bestChi2) {
            bestChi2 = chi2->data.F32[i];
            bestIndex = i;
        }
    }
    psFree(chi2);

    return bestIndex + 1;
}


bool pmSubtractionOrderStamp(psVector *ratios, psVector *mask, const pmSubtractionStampList *stamps,
                             const psArray *models, const psVector *modelSums,
                             int index, float bg1, float bg2)
{
    PS_ASSERT_VECTOR_NON_NULL(ratios, false);
    PS_ASSERT_VECTOR_NON_NULL(mask, false);
    PS_ASSERT_VECTORS_SIZE_EQUAL(ratios, mask, false);
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PS_ASSERT_INT_NONNEGATIVE(index, false);
    PS_ASSERT_INT_LESS_THAN(index, stamps->num, false);
    PS_ASSERT_ARRAY_NON_NULL(models, false);
    PS_ASSERT_VECTOR_NON_NULL(modelSums, false);
    PS_ASSERT_VECTOR_SIZE(modelSums, models->n, false);

    pmSubtractionStamp *stamp = stamps->stamps->data[index]; // Stamp of interest
    psAssert(stamp->status == PM_SUBTRACTION_STAMP_CALCULATE || stamp->status == PM_SUBTRACTION_STAMP_USED,
             "We checked this earlier.");

    // Widths of stars
    int width1 = subtractionOrderWidth(stamp->image1, bg1, stamps->footprint, models, modelSums);
    int width2 = subtractionOrderWidth(stamp->image2, bg2, stamps->footprint, models, modelSums);

    if (width1 == 0 || width2 == 0) {
        ratios->data.F32[index] = NAN;
        mask->data.PS_TYPE_VECTOR_MASK_DATA[index] = 0xff;
    } else {
        ratios->data.F32[index] = (float)width1 / (float)width2;
        mask->data.PS_TYPE_VECTOR_MASK_DATA[index] = 0;
    }

    return true;
}

bool pmSubtractionOrderThread(psThreadJob *job)
{
    PS_ASSERT_THREAD_JOB_NON_NULL(job, false);

    psVector *ratios = job->args->data[0]; // Ratios of widths
    psVector *mask = job->args->data[1]; // Mask for ratios
    const pmSubtractionStampList *stamps = job->args->data[2]; // List of stamps
    const psArray *models = job->args->data[3]; // Gaussian models
    const psVector *modelSums = job->args->data[4]; // Gaussian model sums
    int index = PS_SCALAR_VALUE(job->args->data[5], S32); // Stamp index
    float bg1 = PS_SCALAR_VALUE(job->args->data[6], F32); // Background of image 1
    float bg2 = PS_SCALAR_VALUE(job->args->data[7], F32); // Background of image 2

    return pmSubtractionOrderStamp(ratios, mask, stamps, models, modelSums, index, bg1, bg2);
}

pmSubtractionMode pmSubtractionOrder(pmSubtractionStampList *stamps, float bg1, float bg2)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, PM_SUBTRACTION_MODE_ERR);

    psVector *mask = psVectorAlloc(stamps->num, PS_TYPE_VECTOR_MASK); // Mask for stamps
    psVector *ratios = psVectorAlloc(stamps->num, PS_TYPE_F32); // Ratios of widths

    // Generate models
    int size = stamps->footprint;       // Maximum size
    psArray *models = psArrayAlloc(size); // Gaussian models
    psVector *modelSums = psVectorAlloc(size, PS_TYPE_F64); // Gaussian model sums
    for (int sigma = 0; sigma < size; sigma++) {
        psKernel *model = psKernelAlloc(-size, size, -size, size); // Gaussian model
        float invSigma2 = 1.0 / (float)PS_SQR(1 + sigma); // Inverse sigma squared
        double sumGG = 0.0;         // Sum of square of Gaussian
        for (int y = -size; y <= size; y++) {
            int y2 = PS_SQR(y);     // y squared
            for (int x = -size; x <= size; x++) {
                float rad2 = PS_SQR(x) + y2; // Radius squared
                float value = expf(-rad2 * invSigma2); // Model value
                model->kernel[y][x] = value;
                sumGG += PS_SQR(value);
            }
        }
        models->data[sigma] = model;
        modelSums->data.F64[sigma] = 1.0 / sumGG;
    }

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

        if (pmSubtractionThreaded()) {
            psThreadJob *job = psThreadJobAlloc("PSMODULES_SUBTRACTION_ORDER");
            psArrayAdd(job->args, 1, ratios);
            psArrayAdd(job->args, 1, mask);
            psArrayAdd(job->args, 1, stamps);
            psArrayAdd(job->args, 1, models);
            psArrayAdd(job->args, 1, modelSums);
            PS_ARRAY_ADD_SCALAR(job->args, i, PS_TYPE_S32);
            PS_ARRAY_ADD_SCALAR(job->args, bg1, PS_TYPE_F32);
            PS_ARRAY_ADD_SCALAR(job->args, bg2, PS_TYPE_F32);
            if (!psThreadJobAddPending(job)) {
                psFree(job);
                return false;
            }
            psFree(job);
        } else {
            if (!pmSubtractionOrderStamp(ratios, mask, stamps, models, modelSums, i, bg1, bg2)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to measure PSF width for stamp %d", i);
                psFree(models);
                psFree(modelSums);
                psFree(ratios);
                psFree(mask);
                return false;
            }
        }
    }

    if (!psThreadPoolWait(true)) {
        psError(PS_ERR_UNKNOWN, false, "Error waiting for threads.");
        psFree(models);
        psFree(modelSums);
        psFree(ratios);
        psFree(mask);
            return false;
    }

    psFree(models);
    psFree(modelSums);

    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN);
    if (!psVectorStats(stats, ratios, NULL, mask, 0xff)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to calculate statistics for moments ratio.");
        psFree(mask);
        psFree(ratios);
        psFree(stats);
        return PM_SUBTRACTION_MODE_ERR;
    }
    psFree(ratios);
    psFree(mask);

    psLogMsg("psModules.imcombine", PS_LOG_INFO, "Median width ratio: %lf", stats->robustMedian);
    pmSubtractionMode mode = (stats->robustMedian <= 1.0 ? PM_SUBTRACTION_MODE_1 : PM_SUBTRACTION_MODE_2);
    psFree(stats);

    return mode;
}



