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

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

#include "ppStack.h"

#define ARRAY_BUFFER 16                 // Number to add to array at a time
#define MAG_IGNORE 50                   // Ignore magnitudes fainter than this --- they're not real!
#define FAKE_SIZE 1                     // Size of fake convolution kernel
#define SOURCE_MASK (PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | \
                     PM_SOURCE_MODE_CR_LIMIT) // Mask to apply to input sources

//#define TESTING

bool ppStackMatch(pmReadout *readout, psArray **regions, psArray **kernels,
                  const psArray *sources, const pmPSF *psf, psRandom *rng, const pmConfig *config)
{
    assert(readout);
    assert(regions && !*regions);
    assert(kernels && !*kernels);
    assert(config);

    pmReadout *output = pmReadoutAlloc(NULL); // Output readout, for holding results temporarily

    // Look up appropriate values from the ppSub recipe
    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, "PPSUB"); // PPSUB recipe
    int size = psMetadataLookupS32(NULL, recipe, "KERNEL.SIZE"); // Kernel half-size

    psMaskType maskIn = pmConfigMaskGet("MASK.VALUE", config); // Bits to mask going in to pmSubtractionMatch
    psMaskType maskOut = pmConfigMaskGet("BLANK", config); // Bits to mask after pmSubtractionMatch

    bool mdok;                          // Status of MD lookup
    bool renorm = psMetadataLookupBool(&mdok, config->arguments, "RENORM"); // Renormalise variances?
    psStatsOptions renormMean = psMetadataLookupS32(&mdok, config->arguments,
                                                    "RENORM.MEAN"); // Statistic for mean
    psStatsOptions renormStdev = psMetadataLookupS32(&mdok, config->arguments,
                                                     "RENORM.STDEV"); // Statistic for stdev
    int renormWidth = psMetadataLookupS32(&mdok, config->arguments,
                                          "RENORM.WIDTH"); // Width for renormalisation box
    float penalty = psMetadataLookupF32(NULL, config->arguments, "PENALTY"); // Penalty for wideness

    if (!pmReadoutMaskNonfinite(readout, maskIn)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to mask non-finite pixels in readout.");
        return false;
    }

    if (psMetadataLookupBool(&mdok, config->arguments, "HAVE.PSF")) {
        assert(psf);
        assert(sources);

        int order = psMetadataLookupS32(NULL, recipe, "SPATIAL.ORDER"); // Spatial polynomial order
        float regionSize = psMetadataLookupF32(NULL, recipe, "REGION.SIZE"); // Size of iso-kernel regs
        float spacing = psMetadataLookupF32(NULL, recipe, "STAMP.SPACING"); // Typical stamp spacing
        int footprint = psMetadataLookupS32(NULL, recipe, "STAMP.FOOTPRINT"); // Stamp half-size
        float threshold = psMetadataLookupF32(NULL, recipe, "STAMP.THRESHOLD"); // Threshold for stmps
        int iter = psMetadataLookupS32(NULL, recipe, "ITER"); // Rejection iterations
        float rej = psMetadataLookupF32(NULL, recipe, "REJ"); // Rejection threshold
        pmSubtractionKernelsType type = pmSubtractionKernelsTypeFromString(
            psMetadataLookupStr(NULL, recipe, "KERNEL.TYPE")); // Kernel type
        psVector *widths = psMetadataLookupPtr(NULL, recipe, "ISIS.WIDTHS"); // ISIS Gaussian widths
        psVector *orders = psMetadataLookupPtr(NULL, recipe, "ISIS.ORDERS"); // ISIS Polynomial orders
        int inner = psMetadataLookupS32(NULL, recipe, "INNER"); // Inner radius
        int ringsOrder = psMetadataLookupS32(NULL, recipe, "RINGS.ORDER"); // RINGS polynomial order
        int binning = psMetadataLookupS32(NULL, recipe, "SPAM.BINNING"); // Binning for SPAM kernel
        float badFrac = psMetadataLookupF32(NULL, recipe, "BADFRAC"); // Maximum bad fraction
        bool optimum = psMetadataLookupBool(&mdok, recipe, "OPTIMUM"); // Derive optimum parameters?
        float optMin = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.MIN"); // Minimum width for search
        float optMax = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.MAX"); // Maximum width for search
        float optStep = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.STEP"); // Step for search
        float optThresh = psMetadataLookupF32(&mdok, recipe, "OPTIMUM.TOL"); // Tolerance for search
        int optOrder = psMetadataLookupS32(&mdok, recipe, "OPTIMUM.ORDER"); // Order for search

        // These values are specified specifically for stacking
        const char *stampsName = psMetadataLookupStr(NULL, config->arguments, "STAMPS"); // Stamps filename

        psVector *optWidths = NULL;         // Vector with FWHMs for optimum search
        if (optimum) {
            optWidths = psVectorCreate(optWidths, optMin, optMax, optStep, PS_TYPE_F32);
        }

        // Generate a fake image to match to
        float maxMag = -INFINITY;           // Maximum magnitude of sources
        for (int i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i]; // Source of interest
            if (source->psfMag > maxMag && source->psfMag <= MAG_IGNORE && !(source->mode & SOURCE_MASK)) {
                maxMag = source->psfMag;
            }
        }

        pmReadout *fake = pmReadoutAlloc(NULL); // Fake readout with target PSF

        if (!pmReadoutFakeFromSources(fake, readout->image->numCols, readout->image->numRows, sources, NULL,
                                      NULL, psf, powf(10.0, -0.4 * maxMag), 0, false)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate fake image with target PSF.");
            psFree(fake);
            psFree(optWidths);
            psFree(output);
            return false;
        }

