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

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

#include "ppSub.h"

#define WCS_TOLERANCE 0.001             // Tolerance for WCS
//#define TESTING                         // For test output

#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


// Copy every instance of a single keyword from one metadata to another
static bool metadataCopySingle(psMetadata *target, psMetadata *source, const char *name)
{
    PS_ASSERT_METADATA_NON_NULL(target, false);
    PS_ASSERT_METADATA_NON_NULL(source, false);
    PS_ASSERT_STRING_NON_EMPTY(name, false);

    psString regex = NULL;      // Regular expression
    psStringAppend(&regex, "^%s$", name);
    psMetadataIterator *iter = psMetadataIteratorAlloc(source, PS_LIST_HEAD, regex); // Iterator
    psFree(regex);
    psMetadataItem *item;       // Item from iteration
    while ((item = psMetadataGetAndIncrement(iter))) {
        psMetadataAddItem(target, item, PS_LIST_TAIL, PS_META_DUPLICATE_OK);
    }
    psFree(iter);

    return true;
}


bool ppSubReadout(pmConfig *config, psMetadata *stats, const pmFPAview *view)
{
    pmReadout *inRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT"); // Input readout
    pmReadout *refRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF"); // Reference readout
    pmReadout *sourcesRO = pmFPAfileThisReadout(config->files, view, "PPSUB.SOURCES"); // Readout with sources
    pmReadout *inConv = pmReadoutAlloc(NULL); // Convolved version of input
    pmReadout *refConv = pmReadoutAlloc(NULL); // Convolved version of reference
    pmCell *outCell = pmFPAfileThisCell(config->files, view, "PPSUB.OUTPUT"); // Output cell
    pmReadout *outRO = pmReadoutAlloc(outCell); // Output readout: subtraction
    pmFPA *outFPA = outCell->parent->parent; // Output FPA
    pmHDU *outHDU = outFPA->hdu; // Output HDU
    if (!outHDU->header) {
        outHDU->header = psMetadataAlloc();
    }

    psImage *input = inRO->image;       // Input image
    psImage *reference = refRO->image;  // Reference image
    PS_ASSERT_IMAGES_SIZE_EQUAL(input, reference, false);

    // Look up appropriate values
    bool mdok;                          // Status of MD lookup
    int size = psMetadataLookupS32(NULL, config->arguments, "KERNEL.SIZE"); // Kernel half-size
    int order = psMetadataLookupS32(NULL, config->arguments, "SPATIAL.ORDER"); // Spatial polynomial order
    float regionSize = psMetadataLookupF32(NULL, config->arguments, "REGION.SIZE"); // Size of iso-kernel regs
    float spacing = psMetadataLookupF32(NULL, config->arguments, "STAMP.SPACING"); // Typical stamp spacing
    int footprint = psMetadataLookupS32(NULL, config->arguments, "STAMP.FOOTPRINT"); // Stamp half-size
    float threshold = psMetadataLookupF32(NULL, config->arguments, "STAMP.THRESHOLD"); // Threshold for stmps
    int iter = psMetadataLookupS32(NULL, config->arguments, "ITER"); // Rejection iterations
    float rej = psMetadataLookupF32(NULL, config->arguments, "REJ"); // Rejection threshold
    pmSubtractionKernelsType type = psMetadataLookupS32(NULL, config->arguments,
                                                        "KERNEL.TYPE"); // Kernel type
    bool reverse = psMetadataLookupBool(NULL, config->arguments, "REVERSE"); // Reverse sense of subtraction?
    psVector *widths = psMetadataLookupPtr(NULL, config->arguments, "ISIS.WIDTHS"); // ISIS Gaussian widths
    psVector *orders = psMetadataLookupPtr(NULL, config->arguments, "ISIS.ORDERS"); // ISIS Polynomial orders
    int inner = psMetadataLookupS32(NULL, config->arguments, "INNER"); // Inner radius
    int ringsOrder = psMetadataLookupS32(NULL, config->arguments, "RINGS.ORDER"); // RINGS polynomial order
    int binning = psMetadataLookupS32(NULL, config->arguments, "SPAM.BINNING"); // Binning for SPAM kernel
    float penalty = psMetadataLookupF32(NULL, config->arguments, "PENALTY"); // Penalty for wideness
    psMaskType maskIn = pmConfigMaskGet("MASK.VALUE", config); // Bits to mask going in to pmSubtractionMatch
    psMaskType maskOut = pmConfigMaskGet("BLANK", config); // Bits to mask after pmSubtractionMatch
    float badFrac = psMetadataLookupF32(NULL, config->arguments, "BADFRAC"); // Maximum bad fraction
    const char *stampsName = psMetadataLookupStr(&mdok, config->arguments, "STAMPS"); // Filename for stamps

    bool optimum = psMetadataLookupBool(&mdok, config->arguments, "OPTIMUM"); // Derive optimum parameters?
    float optMin = psMetadataLookupF32(&mdok, config->arguments, "OPTIMUM.MIN"); // Minimum width for search
    float optMax = psMetadataLookupF32(&mdok, config->arguments, "OPTIMUM.MAX"); // Maximum width for search
    float optStep = psMetadataLookupF32(&mdok, config->arguments, "OPTIMUM.STEP"); // Step for search
    float optThresh = psMetadataLookupF32(&mdok, config->arguments, "OPTIMUM.TOL"); // Tolerance for search
    int optOrder = psMetadataLookupS32(&mdok, config->arguments, "OPTIMUM.ORDER"); // Order for search
    bool dual = psMetadataLookupBool(&mdok, config->arguments, "DUAL"); // Dual convolution?
    bool renorm = psMetadataLookupBool(&mdok, config->arguments, "RENORM"); // Renormalise weights?
    int renormWidth = psMetadataLookupS32(&mdok, config->arguments, "RENORM.WIDTH"); // Width for renormalise

    pmSubtractionMode mode = dual ? PM_SUBTRACTION_MODE_DUAL : PM_SUBTRACTION_MODE_UNSURE; // Subtraction mode

    // Generate masks if they don't exist
    int numCols = input->numCols, numRows = input->numRows; // Image dimensions
    if (!inRO->mask) {
        if (psMetadataLookupBool(NULL, config->arguments, "MASK.GENERATE")) {
            pmReadoutSetMask(inRO, pmConfigMaskGet("SAT", config), pmConfigMaskGet("BAD", config));
        } else {
            inRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
            psImageInit(inRO->mask, 0);
        }
    }
    if (!refRO->mask) {
        if (psMetadataLookupBool(NULL, config->arguments, "MASK.GENERATE")) {
            pmReadoutSetMask(refRO, pmConfigMaskGet("SAT", config), pmConfigMaskGet("BAD", config));
        } else {
            refRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
            psImageInit(refRO->mask, 0);
        }
    }

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

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

    psArray *sources = NULL;            // Sources in image
    if (sourcesRO) {
        sources = psMetadataLookupPtr(&mdok, sourcesRO->analysis, "PSPHOT.SOURCES");
    }

    if (!pmSubtractionMatch(inConv, refConv, inRO, refRO, footprint, regionSize, spacing, threshold, sources,
                            stampsName, type, size, order, widths, orders, inner, ringsOrder, binning,
                            penalty, optimum, optWidths, optOrder, optThresh, iter, rej, maskIn,
                            maskOut, badFrac, mode)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to match images.");
        psFree(inConv);
        psFree(refConv);
        psFree(outRO);
        return false;
    }
    psFree(optWidths);

    // Add kernel descrption to header
    pmSubtractionKernels *kernels = psMetadataLookupPtr(&mdok, inConv->analysis,
                                                        PM_SUBTRACTION_ANALYSIS_KERNEL); // Subtraction kernel
    if (!kernels) {
        kernels = psMetadataLookupPtr(&mdok, refConv->analysis, PM_SUBTRACTION_ANALYSIS_KERNEL);
    }
    if (!kernels) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find SUBTRACTION.KERNEL");
        psFree(inConv);
        psFree(refConv);
        psFree(outRO);
        return false;
    }
    psMetadataAddStr(outHDU->header, PS_LIST_TAIL, "PPSUB.KERNEL", 0,
                     "Subtraction kernel", kernels->description);


    // Statistics on the matching
    if (stats) {
        if (psMetadataLookup(inConv->analysis, PM_SUBTRACTION_ANALYSIS_MODE)) {
            metadataCopySingle(stats, inConv->analysis, PM_SUBTRACTION_ANALYSIS_MODE);
            metadataCopySingle(stats, inConv->analysis, PM_SUBTRACTION_ANALYSIS_STAMPS_NUM);
            metadataCopySingle(stats, inConv->analysis, PM_SUBTRACTION_ANALYSIS_STAMPS_RMS);
        }
        if (psMetadataLookup(refConv->analysis, PM_SUBTRACTION_ANALYSIS_MODE)) {
            metadataCopySingle(stats, refConv->analysis, PM_SUBTRACTION_ANALYSIS_MODE);
            metadataCopySingle(stats, refConv->analysis, PM_SUBTRACTION_ANALYSIS_STAMPS_NUM);
            metadataCopySingle(stats, refConv->analysis, PM_SUBTRACTION_ANALYSIS_STAMPS_RMS);
        }
    }

