#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 ppStackReadoutInitialThread(psThreadJob *job)
{
    PS_ASSERT_THREAD_JOB_NON_NULL(job, false);

    psArray *args = job->args;          // Arguments
    ppStackThread *thread = args->data[0]; // Thread
    pmConfig *config = args->data[1];   // Configuration
    pmReadout *outRO = args->data[2];   // Output readout
    psArray *subRegions = args->data[3]; // Regions for PSF-matching
    psArray *subKernels = args->data[4]; // Kernels for PSF-matching
    psVector *weightings = args->data[5]; // Weightings (1/noise^2) for each image
    psVector *addVariance = args->data[6]; // Additional variance when rejecting

    psArray *inspect = ppStackReadoutInitial(config, outRO, thread->readouts, subRegions, subKernels,
                                             weightings, addVariance);

    job->results = inspect;
    thread->busy = false;

    return inspect ? true : false;
}

bool ppStackReadoutFinalThread(psThreadJob *job)
{
    PS_ASSERT_THREAD_JOB_NON_NULL(job, false);

    psArray *args = job->args;          // Arguments
    ppStackThread *thread = args->data[0]; // Thread
    pmConfig *config = args->data[1];   // Configuration
    pmReadout *outRO = args->data[2];   // Output readout
    psArray *rejected = args->data[3];  // Rejected pixels
    psVector *weightings = args->data[4];  // Weighting (1/noise^2) for each image

    bool status = ppStackReadoutFinal(config, outRO, thread->readouts, rejected,
                                      weightings); // Status of operation

    thread->busy = false;

    return status;
}


bool ppStackInspect(psThreadJob *job)
{
    PS_ASSERT_THREAD_JOB_NON_NULL(job, false);

    psArray *args = job->args;  // Input arguments
    psArray *inspect = args->data[0]; // Array of pixel arrays
    int index = PS_SCALAR_VALUE(args->data[1], S32); // Index of interest

    psArray *inputs = inspect->data[index]; // Array of interest
    psPixels *output = NULL;    // Output pixel list
    for (int i = 0; i < inputs->n; i++) {
        psPixels *input = inputs->data[i]; // Input pixel list
        if (!input || input->n == 0) {
            continue;
        }
        output = psPixelsConcatenate(output, input);
    }

    if (!output) {
        // If there are no pixels to inspect, then just fake it
        output = psPixelsAllocEmpty(0);
    }

    psFree(inputs);
    inspect->data[index] = output;

    return true;
}


psArray *ppStackReadoutInitial(const pmConfig *config, pmReadout *outRO, const psArray *readouts,
                               const psArray *regions, const psArray *kernels, const psVector *weightings,
                               const psVector *addVariance)
{
    assert(config);
    assert(outRO);
    assert(readouts);
    assert(regions);
    assert(kernels);
    assert(readouts->n == regions->n);
    assert(regions->n == kernels->n);
    assert(weightings && weightings->n == readouts->n && weightings->type.type == PS_TYPE_F32);
    assert(addVariance && addVariance->n == readouts->n && addVariance->type.type == PS_TYPE_F32);
    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
    float combineSys = psMetadataLookupF32(NULL, recipe, "COMBINE.SYS"); // Combination systematic error
    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
    psImageMaskType 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
    psImageMaskType 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;
        }

#if 0
        // This doesn't seem to work, so getting the weightings directly from a vector
        float weighting = psMetadataLookupF32(&mdok, ro->analysis, "PPSTACK.WEIGHTING"); // Relative weight
        if (!mdok || !isfinite(weighting)) {
            psWarning("No weighting supplied for image %d --- set to unity.", i);
            weighting = 1.0;
        } else {
            psLogMsg("ppStack", PS_LOG_INFO, "Weighting for image %d is %f", i, weighting);
        }
#endif

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

        stack->data[i] = pmStackDataAlloc(ro, weightings->data.F32[i], addVariance->data.F32[i]);
    }

    if (!pmStackCombine(outRO, stack, maskVal | maskBad, maskBad, kernelSize, iter, combineRej, combineSys,
                        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
    psArray *inspect = psArrayAlloc(num); // 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;
        }
        inspect->data[i] = psMemIncrRefCounter(data->inspect);
    }
    psFree(stack);

    sectionNum++;

    return inspect;
}



bool ppStackReadoutFinal(const pmConfig *config, pmReadout *outRO, const psArray *readouts,
                         const psArray *rejected, const psVector *weightings)
{
    assert(config);
    assert(outRO);
    assert(readouts);
    assert(rejected);
    assert(readouts->n == rejected->n);
    assert(weightings && weightings->n == readouts->n && weightings->type.type == PS_TYPE_F32);

    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
    psImageMaskType 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
    psImageMaskType maskBad = pmConfigMaskGet(maskBadStr, config); // Bits to mask for bad pixels

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

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

#if 0
        // This doesn't seem to work, so getting the weightings directly from a vector
        bool mdok;                      // Status of MD lookup
        float weighting = psMetadataLookupF32(&mdok, ro->analysis, "PPSTACK.WEIGHTING"); // Relative weight
        if (!mdok || !isfinite(weighting)) {
            psWarning("No WEIGHTING supplied for image %d --- set to unity.", i);
            weighting = 1.0;
        }
#endif

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

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

    if (!pmStackCombine(outRO, stack, maskVal | maskBad, maskBad, 0, 0, NAN, NAN,
                        numGood != num, useVariance, false)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to combine input readouts.");
        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

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

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

    psFree(stack);

    sectionNum++;

    return true;
}
