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

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

#include "ppStack.h"
#include "ppStackLoop.h"

// Update the value of a concept
#define UPDATE_CONCEPT(SOURCE, NAME, VALUE) { \
    psMetadataItem *item = psMetadataLookup(SOURCE->concepts, NAME); \
    psAssert(item, "Concept should be present"); \
    psAssert(item->type == PS_TYPE_F32, "Concept should be F32"); \
    item->data.F32 = VALUE; \
}



bool ppStackConvolve(ppStackOptions *options, pmConfig *config)
{
    psAssert(options, "Require options");
    psAssert(config, "Require configuration");

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

    int num = options->num;             // Number of inputs
    options->cells = psArrayAlloc(num); // Cells for convolved images --- a handle for reading again
    options->kernels = psArrayAlloc(num); // PSF-matching kernels --- required in the stacking
    options->regions = psArrayAlloc(num); // PSF-matching regions --- required in the stacking
    int numGood = 0;                    // Number of good frames
    options->numCols = 0;
    options->numRows = 0;
    options->matchChi2 = psVectorAlloc(num, PS_TYPE_F32); // chi^2 for stamps when matching
    psVectorInit(options->matchChi2, NAN);
    options->weightings = psVectorAlloc(num, PS_TYPE_F32); // Combination weightings for images (1/noise^2)
    psVectorInit(options->weightings, 0.0);
    options->covariances = psArrayAlloc(num); // Covariance matrices

    psList *fpaList = psListAlloc(NULL); // List of input FPAs, for concept averaging
    psList *cellList = psListAlloc(NULL); // List of input cells, for concept averaging
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS); // Random number generator

    for (int i = 0; i < num; i++) {
        if (options->inputMask->data.U8[i]) {
            continue;
        }
        psTrace("ppStack", 2, "Convolving input %d of %d to target PSF....\n", i, num);
        pmFPAfileActivate(config->files, false, NULL);
        ppStackFileActivationSingle(config, PPSTACK_FILES_CONVOLVE, true, i);
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.INPUT", i); // File of interest
        pmFPAview *view = ppStackFilesIterateDown(config);
        if (!view) {
            psFree(rng);
            psFree(fpaList);
            psFree(cellList);
            return false;
        }

        pmReadout *readout = pmFPAviewThisReadout(view, file->fpa); // Input readout
        psFree(view);

        if (options->numCols == 0 && options->numRows == 0) {
            options->numCols = readout->image->numCols;
            options->numRows = readout->image->numRows;
        } else if (options->numCols != readout->image->numCols ||
                   options->numRows != readout->image->numRows) {
            psError(PS_ERR_UNKNOWN, true, "Sizes of input images don't match: %dx%d vs %dx%d",
                    readout->image->numCols, readout->image->numRows, options->numCols, options->numRows);
            psFree(rng);
            psFree(fpaList);
            psFree(cellList);
            return false;
        }

        // Background subtraction, scaling and normalisation is performed automatically by the image matching
        psTimerStart("PPSTACK_MATCH");
        if (!ppStackMatch(readout, options, i, config)) {
            psErrorStackPrint(stderr, "Unable to match image %d --- ignoring.", i);
            options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PPSTACK_MASK_MATCH;
            psErrorClear();
            continue;
        }
        options->covariances->data[i] = psMemIncrRefCounter(readout->covariance);

        if (options->stats) {
            psMetadataAddF32(options->stats, PS_LIST_TAIL, "TIME_MATCH", PS_META_DUPLICATE_OK,
                             "Time to match PSF", psTimerMark("PPSTACK_MATCH"));
            psMetadataAddF32(options->stats, PS_LIST_TAIL, "PPSTACK.WEIGHTING", PS_META_DUPLICATE_OK,
                             "Weighting for image", options->weightings->data.F32[i]);

            if (options->convolve) {
                // Pull parameters out of convolution kernel
                pmSubtractionKernels *kernels = psMetadataLookupPtr(NULL, readout->analysis,
                                                                    PM_SUBTRACTION_ANALYSIS_KERNEL);
                psMetadataAddF32(options->stats, PS_LIST_TAIL, "STAMP.MEAN", PS_META_DUPLICATE_OK,
                                 "Mean deviation for stamps", kernels->mean);
                psMetadataAddF32(options->stats, PS_LIST_TAIL, "STAMP.RMS", PS_META_DUPLICATE_OK,
                                 "RMS deviation for stamps", kernels->rms);
                psMetadataAddF32(options->stats, PS_LIST_TAIL, "STAMP.NUM", PS_META_DUPLICATE_OK,
                                 "Number of stamps", kernels->numStamps);
                float deconv = psMetadataLookupF32(NULL, readout->analysis,
                                                   PM_SUBTRACTION_ANALYSIS_DECONV_MAX);
                psMetadataAddF32(options->stats, PS_LIST_TAIL, "KERNEL.DECONV", PS_META_DUPLICATE_OK,
                                 "Deconvolution fraction for kernel", deconv);
            }
        }
        psLogMsg("ppStack", PS_LOG_INFO, "Time to match image %d: %f sec", i, psTimerClear("PPSTACK_MATCH"));

        // Write the temporary convolved files
        pmHDU *hdu = readout->parent->parent->parent->hdu; // HDU for convolved image
        assert(hdu);
        ppStackWriteImage(options->imageNames->data[i], hdu->header, readout->image, config);
        psMetadata *maskHeader = psMetadataCopy(NULL, hdu->header); // Copy of header, for mask
        pmConfigMaskWriteHeader(config, maskHeader);
        ppStackWriteImage(options->maskNames->data[i], maskHeader, readout->mask, config);
        psFree(maskHeader);
        psImageCovarianceTransfer(readout->variance, readout->covariance);
        ppStackWriteImage(options->varianceNames->data[i], hdu->header, readout->variance, config);
#ifdef TESTING
        {
            psString name = NULL;
            psStringAppend(&name, "covariance_%d.fits", i);
            ppStackWriteImage(name, hdu->header, readout->covariance->image, config);
            pmStackVisualPlotTestImage(readout->covariance->image, name);
            psFree(name);
        }
#endif

        pmCell *inCell = readout->parent; // Input cell

        psListAdd(cellList, PS_LIST_TAIL, inCell);
        psListAdd(fpaList, PS_LIST_TAIL, inCell->parent->parent);

        options->cells->data[i] = psMemIncrRefCounter(inCell);
        if (!ppStackFilesIterateUp(config)) {
            psFree(fpaList);
            psFree(cellList);
            psFree(rng);
            return false;
        }
        numGood++;

        ppStackMemDump("match");
    }
    psFree(rng);

    psFree(options->norm); options->norm = NULL;
    psFree(options->sourceLists); options->sourceLists = NULL;
    psFree(options->psf); options->psf = NULL;


    if (numGood == 0) {
        psError(PS_ERR_UNKNOWN, false, "No good images to combine.");
        psFree(fpaList);
        psFree(cellList);
        return false;
    }

    // Update concepts for output
    // XXX This should probably be placed later, so that it's not influenced
    // by images that get rejected later.
    {
        pmFPAview view;                 // View for output
        view.chip = view.cell = view.readout = 0;
        pmCell *outCell = pmFPAfileThisCell(config->files, &view, "PPSTACK.OUTPUT"); // Output cell
        pmFPA *outFPA = outCell->parent->parent; // Output FPA
        pmConceptsAverageFPAs(outFPA, fpaList);
        pmConceptsAverageCells(outCell, cellList, NULL, NULL, true);
        psFree(fpaList);
        psFree(cellList);

        UPDATE_CONCEPT(outFPA,  "FPA.EXPOSURE",  options->sumExposure);
        UPDATE_CONCEPT(outCell, "CELL.EXPOSURE", options->sumExposure);
        UPDATE_CONCEPT(outCell, "CELL.DARKTIME", NAN);
    }


    // Reject images out-of-hand on the basis of their match chi^2
    if (options->convolve) {
        psVector *values = psVectorAllocEmpty(num, PS_TYPE_F32); // Values to sort
        for (int i = 0; i < num; i++) {
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_ALL) {
                continue;
            }
            values->data.F32[values->n++] = options->matchChi2->data.F32[i];
        }
        assert(values->n == numGood);
        if (!psVectorSortInPlace(values)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to sort vector.");
            psFree(values);
            return false;
        }
        float median = numGood % 2 ? values->data.F32[numGood / 2] :
            0.5 * (values->data.F32[numGood / 2 - 1] + values->data.F32[numGood / 2]);

        float rms = 0.74 * (values->data.F32[numGood * 3 / 4] -
                            values->data.F32[numGood / 4]); // Estimated RMS from interquartile range
        psFree(values);

        float rej = psMetadataLookupF32(NULL, recipe, "MATCH.REJ"); // Rejection threshold (stdevs)
        if (isfinite(rej)) {
            float thresh = median + rej * rms; // Threshold for rejection
            psLogMsg("ppStack", PS_LOG_INFO, "chi^2 rejection threshold = %f + %f * %f = %f",
                     median, rej, rms, thresh);

            int numRej = 0;                 // Number rejected
            numGood = 0;                    // Number of good images
            for (int i = 0; i < num; i++) {
              if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_ALL) {
                    continue;
                }
                if (options->matchChi2->data.F32[i] > thresh) {
                    numRej++;
                    options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PPSTACK_MASK_CHI2;
                    psLogMsg("ppStack", PS_LOG_INFO, "Rejecting image %d because of large matching chi^2: %f",
                             i, options->matchChi2->data.F32[i]);
                } else {
                    psLogMsg("ppStack", PS_LOG_INFO, "Image %d has matching chi^2: %f",
                             i, options->matchChi2->data.F32[i]);
                    numGood++;
                }
            }
        }
    }

    if (numGood == 0) {
        psError(PS_ERR_UNKNOWN, false, "No good images to combine.");
        return false;
    }

    return true;
}
