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

#include <stdio.h>
#include <pslib.h>
#include <psmodules.h>
#include <psphot.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 FAINT_SOURCE_FRAC 1.0e-4         // Set minimum flux to this fraction of faintest source flux
#define COVAR_FRAC 0.01                 // Truncation fraction for covariance matrix

//#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

// Get coordinates from a source
static void coordsFromSource(float *x, float *y, const pmSource *source)
{
    assert(x && y);
    assert(source);

    if (source->modelPSF) {
        *x = source->modelPSF->params->data.F32[PM_PAR_XPOS];
        *y = source->modelPSF->params->data.F32[PM_PAR_YPOS];
    } else {
        *x = source->peak->xf;
        *y = source->peak->yf;
    }
    return;
}

static psArray *stackSourcesFilter(psArray *sources, // Source list to filter
                                   int exclusion // Exclusion zone, pixels
    )
{
    psAssert(sources && sources->n > 0, "Require array of sources");
    if (exclusion <= 0) {
        return psMemIncrRefCounter(sources);
    }

    int num = sources->n;               // Number of sources
    psVector *x = psVectorAlloc(num, PS_TYPE_F32), *y = psVectorAlloc(num, PS_TYPE_F32); // Coordinates
    int numGood = 0;                    // Number of good sources
    for (int i = 0; i < num; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (!source) {
            continue;
        }
        coordsFromSource(&x->data.F32[numGood], &y->data.F32[numGood], source);
        numGood++;
    }
    x->n = y->n = numGood;

    psTree *tree = psTreePlant(2, 2, x, y); // kd tree

    psArray *filtered = psArrayAllocEmpty(numGood); // Filtered list of sources
    psVector *coords = psVectorAlloc(2, PS_TYPE_F64); // Coordinates of source
    int numFiltered = 0;                // Number of filtered sources
    for (int i = 0; i < num; i++) {
        pmSource *source = sources->data[i]; // Source of interest
        if (!source) {
            continue;
        }
        float xSource, ySource;         // Coordinates of source
        coordsFromSource(&xSource, &ySource, source);

        coords->data.F64[0] = xSource;
        coords->data.F64[1] = ySource;

        long numWithin = psTreeWithin(tree, coords, exclusion); // Number within exclusion zone
        psTrace("ppStack", 9, "Source at %.0lf,%.0lf has %ld sources in exclusion zone",
                coords->data.F64[0], coords->data.F64[1], numWithin);
        if (numWithin == 1) {
            // Only itself inside the exclusion zone
            filtered = psArrayAdd(filtered, filtered->n, source);
        } else {
            numFiltered++;
        }
    }
    psFree(coords);
    psFree(tree);

    psLogMsg("ppStack", PS_LOG_INFO, "Filtered out %d of %d sources", numFiltered, numGood);

    return filtered;
}

// Add background into the fake image
// Based on ppSubBackground()
static psImage *stackBackgroundModel(pmReadout *ro, // Readout for which to generate background model
                                     const pmConfig *config // Configuration
    )
{
    psAssert(ro && ro->image, "Need readout image");
    psAssert(config, "Need configuration");

    psImage *image = ro->image;         // Image of interest
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    psMetadata *ppStackRecipe = psMetadataLookupPtr(NULL, config->recipes, PPSTACK_RECIPE);
    psAssert(ppStackRecipe, "Need PPSTACK recipe");
    psMetadata *psphotRecipe = psMetadataLookupPtr(NULL, config->recipes, PSPHOT_RECIPE);
    psAssert(psphotRecipe, "Need PSPHOT recipe");

    psString maskBadStr = psMetadataLookupStr(NULL, ppStackRecipe, "MASK.BAD");// Name of bits to mask for bad
    psMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

    // user-defined masks to test for good/bad pixels (build from recipe list if not yet set)
    psMetadataAddU8(psphotRecipe, PS_LIST_TAIL, "MASK.PSPHOT", PS_META_REPLACE, "user-defined mask", maskBad);

    psImage *binned = psphotBackgroundModel(ro, config); // Binned background model
    psImageBinning *binning = psMetadataLookupPtr(NULL, ro->analysis,
                                                  "PSPHOT.BACKGROUND.BINNING"); // Binning for model
    psAssert(binning, "Need binning parameters");
    psImage *unbinned = psImageAlloc(numCols, numRows, PS_TYPE_F32); // Unbinned background model
    if (!psImageUnbin(unbinned, binned, binning)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to unbin background model");
        psFree(unbinned);
        return NULL;
    }

    // XXX should these really be here?? (probably not...)
    // pmFPAfileDropInternal(config->files, "PSPHOT.BACKMDL");
    // pmFPAfileDropInternal(config->files, "PSPHOT.BACKMDL.STDEV");

     return unbinned;
}


