#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                         // Enable debugging output



#ifdef TESTING
// Read a FITS image
static bool readImage(psImage **target, // Target for image
                      const char *name, // Name of FITS file
                      const pmConfig *config // Configuration
    )
{
    psString resolved = pmConfigConvertFilename(name, config, false, false); // Resolved filename
    psFits *fits = psFitsOpen(resolved, "r");
    psFree(resolved);
    if (!fits) {
        psError(PS_ERR_IO, false, "Unable to open previously produced image: %s", name);
        return false;
    }
    psImage *image = psFitsReadImage(fits, psRegionSet(0,0,0,0), 0); // Image of interest
    if (!image) {
        psError(PS_ERR_IO, false, "Unable to read previously produced image: %s", name);
        psFitsClose(fits);
        return false;
    }
    psFitsClose(fits);

    psFree(*target);
    *target = image;

    return true;
}
#endif

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

    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
    psAssert(recipe, "We've thrown an error on this before.");

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

    psString maskValStr = psMetadataLookupStr(NULL, recipe, "MASK.VAL"); // Name of bits to mask going in
    psMaskType maskVal = pmConfigMaskGet(maskValStr, config); // Bits to mask going in to pmSubtractionMatch
    psString maskPoorStr = psMetadataLookupStr(NULL, recipe, "MASK.POOR"); // Name of bits to mask for poor
    psMaskType maskPoor = pmConfigMaskGet(maskPoorStr, config); // Bits to mask for poor pixels
    psString maskBadStr = psMetadataLookupStr(NULL, recipe, "MASK.BAD"); // Name of bits to mask for bad
    psMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

    bool mdok;                          // Status of MD lookup
    bool renorm = psMetadataLookupBool(&mdok, recipe, "RENORM"); // Renormalise variances?
    float penalty = psMetadataLookupF32(NULL, ppsub, "PENALTY"); // Penalty for wideness
    int threads = psMetadataLookupS32(NULL, config->arguments, "-threads"); // Number of threads

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

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

    static int numInput = -1;            // Index of input file
    numInput++;