#ifdef TESTING
    psImage *kernelImage = psMetadataLookupPtr(&mdok, inConv->analysis,
                                               "SUBTRACTION.KERNEL.IMAGE"); // Image of the kernels
    if (!kernelImage) {
        kernelImage = psMetadataLookupPtr(&mdok, refConv->analysis, "SUBTRACTION.KERNEL.IMAGE");
    }
    psFits *fits = psFitsOpen("kernel.fits", "w");
    psFitsWriteImage(fits, NULL, kernelImage, 0, NULL);
    psFitsClose(fits);
#endif

    // Subtraction is: minuend - subtrahend
    pmReadout *minuend = inConv;
    pmReadout *subtrahend = refConv;

    if (reverse) {
        pmReadout *temp = subtrahend;
        subtrahend = minuend;
        minuend = temp;
    }

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

    outRO->mask = (psImage*)psBinaryOp(outRO->mask, inConv->mask, "|", refConv->mask);
    outRO->data_exists = outCell->data_exists = outCell->parent->data_exists = true; // It'll be there soon

    // Photometry is to be performed in two stages:
    // 1. Measure the PSF using the PSF-matched images
    // 2. Find and measure sources on the subtracted image

    // Photometry stage 1: measure the PSF
    pmPSF *psf = NULL;                  // PSF for photometry
    if (psMetadataLookupBool(NULL, config->arguments, "PHOTOMETRY")) {
        outRO->image = psImageCopy(outRO->image, minuend->image, PS_TYPE_F32);
        if (minuend->weight) {
            outRO->weight = psImageCopy(outRO->weight, minuend->weight, PS_TYPE_F32);
        }

        pmFPAfile *photFile = psMetadataLookupPtr(NULL, config->files, "PSPHOT.INPUT");
        pmFPACopy(photFile->fpa, outRO->parent->parent->parent);

        // We have a list of sources, so we could use those to redetermine the PSF model.
        // Until Gene makes the necessary adaptations to psphot, we will simply redetermine the PSF model from
        // scratch.

        // Need to ensure aperture residual is not calculated
        psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PSPHOT_RECIPE); // Recipe
        if (!recipe) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find %s recipe.", PSPHOT_RECIPE);
            return false;
        }
        const char *breakpoint = psMetadataLookupStr(&mdok, recipe, "BREAK_POINT"); // Current break point
        psMetadataAddStr(recipe, PS_LIST_TAIL, "BREAK_POINT", PS_META_REPLACE,
                         "ALTERED break point for psphot operations", "PSFMODEL");

        if (!psphotReadout(config, view)) {
            psWarning("Unable to perform photometry on subtracted image.");
            psErrorStackPrint(stderr, "Error stack from photometry:");
            psErrorClear();
        }

        if (breakpoint) {
            psMetadataAddStr(recipe, PS_LIST_TAIL, "BREAK_POINT", PS_META_REPLACE,
                             "RESTORED break point for psphot operations", breakpoint);
        }

        // Blow away the sources psphot found --- they're irrelevant for the subtraction
        pmReadout *photRO = pmFPAviewThisReadout(view, photFile->fpa); // Readout with sources
        psMetadataRemoveKey(photRO->analysis, "PSPHOT.SOURCES");
        psMetadataRemoveKey(photRO->analysis, "PSPHOT.HEADER");

        psf = psMemIncrRefCounter(psMetadataLookupPtr(NULL, photRO->parent->parent->analysis, "PSPHOT.PSF"));
        if (!psf) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find PSF from psphot");
            return false;
        }

        pmCell *photCell = pmFPAfileThisCell(config->files, view, "PSPHOT.INPUT");
        pmCellFreeReadouts(photCell);
    }

    // Do the actual subtraction
    outRO->image = (psImage*)psBinaryOp(outRO->image, minuend->image, "-", subtrahend->image);
    if (inConv->weight && refConv->weight) {
        outRO->weight = (psImage*)psBinaryOp(outRO->weight, inConv->weight, "+", refConv->weight);
    }
    outRO->data_exists = outCell->data_exists = outCell->parent->data_exists = true;

    pmReadoutMaskApply(outRO, maskOut);

    psFree(inConv);
    psFree(refConv);

