#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 WCS_TOLERANCE 0.001             // Tolerance for WCS

//#define NO_CONVOLUTION                  // Don't perform convolutions?
#define CONVOLUTION_FILES               // Write convolutions?
#define REJECTION_FILES                 // Write rejection mask?
#define INSPECTION_FILES                // Write inspection mask?

bool ppStackReadout(pmConfig *config, const pmFPAview *view)
{
    // Get the recipe values
    bool mdok;                          // Status of MD lookup
    int iter = psMetadataLookupS32(NULL, config->arguments, "ITER"); // Rejection iterations
    float combineRej = psMetadataLookupF32(NULL, config->arguments, "COMBINE.REJ"); // Combination threshold
    psMaskType maskBad = psMetadataLookupU8(NULL, config->arguments, "MASK.BAD"); // Value to mask
    psMaskType maskBlank = psMetadataLookupU8(NULL, config->arguments, "MASK.BLANK"); // Mask for blank reg.
    float threshold = psMetadataLookupF32(NULL, config->arguments, "THRESHOLD.MASK"); // Threshold for mask deconvolution
    bool photometry = psMetadataLookupBool(&mdok, config->arguments, "PHOTOMETRY"); // Perform photometry?

    // Get the output target
    pmCell *outCell = pmFPAfileThisCell(config->files, view, "PPSTACK.OUTPUT"); // Output cell
    pmReadout *outRO = pmReadoutAlloc(outCell); // Output readout
    pmFPA *outFPA = outCell->parent->parent; // Output FPA

    int num = psMetadataLookupS32(NULL, config->arguments, "INPUTS.NUM"); // Number of inputs
    assert(num > 0);

    // Get the input sources
    psArray *stack = psArrayAllocEmpty(ARRAY_BUFFER); // The stack of inputs
    pmReadout *sources = pmFPAfileThisReadout(config->files, view, "PPSTACK.SOURCES"); // Sources for stamps
    psMetadataIterator *fileIter = psMetadataIteratorAlloc(config->files, PS_LIST_HEAD,
                                                           "^PPSTACK.INPUT$"); // Iterator over input files
    psMetadataItem *fileItem;           // Item from iteration
    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN); // Statistics
    psRandom *rng = psRandomAlloc(0, PS_RANDOM_TAUS); // Random number generator
    int fileNum = 0;                    // Number of file
    float totExposure = 0.0;            // Total exposure time
    psList *fpaList = psListAlloc(NULL); // List of input FPAs, for concept averaging
    psList *cellList = psListAlloc(NULL); // List of input cells, for concept averaging
    while ((fileItem = psMetadataGetAndIncrement(fileIter))) {
        assert(fileItem->type == PS_DATA_UNKNOWN);
        pmFPAfile *inputFile = fileItem->data.V; // An input file
        pmReadout *ro = pmFPAviewThisReadout(view, inputFile->fpa); // Corresponding readout

        psListAdd(fpaList, PS_LIST_TAIL, inputFile->fpa);
        psListAdd(cellList, PS_LIST_TAIL, ro->parent);

        bool mdok;                      // Status of MD lookup
        float weighting = psMetadataLookupF32(&mdok, inputFile->fpa->analysis,
                                              "PPSTACK.WEIGHTING"); // Relative weighting
        if (!mdok || !isfinite(weighting)) {
            psWarning("No WEIGHTING supplied for image %d --- set to unity.", fileNum);
            weighting = 1.0;
        }
        float scale = psMetadataLookupF32(&mdok, inputFile->fpa->analysis, "PPSTACK.SCALE"); // Rel. scale
        if (!mdok || !isfinite(scale)) {
            psWarning("No SCALE supplied for image %d --- set to unity.", fileNum);
            scale = 1.0;
        }

        float exposure = psMetadataLookupF32(NULL, ro->parent->concepts, "CELL.EXPOSURE"); // Exposure time
#if 0
        if (!isfinite(exposure)) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    "CELL.EXPOSURE is not set for input file %ld", stack->n);
            psFree(stats);
            psFree(rng);
            psFree(fileIter);
            psFree(fpaList);
            psFree(cellList);
            psFree(outRO);
            psFree(stack);
            return false;
        }
#endif
        totExposure += exposure;        // Total exposure time

        // Generate convolved version of input
        pmReadout *convolved = pmReadoutAlloc(NULL); // Convolved version of input readout