#ifdef TESTING
    // Read previously produced kernel
    if (psMetadataLookupBool(NULL, config->arguments, "PPSTACK.DEBUG.STACK")) {
        const char *outName = psMetadataLookupStr(NULL, config->arguments, "OUTPUT"); // Output root
        assert(outName);
        // Read convolution kernel
        {
            psString filename = NULL;   // Output filename
            psStringAppend(&filename, "%s.%d.kernel", outName, numInput);
            psString resolved = pmConfigConvertFilename(filename, config, false, false); // Resolved filename
            psFree(filename);
            psFits *fits = psFitsOpen(resolved, "r"); // FITS file for subtraction kernel
            psFree(resolved);
            if (!fits || !pmReadoutReadSubtractionKernels(output, fits)) {
                psError(PS_ERR_IO, false, "Unable to read previously produced kernel");
                psFitsClose(fits);
                return false;
            }
            psFitsClose(fits);

            // Add in variance factor
            pmSubtractionKernels *kernels = psMetadataLookupPtr(NULL, output->analysis,
                                                                PM_SUBTRACTION_ANALYSIS_KERNEL); // Kernels
            float vf = pmSubtractionVarianceFactor(kernels, 0.0, 0.0, false); // Variance factor
            psMetadataItem *vfItem = psMetadataLookup(readout->parent->concepts, "CELL.VARFACTOR");
            if (!isfinite(vf)) {
                vf = 1.0;
            }
            if (isfinite(vfItem->data.F32)) {
                vfItem->data.F32 *= vf;
            } else {
                vfItem->data.F32 = vf;
            }
        }

        // Read image, mask, weight
        const char *tempImage = psMetadataLookupStr(NULL, recipe, "TEMP.IMAGE"); // Suffix for image
        const char *tempMask = psMetadataLookupStr(NULL, recipe, "TEMP.MASK"); // Suffix for mask
        const char *tempWeight = psMetadataLookupStr(NULL, recipe, "TEMP.WEIGHT"); // Suffix for weight map
        psString imageName = NULL, maskName = NULL, weightName = NULL; // Names for convolved images
        psStringAppend(&imageName, "%s.%d.%s", outName, numInput, tempImage);
        psStringAppend(&maskName, "%s.%d.%s", outName, numInput, tempMask);
        psStringAppend(&weightName, "%s.%d.%s", outName, numInput, tempWeight);

        if (!readImage(&readout->image, imageName, config) || !readImage(&readout->mask, maskName, config) ||
            !readImage(&readout->weight, weightName, config)) {
            psError(PS_ERR_IO, false, "Unable to read previously produced image.");
            psFree(imageName);
            psFree(maskName);
            psFree(weightName);
            return false;
        }
        psFree(imageName);
        psFree(maskName);
        psFree(weightName);
    } else {
#endif

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

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

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

            float minFlux = INFINITY;       // Minimum flux for fake image
            {
                psStats *bg = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV); // Statistics for bg
                if (!psImageBackground(bg, NULL, readout->image, readout->mask, maskVal, rng)) {
                    psWarning("Can't measure background for image.");
                    psErrorClear();
                } else {
                    minFlux = 0.1 * bg->robustStdev;
                }
                psFree(bg);
            }

            // Generate a fake image to match to
            if (!isfinite(minFlux)) {
                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;
                    }
                }
                minFlux = 0.01 * powf(10.0, -0.4 * maxMag);
            }

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

            if (!pmReadoutFakeFromSources(fake, readout->image->numCols, readout->image->numRows, sources,
                                          NULL, NULL, psf, minFlux, 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
            {
                psString name = NULL;
                psStringAppend(&name, "fake_%03d.fits", numInput);
                psFits *fits = psFitsOpen(name, "w");
                psFree(name);
                psFitsWriteImage(fits, NULL, fake->image, 0, NULL);
                psFitsClose(fits);
            }
#endif

            if (threads > 0) {
                pmSubtractionThreadsInit(readout, fake);
            }

            // 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,
                                    maskVal, maskBad, maskPoor, poorFrac, 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);

            if (threads > 0) {
                pmSubtractionThreadsFinalize(readout, fake);
            }

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

            readout->analysis = psMetadataCopy(readout->analysis, output->analysis);
        } 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);
        }

        // Write convolution kernel
        {
            const char *outName = psMetadataLookupStr(NULL, config->arguments, "OUTPUT"); // Output root
            assert(outName);

            psString filename = NULL;   // Output filename
            psStringAppend(&filename, "%s.%d.kernel", outName, numInput);
            psString resolved = pmConfigConvertFilename(filename, config, true, false); // Resolved filename
            psFree(filename);
            psFits *fits = psFitsOpen(resolved, "w"); // FITS file for subtraction kernel
            psFree(resolved);
            pmReadoutWriteSubtractionKernels(output, fits);
            psFitsClose(fits);
        }
#ifdef TESTING
    }
#endif

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

    // Correct the variance for the chi^2
    {
        *chi2 = 0.0;
        int num = 0;                    // Number of measurements of chi^2
        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);
            pmSubtractionKernels *kernels = item->data.V; // Convolution kernels
            *chi2 += kernels->mean;
            num++;
        }
        psFree(iter);

        float vf = psMetadataLookupF32(NULL, readout->parent->concepts, "CELL.VARFACTOR"); // Variance factor
        *chi2 /= vf * num;
    }

    // Renormalise the variances if desired
    if (renorm) {
        // Statistics for renormalisation
        psStatsOptions renormMean = psStatsOptionFromString(psMetadataLookupStr(&mdok, recipe,
                                                                                "RENORM.MEAN"));
        psStatsOptions renormStdev = psStatsOptionFromString(psMetadataLookupStr(&mdok, recipe,
                                                                                 "RENORM.STDEV"));
        int renormWidth = psMetadataLookupS32(&mdok, recipe, "RENORM.WIDTH"); // Width for renormalisation box

        psLogMsg("ppStack", PS_LOG_INFO, "Renormalising variance map.");
        if (!pmReadoutWeightRenorm(readout, maskBad, 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, maskVal | maskBad, 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;
}
