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

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

#include "ppStack.h"

//#define TESTING                  // Write debugging output?


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

    bool mdok;                          // Status of MD lookup
    int iter = psMetadataLookupS32(NULL, recipe, "ITER"); // Rejection iterations
    float combineRej = psMetadataLookupF32(NULL, recipe, "COMBINE.REJ"); // Combination threshold
    bool useVariance = psMetadataLookupBool(&mdok, recipe, "VARIANCE"); // Use variance for rejection?
    bool safe = psMetadataLookupBool(&mdok, recipe, "SAFE"); // Be 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

    psString maskValStr = psMetadataLookupStr(NULL, recipe, "MASK.VAL"); // Name of bits to mask going in
    psMaskType maskVal = pmConfigMaskGet(maskValStr, config); // Bits to mask going in to pmSubtractionMatch
    psString maskBadStr = psMetadataLookupStr(NULL, recipe, "MASK.BAD"); // Name of bits to mask for bad
    psMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

    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 | maskBad, maskBad, kernelSize, iter, combineRej, true,
                        useVariance, safe)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts with rejection.");
        psFree(stack);
        return false;
    }

#ifdef TESTING
    {
        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

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

    bool mdok;                          // Status of MD lookup
    bool useVariance = psMetadataLookupBool(&mdok, recipe, "VARIANCE"); // Use variance for rejection?

    psString maskValStr = psMetadataLookupStr(NULL, recipe, "MASK.VAL"); // Name of bits to mask going in
    psMaskType maskVal = pmConfigMaskGet(maskValStr, config); // Bits to mask going in to pmSubtractionMatch
    psString maskBadStr = psMetadataLookupStr(NULL, recipe, "MASK.BAD"); // Name of bits to mask for bad
    psMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

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

    if (!pmStackCombine(outRO, stack, maskVal | maskBad, maskBad, 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 TESTING
    {
        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;
}