bool ppStackMatch(pmReadout *readout, ppStackOptions *options, int index, const pmConfig *config)
{
    assert(readout);
    assert(options);
    assert(config);

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

    float deconvLimit = psMetadataLookupF32(NULL, recipe, "DECONV.LIMIT"); // Limit on deconvolution fraction

    // 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
    psImageMaskType 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
    psImageMaskType maskPoor = pmConfigMaskGet(maskPoorStr, config); // Bits to mask for poor pixels
    psString maskBadStr = psMetadataLookupStr(NULL, recipe, "MASK.BAD"); // Name of bits to mask for bad
    psImageMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

    bool mdok;                          // Status of MD lookup
    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;
    }

    // Match the PSF
    if (options->convolve) {
        pmReadout *conv = pmReadoutAlloc(NULL); // Conv readout, for holding results temporarily
#ifdef TESTING
        // This is a hack to use the temporary convolved images and kernel generated previously.
        // This makes the 'matching' operation much faster, allowing debugging of the stack process easier.
        // It implicitly assumes the output root name is the same between invocations.

        // Read previously produced kernel
        if (psMetadataLookupBool(NULL, config->arguments, "PPSTACK.DEBUG.STACK")) {
            pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.CONV.KERNEL", index);
            psAssert(file, "Require file");

            pmFPAview *view = pmFPAviewAlloc(0); // View to readout of interest
            view->chip = view->cell = view->readout = 0;
            psString filename = pmFPAfileNameFromRule(file->filerule, file, view); // Filename of interest

            // Read convolution kernel
            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(conv, fits)) {
                psError(PS_ERR_IO, false, "Unable to read previously produced kernel");
                psFitsClose(fits);
                return false;
            }
            psFitsClose(fits);

            if (!readImage(&readout->image, options->imageNames->data[index], config) ||
                !readImage(&readout->mask, options->maskNames->data[index], config) ||
                !readImage(&readout->variance, options->varianceNames->data[index], config)) {
                psError(PS_ERR_IO, false, "Unable to read previously produced image.");
                return false;
            }

            psRegion *region = psMetadataLookupPtr(NULL, conv->analysis,
                                                   PM_SUBTRACTION_ANALYSIS_REGION); // Convolution region

            pmSubtractionAnalysis(readout->analysis, kernels, region,
                                  readout->image->numCols, readout->image->numRows);

            psKernel *kernel = pmSubtractionKernel(kernels, 0.0, 0.0, false); // Convolution kernel
            psKernel *covar = psImageCovarianceCalculate(kernel, readout->covariance); // Covariance matrix
            psFree(readout->covariance);
            readout->covariance = covar;
            psFree(kernel);
        } else {
#endif

            // Normal operations here
            psAssert(options->psf, "Require target PSF");
            psAssert(options->sourceLists && options->sourceLists->data[index], "Require source list");

            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 stride = psMetadataLookupS32(NULL, ppsub, "STRIDE"); // Size of convolution patches
            int iter = psMetadataLookupS32(NULL, ppsub, "ITER"); // Rejection iterations
            float rej = psMetadataLookupF32(NULL, ppsub, "REJ"); // Rejection threshold
            float sysError = psMetadataLookupF32(NULL, ppsub, "SYS"); // Relative systematic error in kernel
            const char *typeStr = psMetadataLookupStr(NULL, ppsub, "KERNEL.TYPE"); // Kernel type
            pmSubtractionKernelsType type = pmSubtractionKernelsTypeFromString(typeStr); // 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);
            }

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

            // For the sake of stamps, remove nearby sources
            psArray *stampSources = stackSourcesFilter(options->sourceLists->data[index],
                                                       footprint); // Filtered list of sources

            if (!pmReadoutFakeFromSources(fake, readout->image->numCols, readout->image->numRows,
                                          stampSources, NULL, NULL, options->psf, NAN, footprint + size,
                                          false, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to generate fake image with target PSF.");
                psFree(fake);
                psFree(optWidths);
                psFree(conv);
                return false;
            }

            fake->mask = psImageCopy(NULL, readout->mask, PS_TYPE_IMAGE_MASK);

            // Add the background into the target image
            psImage *bgImage = stackBackgroundModel(readout, config); // Image of background
            psBinaryOp(fake->image, fake->image, "+", bgImage);
            psFree(bgImage);

#ifdef TESTING
            {
                pmHDU *hdu = pmHDUFromCell(readout->parent);
                psString name = NULL;
                psStringAppend(&name, "fake_%03d.fits", index);
                pmStackVisualPlotTestImage(fake->image, name);
                psFits *fits = psFitsOpen(name, "w");
                psFree(name);
                psFitsWriteImage(fits, hdu->header, fake->image, 0, NULL);
                psFitsClose(fits);
            }
            {
                pmHDU *hdu = pmHDUFromCell(readout->parent);
                psString name = NULL;
                psStringAppend(&name, "real_%03d.fits", index);
                pmStackVisualPlotTestImage(readout->image, name);
                psFits *fits = psFitsOpen(name, "w");
                psFree(name);
                psFitsWriteImage(fits, hdu->header, readout->image, 0, NULL);
                psFitsClose(fits);
            }
#endif

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

            // Do the image matching
            pmSubtractionKernels *kernel = psMetadataLookupPtr(&mdok, readout->analysis,
                                                               PM_SUBTRACTION_ANALYSIS_KERNEL); // Conv kernel
            if (kernel) {
                if (!pmSubtractionMatchPrecalc(conv, NULL, readout, fake, readout->analysis,
                                               stride, sysError, maskVal, maskBad, maskPoor,
                                               poorFrac, badFrac)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to convolve images.");
                    psFree(fake);
                    psFree(optWidths);
                    psFree(stampSources);
                    psFree(conv);
                    if (threads > 0) {
                        pmSubtractionThreadsFinalize(readout, fake);
                    }
                    return false;
                }
            } else {
                if (!pmSubtractionMatch(conv, NULL, readout, fake, footprint, stride, regionSize, spacing,
                                        threshold, stampSources, stampsName, type, size, order, widths,
                                        orders, inner, ringsOrder, binning, penalty,
                                        optimum, optWidths, optOrder, optThresh, iter, rej, sysError,
                                        maskVal, maskBad, maskPoor, poorFrac, badFrac,
                                        PM_SUBTRACTION_MODE_1)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to match images.");
                    psFree(fake);
                    psFree(optWidths);
                    psFree(stampSources);
                    psFree(conv);
                    if (threads > 0) {
                        pmSubtractionThreadsFinalize(readout, fake);
                    }
                    return false;
                }
            }

#ifdef TESTING
            {
                pmHDU *hdu = pmHDUFromCell(readout->parent);
                psString name = NULL;
                psStringAppend(&name, "conv_%03d.fits", index);
                pmStackVisualPlotTestImage(conv->image, name);
                psFits *fits = psFitsOpen(name, "w");
                psFree(name);
                psFitsWriteImage(fits, hdu->header, conv->image, 0, NULL);
                psFitsClose(fits);
            }
            {
                pmHDU *hdu = pmHDUFromCell(readout->parent);
                psString name = NULL;
                psStringAppend(&name, "diff_%03d.fits", index);
                pmStackVisualPlotTestImage(fake->image, name);
                psFits *fits = psFitsOpen(name, "w");
                psFree(name);
                psBinaryOp(fake->image, conv->image, "-", fake->image);
                psFitsWriteImage(fits, hdu->header, fake->image, 0, NULL);
                psFitsClose(fits);
            }
#endif

            psFree(fake);
            psFree(optWidths);
            psFree(stampSources);

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

            // Replace original images with convolved
            psFree(readout->image);
            psFree(readout->mask);
            psFree(readout->variance);
            psFree(readout->covariance);
            readout->image  = psMemIncrRefCounter(conv->image);
            readout->mask   = psMemIncrRefCounter(conv->mask);
            readout->variance = psMemIncrRefCounter(conv->variance);
            readout->covariance = psImageCovarianceTruncate(conv->covariance, COVAR_FRAC);
#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
        psArray *regions = options->regions->data[index] = psArrayAllocEmpty(ARRAY_BUFFER); // Match regions
        {
            psString regex = NULL;          // Regular expression
            psStringAppend(&regex, "^%s$", PM_SUBTRACTION_ANALYSIS_REGION);
            psMetadataIterator *iter = psMetadataIteratorAlloc(conv->analysis, PS_LIST_HEAD, regex);
            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);
        }
        psArray *kernels = options->kernels->data[index] = psArrayAllocEmpty(ARRAY_BUFFER); // Match kernels
        {
            psString regex = NULL;          // Regular expression
            psStringAppend(&regex, "^%s$", PM_SUBTRACTION_ANALYSIS_KERNEL);
            psMetadataIterator *iter = psMetadataIteratorAlloc(conv->analysis, PS_LIST_HEAD, regex);
            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);
        }
        psAssert((regions)->n == (kernels)->n, "Number of match regions and kernels should match");

        // Record chi^2
        {
            double sum = 0.0;           // Sum of chi^2
            int num = 0;                // Number of measurements of chi^2
            psString regex = NULL;      // Regular expression
            psStringAppend(&regex, "^%s$", PM_SUBTRACTION_ANALYSIS_KERNEL);
            psMetadataIterator *iter = psMetadataIteratorAlloc(conv->analysis, PS_LIST_HEAD, regex);
            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
                sum += kernels->mean;
                num++;
            }
            psFree(iter);
            options->matchChi2->data.F32[index] = sum / (psImageCovarianceFactor(readout->covariance) * num);
        }

        // Reject image completely if the maximum deconvolution fraction exceeds the limit
        float deconv = psMetadataLookupF32(NULL, conv->analysis,
                                           PM_SUBTRACTION_ANALYSIS_DECONV_MAX); // Max deconvolution fraction
        if (deconv > deconvLimit) {
            psWarning("Maximum deconvolution fraction (%f) exceeds limit (%f) --- rejecting\n",
                      deconv, deconvLimit);
            psFree(conv);
            return NULL;
        }

        readout->analysis = psMetadataCopy(readout->analysis, conv->analysis);

        psFree(conv);
    } else {
        // Match the normalisation
        float norm = powf(10.0, -0.4 * options->norm->data.F32[index]); // Normalisation
        psBinaryOp(readout->image, readout->image, "*", psScalarAlloc(norm, PS_TYPE_F32));
        psBinaryOp(readout->variance, readout->variance, "*", psScalarAlloc(PS_SQR(norm), PS_TYPE_F32));
    }

    // Ensure the background value is zero
    psStats *bg = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV); // Statistics for background
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS); // Random number generator
    if (!psImageBackground(bg, NULL, readout->image, readout->mask, maskVal | maskBad, rng)) {
      psWarning("Can't measure background for image.");
      psErrorClear();
    } else {
      if (!psMetadataLookupBool(NULL, config->arguments, "PPSTACK.SKIP.BG.SUB")) {
        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));
      }
    }


    // Measure the variance level for the weighting
    if (!psImageBackground(bg, NULL, readout->variance, readout->mask, maskVal | maskBad, rng)) {
        psError(PS_ERR_UNKNOWN, false, "Can't measure mean variance for image.");
        psFree(rng);
        psFree(bg);
        return false;
    }
    options->weightings->data.F32[index] = 1.0 / (psStatsGetValue(bg, PS_STAT_ROBUST_MEDIAN) *
                                                  psImageCovarianceFactor(readout->covariance));
    psMetadataAddF32(readout->analysis, PS_LIST_TAIL, "PPSTACK.WEIGHTING", 0,
                     "Weighting by 1/noise^2 for stack", options->weightings->data.F32[index]);

    psFree(rng);
    psFree(bg);

#ifdef TESTING
    {
        pmHDU *hdu = pmHDUFromCell(readout->parent);
        psString name = NULL;
        psStringAppend(&name, "convolved_%03d.fits", index);
        pmStackVisualPlotTestImage(readout->image, name);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, hdu->header, readout->image, 0, NULL);
        psFitsClose(fits);
    }
#endif

    return true;
}