#ifdef TESTING
    for (int y = 0; y < outRO->image->numRows; y++) {
        for (int x = 0; x < outRO->image->numCols; x++) {
            if (isnan(outRO->image->data.F32[y][x]) && !(outRO->mask->data.U8[y][x] & maskOut)) {
                printf("Unmasked NAN at %d %d --> %d\n", x, y, outRO->mask->data.U8[y][x]);
            }
        }
    }
#endif

    if (!pmFPACopyConcepts(outCell->parent->parent, inRO->parent->parent->parent)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to copy concepts from input to output.");
        psFree(outRO);
        return false;
    }

    if (renorm) {
        psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0); // Random number generator
        if (!pmReadoutWeightRenorm(outRO, maskOut, PS_STAT_ROBUST_MEDIAN, PS_STAT_ROBUST_STDEV,
                                   renormWidth, rng)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to renormalise weights for photometry.");
            psFree(outRO);
            return false;
        }
        psFree(rng);
    }

#if 0
    // XXX Under development

    // QA: photometry of known sources with measured PSF
    if (sourcesRO && psMetadataLookupBool(NULL, config->arguments, "PHOTOMETRY")) {
        sources = psMetadataLookupPtr(&mdok, sourcesRO->analysis, "PSPHOT.SOURCES");



        // For each source, need to:
        // * generate appropriate model
        // * select pixels
        // * define fit radius
        // Don't need to measure sky (psphotFitSourcesLinear does that)

        psArray *dummy = psArrayAlloc(1); // Dummy array of sources, for psphotFitSourcesLinear with 1 source
        for (long i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i];
            if (!source || (source->mode & SOURCE_MASK)) {
                continue;
            }

            float origMag = source->psfMag; // Original magnitude

            // Coordinates of source
            float x = source->modelPSF->params->data.F32[PM_PAR_XPOS];
            float y = source->modelPSF->params->data.F32[PM_PAR_YPOS];

            model = pmModelFromPSFforXY(psf, x, y, NAN);
            if (!fakeModel) {
                psWarning("Can't generate model");
                continue;
            }
            if (!pmSourceDefinePixels(fakeSource, readout, x, y, fakeRadius)) {
                psErrorClear();
                continue;
            }
pmModel *pmModelFromPSFforXY (
    const pmPSF *psf,
    float Xo,
    float Yo,
    float Io
    );

bool pmSourceDefinePixels(
    pmSource *mySource,                 ///< source to be re-defined
    const pmReadout *readout,  ///< base the source on this readout
    psF32 x,                            ///< center coords of source
    psF32 y,                            ///< center coords of source
    psF32 Radius                        ///< size of box on source
);


bool psphotFitSourcesLinear (pmReadout *readout, psArray *sources, psMetadata *recipe, pmPSF *psf, bool final) {

    }
#endif

    // Photometry stage 2: find and measure sources on the subtracted image
    if (psMetadataLookupBool(NULL, config->arguments, "PHOTOMETRY")) {
        // The PSF should already be stored for the readout
        pmFPAfile *photFile = psMetadataLookupPtr(NULL, config->files, "PSPHOT.INPUT");
        pmFPACopy(photFile->fpa, outRO->parent->parent->parent);

        pmReadout *psfRO = pmFPAfileThisReadout(config->files, view, "PSPHOT.PSF.LOAD");
        if (!psfRO) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to find readout for file PSPHOT.PSF.LOAD");
            return false;
        }
        psMetadataAddPtr(psfRO->parent->parent->analysis, PS_LIST_TAIL, "PSPHOT.PSF", PS_DATA_UNKNOWN,
                         "PSF from matched addition", psf);
        psFree(psf);

        // Need to ensure aperture residual is not calculated
        psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PSPHOT_RECIPE); // Recipe
        psMetadataItem *item = psMetadataLookup(recipe, "MEASURE.APTREND"); // Item determining aptrend
        if (!item) {
            psWarning("Unable to find MEASURE.APTREND in psphot recipe");
            psErrorClear();
        } else {
            item->data.B = false;
        }

        if (!psphotReadout(config, view)) {
            psWarning("Unable to perform photometry on subtracted image.");
            psErrorStackPrint(stderr, "Error stack from photometry:");
            psErrorClear();
        }

        pmFPAfileActivate(config->files, false, "PSPHOT.INPUT");
        pmFPAfileActivate(config->files, false, "PSPHOT.LOAD.PSF");

        if (stats) {
            pmReadout *photRO = pmFPAviewThisReadout(view, photFile->fpa); // Readout with the sources
            psArray *sources = psMetadataLookupPtr(NULL, photRO->analysis, "PSPHOT.SOURCES"); // Sources
            psMetadataAddS32(stats, PS_LIST_TAIL, "NUM_SOURCES", 0, "Number of sources detected", sources->n);
        }
    }

    // Copy astrometry over
    pmFPA *refFPA = refRO->parent->parent->parent; // Reference FPA
    pmHDU *refHDU = refFPA->hdu;        // Reference HDU
    if (!outHDU || !refHDU) {
        psWarning("Unable to find HDU at FPA level to copy astrometry.");
    } else {
        if (!pmAstromReadWCS(outFPA, outCell->parent, refHDU->header, 1.0)) {
            psWarning("Unable to read WCS astrometry from reference FPA.");
            psErrorClear();
        } else if (!pmAstromWriteWCS(outHDU->header, outFPA, outCell->parent, WCS_TOLERANCE)) {
            psWarning("Unable to write WCS astrometry to output FPA.");
            psErrorClear();
        }
    }

    // Add additional data to the header
    pmFPAfile *refFile = psMetadataLookupPtr(NULL, config->files, "PPSUB.REF"); // Reference file
    pmFPAfile *inFile = psMetadataLookupPtr(NULL, config->files, "PPSUB.INPUT"); // Input file
    psMetadataAddStr(outHDU->header, PS_LIST_TAIL, "PPSUB.REFERENCE", 0,
                     "Subtraction reference", refFile->filename);
    psMetadataAddStr(outHDU->header, PS_LIST_TAIL, "PPSUB.INPUT", 0,
                     "Subtraction input", inFile->filename);

    psFree(outRO);

    return true;
}