#ifdef TESTING
        {
            psFits *fits = psFitsOpen("fake.fits", "w");
            psFitsWriteImage(fits, NULL, fake->image, 0, NULL);
            psFitsClose(fits);
        }
#endif

        // Do the image matching
        if (!pmSubtractionMatch(output, NULL, readout, fake, footprint, regionSize, spacing, threshold,
                                sources, stampsName, type, size, order, widths, orders, inner, ringsOrder,
                                binning, penalty, optimum, optWidths, optOrder, optThresh, iter, rej,
                                maskIn, maskOut, badFrac, PM_SUBTRACTION_MODE_1)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to match images.");
            psFree(fake);
            psFree(optWidths);
            psFree(output);
            return false;
        }
        psFree(fake);
        psFree(optWidths);

        // Set the variance factor
        psMetadataItem *vfItem = psMetadataLookup(readout->parent->concepts, "CELL.VARFACTOR");
        float vf = psMetadataLookupF32(NULL, output->analysis, PM_SUBTRACTION_ANALYSIS_VARFACTOR);
        if (!isfinite(vf)) {
            vf = 1.0;
        }
        if (isfinite(vfItem->data.F32)) {
            vfItem->data.F32 *= vf;
        } else {
            vfItem->data.F32 = vf;
        }

        // Replace original images with convolved
        psFree(readout->image);
        psFree(readout->mask);
        psFree(readout->weight);
        readout->image  = psMemIncrRefCounter(output->image);
        readout->mask   = psMemIncrRefCounter(output->mask);
        readout->weight = psMemIncrRefCounter(output->weight);
    } else {
        // Fake the convolution
        psRegion *region = psRegionAlloc(0, readout->image->numCols - 1, 0, readout->image->numRows - 1);
        psMetadataAddPtr(output->analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_REGION,
                         PS_DATA_REGION | PS_META_DUPLICATE_OK, "Fake subtraction region", region);
        psFree(region);
        pmSubtractionKernels *kernels = pmSubtractionKernelsPOIS(FAKE_SIZE, 0, penalty,
                                                                 PM_SUBTRACTION_MODE_1);
        // Set solution to delta function
        kernels->solution1 = psVectorAlloc(kernels->num + 2, PS_TYPE_F64);
        psVectorInit(kernels->solution1, 0.0);
        int normIndex = PM_SUBTRACTION_INDEX_NORM(kernels); // Index for normalisation
        kernels->solution1->data.F64[normIndex] = 1.0;
        psMetadataAddPtr(output->analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_KERNEL,
                         PS_DATA_UNKNOWN | PS_META_DUPLICATE_OK, "Fake subtraction kernel", kernels);
        psFree(kernels);
    }

    // Extract the regions and solutions used in the image matching
    // This stops them from being freed when we iterate back up the FPA
    *regions = psArrayAllocEmpty(ARRAY_BUFFER); // Array of regions
    {
        psString regex = NULL;          // Regular expression
        psStringAppend(&regex, "^%s$", PM_SUBTRACTION_ANALYSIS_REGION);
        psMetadataIterator *iter = psMetadataIteratorAlloc(output->analysis, PS_LIST_HEAD, regex); // Iterator
        psFree(regex);
        psMetadataItem *item = NULL;// Item from iteration
        while ((item = psMetadataGetAndIncrement(iter))) {
            assert(item->type == PS_DATA_REGION);
            *regions = psArrayAdd(*regions, ARRAY_BUFFER, item->data.V);
        }
        psFree(iter);
    }
    *kernels = psArrayAllocEmpty(ARRAY_BUFFER); // Array of kernels
    {
        psString regex = NULL;          // Regular expression
        psStringAppend(&regex, "^%s$", PM_SUBTRACTION_ANALYSIS_KERNEL);
        psMetadataIterator *iter = psMetadataIteratorAlloc(output->analysis, PS_LIST_HEAD, regex); // Iterator
        psFree(regex);
        psMetadataItem *item = NULL;// Item from iteration
        while ((item = psMetadataGetAndIncrement(iter))) {
            assert(item->type == PS_DATA_UNKNOWN);
            // Set the normalisation dimensions, since these will be otherwise unavailable when reading the
            // images by scans.
            pmSubtractionKernels *kernel = item->data.V; // Kernel used in subtraction
            kernel->numCols = readout->image->numCols;
            kernel->numRows = readout->image->numRows;

            *kernels = psArrayAdd(*kernels, ARRAY_BUFFER, kernel);
        }
        psFree(iter);
    }
    assert((*regions)->n == (*kernels)->n);

    // Renormalise the variances if desired
    if (renorm) {
        psLogMsg("ppStack", PS_LOG_INFO, "Renormalising variance map.");
        if (!pmReadoutWeightRenorm(readout, maskOut, renormMean, renormStdev, renormWidth, rng)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to renormalise variances.");
            psFree(output);
            return false;
        }
    }

    // Ensure the background value is zero
    psStats *bg = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV); // Statistics for background
    if (!psImageBackground(bg, NULL, readout->image, readout->mask, maskOut, rng)) {
        psWarning("Can't measure background for image.");
        psErrorClear();
    } else {
        psLogMsg("ppStack", PS_LOG_INFO, "Correcting convolved image background by %lf (+/- %lf)",
                 psStatsGetValue(bg, PS_STAT_ROBUST_MEDIAN), psStatsGetValue(bg, PS_STAT_ROBUST_STDEV));
        (void)psBinaryOp(readout->image, readout->image, "+",
                         psScalarAlloc(- psStatsGetValue(bg, PS_STAT_ROBUST_MEDIAN), PS_TYPE_F32));
    }
    psFree(bg);

    psFree(output);

    return true;
}
