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

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

#include "ppStack.h"

//#define REJECTION_FILES                 // Write rejection mask?
//#define INSPECTION_FILES                // Write inspection mask?
//#define COMBINED_FILES                  // Write combined images?


bool ppStackReadoutInitial(const pmConfig *config, pmReadout *outRO, const psArray *readouts,
                           const psArray *regions, const psArray *kernels)
{
    assert(config);
    assert(outRO);
    assert(readouts);
    assert(regions);
    assert(kernels);
    assert(readouts->n == regions->n);
    assert(regions->n == kernels->n);
    static int sectionNum = 0;          // Section number; for debugging outputs


    // 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 maskVal = pmConfigMaskGet("BLANK", config); // Value to mask
    bool useVariance = psMetadataLookupBool(&mdok, config->arguments, "VARIANCE"); // Use variance for rejection?
    bool safe = psMetadataLookupBool(&mdok, config->arguments, "SAFE"); // Play safe when combining small numbers of pixels?
    psMetadata *ppsub = psMetadataLookupMetadata(NULL, config->recipes, "PPSUB"); // PPSUB recipe
    int kernelSize = psMetadataLookupS32(NULL, ppsub, "KERNEL.SIZE"); // Kernel half-size


    int num = readouts->n;              // Number of inputs
    psArray *stack = psArrayAlloc(num); // Array for stacking

    for (int i = 0; i < num; i++) {
        pmReadout *ro = readouts->data[i];
        if (!ro) {
            // Bad image
            continue;
        }
        pmFPA *fpa = ro->parent->parent->parent; // Parent FPA

        float weighting = psMetadataLookupF32(&mdok, fpa->analysis, "PPSTACK.WEIGHTING"); // Relative weight
        if (!mdok || !isfinite(weighting)) {
            psWarning("No WEIGHTING supplied for image %d --- set to unity.", i);
            weighting = 1.0;
        }

        // 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);
        }

        stack->data[i] = pmStackDataAlloc(ro, weighting);
    }

    if (!pmStackCombine(outRO, stack, maskVal, maskVal, kernelSize, iter, combineRej, true,
                        useVariance, safe)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts with rejection.");
        psFree(stack);
        return false;
    }

#ifdef COMBINED_FILES
    {
        psString name = NULL;           // Name of image
        psStringAppend(&name, "combined_initial_%03d.fits", sectionNum);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, outRO->image, 0, NULL);
        psFitsClose(fits);
    }
#endif
#ifdef INSPECTION_FILES
    for (int i = 0; i < stack->n; i++) {
        pmStackData *data = stack->data[i]; // Data for this image
        if (!data) {
            continue;
        }
        psImage *inspected = psPixelsToMask(NULL, data->inspect,
                                            psRegionSet(0, outRO->image->numCols - 1,
                                                        0, outRO->image->numRows - 1),
                                            maskVal);
        psString name = NULL;           // Name of image
        psStringAppend(&name, "inspect_%03d_%03d.fits", sectionNum, i);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, inspected, 0, NULL);
        psFree(inspected);
        psFitsClose(fits);
    }
#endif

    // Save list of pixels to inspect
    for (int i = 0; i < num; i++) {
        pmStackData *data = stack->data[i]; // Data for this image
        if (!data) {
            continue;
        }
        pmReadout *readout = data->readout; // Readout of interest
        if (!readout) {
            continue;
        }
        psMetadataAddPtr(readout->analysis, PS_LIST_TAIL, PPSTACK_INSPECT_PIXELS,
                         PS_DATA_PIXELS | PS_META_DUPLICATE_OK, "Pixels to inspect from initial combination",
                         data->inspect);
    }
    psFree(stack);

    sectionNum++;

    return true;
}



bool ppStackReadoutFinal(const pmConfig *config, pmReadout *outRO, const psArray *readouts,
                         const psArray *rejected)
{
    assert(config);
    assert(outRO);
    assert(readouts);
    assert(rejected);
    assert(readouts->n == rejected->n);

    static int sectionNum = 0;

    // Get the recipe values
    bool mdok;                          // Status of MD lookup
    psMaskType maskVal = pmConfigMaskGet("BLANK", config); // Value to mask
    bool useVariance = psMetadataLookupBool(&mdok, config->arguments, "VARIANCE"); // Use variance for rejection?

    int num = readouts->n;              // Number of inputs
    psArray *stack = psArrayAlloc(num); // Array for stacking

    pmCell *outCell = outRO->parent;    // Output cell
    pmChip *outChip = outCell->parent;  // Output chip
    pmFPA *outFPA = outChip->parent;    // Output FPA

    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
    int numGood = num;                  // Number of good inputs: images that haven't been completely rejected
    for (int i = 0; i < num; i++) {
        if (!rejected->data[i]) {
            // Image completely rejected
            numGood--;
            continue;
        }

        pmReadout *ro = readouts->data[i];
        assert(ro);
        pmFPA *fpa = ro->parent->parent->parent; // Parent FPA

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

        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(fpaList);
            psFree(cellList);
            psFree(stack);
            return false;
        }
        totExposure += exposure;        // Total exposure time

        // 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);
        }

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

        pmStackData *data = pmStackDataAlloc(ro, weighting);
        data->reject = psMemIncrRefCounter(rejected->data[i]);
        stack->data[i] = data;
    }

#ifdef REJECTION_FILES
    if (sectionNum == 0) {
        for (int i = 0; i < stack->n; i++) {
            pmStackData *data = stack->data[i]; // Data for this image
            if (!data) {
                continue;
            }
            psImage *reject = psPixelsToMask(NULL, data->reject,
                                             psRegionSet(0, outRO->image->numCols - 1,
                                                         0, outRO->image->numRows - 1),
                                             maskVal);
            psString name = NULL;           // Name of image
            psStringAppend(&name, "reject_%03d.fits", i);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, reject, 0, NULL);
            psFree(reject);
            psFitsClose(fits);
        }
    }
#endif

    if (!pmStackCombine(outRO, stack, maskVal, maskVal, 0, 0, NAN, numGood != num, useVariance, false)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts.");
        psFree(fpaList);
        psFree(cellList);
        psFree(stack);
        return false;
    }

#ifdef COMBINED_FILES
    {
        psString name = NULL;           // Name of image
        psStringAppend(&name, "combined_final_%03d.fits", sectionNum);
        psFits *fits = psFitsOpen(name, "w");
        psFree(name);
        psFitsWriteImage(fits, NULL, outRO->image, 0, NULL);
        psFitsClose(fits);
    }
#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);

    psFree(stack);

    sectionNum++;

    return true;
}