#ifndef NO_CONVOLUTION
        if (!ppStackMatch(convolved, ro, sources, config)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to match image %d.", fileNum);
            psFree(stats);
            psFree(rng);
            psFree(fileIter);
            psFree(fpaList);
            psFree(cellList);
            psFree(stack);
            psFree(outRO);
            psFree(convolved);
            return false;
        }

#ifdef CONVOLUTION_FILES
        if (convolved->image) {
            psString name = NULL;           // Name of image
            psStringAppend(&name, "convolved%03d_image.fits", fileNum);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, convolved->image, 0, NULL);
            psFitsClose(fits);
        }
        if (convolved->mask) {
            psString name = NULL;           // Name of image
            psStringAppend(&name, "convolved%03d_mask.fits", fileNum);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, convolved->mask, 0, NULL);
            psFitsClose(fits);
        }
        if (convolved->weight) {
            psString name = NULL;           // Name of image
            psStringAppend(&name, "convolved%03d_weight.fits", fileNum);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, convolved->weight, 0, NULL);
            psFitsClose(fits);
        }
#endif // CONVOLUTION_FILES

#else
        convolved = psMemIncrRefCounter(ro);
        sources = NULL;
#endif // NO_CONVOLVUTION


#if 0
        // Background subtraction, scaling and normalisation is performed automatically by the image matching

        // Brain-dead background subtraction
        if (!psImageBackground(stats, ro->image, ro->mask, maskBad, rng)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to get image background on image %ld", stack->n);
            psFree(stats);
            psFree(rng);
            psFree(fileIter);
            psFree(fpaList);
            psFree(cellList);
            psFree(stack);
            psFree(outRO);
            psFree(convolved);
            return false;
        }
        psTrace("ppStack", 3, "Background for image %d is %f\n", fileNum, stats->robustMedian);
        (void)psBinaryOp(ro->image, ro->image, "-", psScalarAlloc(stats->robustMedian, PS_TYPE_F32));

        // Apply scaling
        (void)psBinaryOp(ro->image, ro->image, "*", psScalarAlloc(1.0 / scale, PS_TYPE_F32));

        // Normalise each input by the exposure time
        float exposure = psMetadataLookupF32(NULL, ro->parent->concepts, "CELL.EXPOSURE"); // Exposure time
        if (!isfinite(exposure)) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    "CELL.EXPOSURE is not set for input file %ld", stack->n);
            psFree(stats);
            psFree(rng);
            psFree(fileIter);
            psFree(fpaList);
            psFree(cellList);
            psFree(outRO);
            psFree(convolved);
            psFree(stack);
            return false;
        }

        (void)psBinaryOp(ro->image, ro->image, "/", psScalarAlloc(exposure, PS_TYPE_F32));
        if (ro->weight) {
            (void)psBinaryOp(ro->weight, ro->weight, "/", psScalarAlloc(exposure, PS_TYPE_F32));
        }
#endif

        if (fileNum == 0) {
            // Copy astrometry over
            pmFPA *fpa = ro->parent->parent->parent; // Template FPA
            pmHDU *hdu = fpa->hdu; // Template HDU
            pmHDU *outHDU = outFPA->hdu; // Output HDU
            if (!outHDU || !hdu) {
                psWarning("Unable to find HDU at FPA level to copy astrometry.");
            } else {
                if (!pmAstromReadWCS(outFPA, outCell->parent, hdu->header, 1.0)) {
                    psErrorClear();
                    psWarning("Unable to read WCS astrometry from input FPA.");
                } else {
                    if (!outHDU->header) {
                        outHDU->header = psMetadataAlloc();
                    }
                    if (!pmAstromWriteWCS(outHDU->header, outFPA, outCell->parent, WCS_TOLERANCE)) {
                        psErrorClear();
                        psWarning("Unable to write WCS astrometry to output FPA.");
                    }
                }
            }
        }

        // Don't need the original any more!
        psFree(ro);

        // Ensure there is a mask, or pmStackCombine will complain
        if (!convolved->mask) {
            convolved->mask = psImageAlloc(convolved->image->numCols, convolved->image->numRows,
                                           PS_TYPE_MASK);
            psImageInit(convolved->mask, 0);
        }

        pmStackData *data = pmStackDataAlloc(convolved, weighting); // Data to stack
        psFree(convolved);
        psArrayAdd(stack, ARRAY_BUFFER, data);
        psFree(data);                   // Drop reference

        fileNum++;
    }
    psFree(fileIter);
    psFree(stats);
    psFree(rng);

    if (!pmStackCombine(outRO, stack, maskBad, maskBlank, iter, combineRej)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts with rejection.");
        psFree(fpaList);
        psFree(cellList);
        psFree(stack);
        psFree(outRO);
        return false;
    }

