#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

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
    float convolveRej = psMetadataLookupF32(NULL, config->arguments, "CONVOLVE.REJ"); // Convolution threshold
    float extent = psMetadataLookupF32(NULL, config->arguments, "EXTENT"); // Extent of convolution
    psMaskType maskBad = psMetadataLookupU8(NULL, config->arguments, "MASK.BAD"); // Value to mask
    psMaskType maskBlank = psMetadataLookupU8(NULL, config->arguments, "MASK.BLANK"); // Mask for blank reg.
    psVector *seeing = psMetadataLookupPtr(&mdok, config->arguments, "SEEING"); // Seeing in each image

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

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

        float seeing = psMetadataLookupF32(NULL, inputFile->fpa->analysis, "PPSTACK.SEEING"); // Seeing FWHM
        float weight = psMetadataLookupF32(NULL, inputFile->fpa->analysis, "PPSTACK.WEIGHT"); // Rel. weight
        float scale = psMetadataLookupF32(NULL, inputFile->fpa->analysis, "PPSTACK.SCALE"); // Rel. scale

        // 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(stack);
            psFree(outRO);
            return false;
        }
        (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(outRO);
            psFree(stack);
            return false;
        }
        totExposure += exposure;        // Total exposure time
        (void)psBinaryOp(ro->image, ro->image, "/", psScalarAlloc(exposure, PS_TYPE_F32));
        pmStackData *data = pmStackDataAlloc(ro, seeing, weight); // Data to stack
        psArrayAdd(stack, ARRAY_BUFFER, data);
        psFree(data);                   // Drop reference

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

        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(stack);
        psFree(outRO);
        return false;
    }

#if 1
    {
        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->sky->image->numCols,
                                                        0, data->sky->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

    // Only perform the additional rejection if we have seeing information
    if (seeing && !pmStackReject(stack, maskBad, extent, convolveRej)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to reject input pixels.");
        psFree(stack);
        psFree(outRO);
        return false;
    }

#if 1
    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->sky->image->numCols,
                                                       0, data->sky->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(stack);
        psFree(outRO);
        return false;
    }

    // Restore image to counts using the total exposure time
    (void)psBinaryOp(outRO->image, outRO->image, "*", psScalarAlloc(totExposure, PS_TYPE_F32));
    psMetadataAddF32(outRO->parent->concepts, PS_LIST_TAIL, "CELL.EXPOSURE", PS_META_REPLACE,
                     "Summed exposure time (sec)", totExposure);

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

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

    return true;
}