#ifdef INSPECTION_FILES
    {
        psFits *fits = psFitsOpen("combined.fits", "w");
        psFitsWriteImage(fits, NULL, outRO->image, 0, NULL);
        psFitsClose(fits);
    }

    for (int i = 0; i < num; i++) {
        pmStackData *data = stack->data[i]; // Data for this image
        psImage *inspected = psPixelsToMask(NULL, data->pixels,
                                            psRegionSet(0, data->readout->image->numCols,
                                                        0, data->readout->image->numRows),
                                            maskBlank);
        psString name = NULL;           // Name of image
        psStringAppend(&name, "inspect%03d.fits", i);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, inspected, 0, NULL);
        psFree(inspected);
        psFitsClose(fits);
    }
#endif

    for (int i = 0; i < num; i++) {
        pmStackData *data = stack->data[i]; // Data for this image
        pmReadout *readout = data->readout; // Readout for this image

        // Extract the regions and solutions used in the image matching
        psArray *regions = psArrayAllocEmpty(ARRAY_BUFFER); // Array of regions
        {
            psMetadataIterator *iter = psMetadataIteratorAlloc(readout->analysis, PS_LIST_HEAD,
                                                               "^SUBTRACTION.REGION$"); // Iterator
            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 *solutions = psArrayAllocEmpty(ARRAY_BUFFER); // Array of solutions
        {
            psMetadataIterator *iter = psMetadataIteratorAlloc(readout->analysis, PS_LIST_HEAD,
                                                               "^SUBTRACTION.SOLUTION$"); // Iterator
            psMetadataItem *item = NULL;// Item from iteration
            while ((item = psMetadataGetAndIncrement(iter))) {
                assert(item->type == PS_DATA_VECTOR);
                solutions = psArrayAdd(solutions, ARRAY_BUFFER, item->data.V);
            }
            psFree(iter);
        }
        assert(regions->n == solutions->n);

        pmSubtractionKernels *kernels = psMetadataLookupPtr(NULL, readout->analysis,
                                                            "SUBTRACTION.KERNEL"); // Kernels

        psPixels *reject = pmStackReject(data->pixels, threshold, regions,
                                         solutions, kernels); // List of pixels to reject
        psFree(data->pixels);
        data->pixels = reject;
    }

#ifdef REJECTION_FILES
    for (int i = 0; i < num; i++) {
        pmStackData *data = stack->data[i]; // Data for this image
        psImage *rejected = psPixelsToMask(NULL, data->pixels,
                                           psRegionSet(0, data->readout->image->numCols,
                                                       0, data->readout->image->numRows),
                                           maskBlank);
        psString name = NULL;           // Name of image
        psStringAppend(&name, "reject%03d.fits", i);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, rejected, 0, NULL);
        psFree(rejected);
        psFitsClose(fits);
    }
#endif

    if (!pmStackCombine(outRO, stack, maskBad, maskBlank, 0, combineRej)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts.");
        psFree(fpaList);
        psFree(cellList);
        psFree(stack);
        psFree(outRO);
        return false;
    }

#if 0
    // Restore image to counts using the total exposure time
    (void)psBinaryOp(outRO->image, outRO->image, "*", psScalarAlloc(totExposure, PS_TYPE_F32));
    if (outRO->weight) {
        (void)psBinaryOp(outRO->weight, outRO->weight, "*", psScalarAlloc(totExposure, PS_TYPE_F32));
    }
#endif
    psMetadataAddF32(outCell->concepts, PS_LIST_TAIL, "CELL.EXPOSURE", PS_META_REPLACE,
                     "Summed exposure time (sec)", totExposure);
    psMetadataAddF32(outCell->parent->parent->concepts, PS_LIST_TAIL, "FPA.EXPOSURE", PS_META_REPLACE,
                     "Summed exposure time (sec)", totExposure);

    outRO->data_exists = true;
    outCell->data_exists = true;
    outCell->parent->data_exists = true;

    // Copy other concepts
    pmConceptsAverageFPAs(outFPA, fpaList);
    pmConceptsAverageCells(outCell, cellList, NULL, NULL, true);
    psFree(fpaList);
    psFree(cellList);

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

        if (!psphotReadout(config, view)) {
            psError(psErrorCodeLast(), false, "Unable to perform photometry on stacked image.\n");
            return false;
        }

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

    psFree(stack);
    psFree(outRO);                      // Drop reference

    return true;
}
