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

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

#include "ppStack.h"
#include "ppStackLoop.h"

#define WCS_TOLERANCE 0.001             // Tolerance for WCS

//#define TESTING                         // Enable test output

bool ppStackCombineFinal(ppStackThreadData *stack, psArray *covariances, ppStackOptions *options,
                         pmConfig *config, bool safe, bool normalise, bool grow)
{
    psAssert(stack, "Require stack");
    psAssert(options, "Require options");
    psAssert(config, "Require configuration");

    pmReadout *outRO = options->outRO;                                      // Output readout
    pmReadout *expRO = options->expRO;                                      // Exposure readout
    int numCols = outRO->image->numCols, numRows = outRO->image->numRows; // Size of image

    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
    psAssert(recipe, "We've thrown an error on this before.");
    float poorFrac = psMetadataLookupF32(NULL, recipe, "POOR.FRACTION"); // Fraction for "poor"

    // Grow the list of rejected pixels, if desired
    psArray *reject = psArrayAlloc(options->num); // Pixels rejected for each image
    for (int i = 0; i < options->num; i++) {
        if (options->inputMask->data.U8[i]) {
            continue;
        }
        if (grow) {
            reject->data[i] = pmStackRejectGrow(options->rejected->data[i], numCols, numRows, poorFrac,
                                                options->regions->data[i], options->kernels->data[i]);
            if (!reject->data[i]) {
                psError(psErrorCodeLast(), false, "Unable to grow rejected pixels for image %d", i);
                psFree(reject);
                return false;
            }
        } else {
            reject->data[i] = psMemIncrRefCounter(options->rejected->data[i]);
        }
    }

    if (!outRO->mask) {
        outRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_IMAGE_MASK);
    }
    if (!expRO->mask) {
        expRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_IMAGE_MASK);
    }

    stack->lastScan = 0;            // Reset read
    bool status;                    // Status of read
    for (int numChunk = 0; true; numChunk++) {
        ppStackThread *thread = ppStackThreadRead(&status, stack, config, numChunk, 0);
        if (!status) {
            // Something went wrong
            psError(psErrorCodeLast(), false, "Unable to read chunk %d", numChunk);
            psFree(reject);
            return false;
        }
        if (!thread) {
            // Nothing more to read
            break;
        }

        // call: ppStackReadoutFinal(config, outRO, readouts, rejected)
        psThreadJob *job = psThreadJobAlloc("PPSTACK_FINAL_COMBINE"); // Job to start
        psArrayAdd(job->args, 1, thread);
        psArrayAdd(job->args, 1, reject);
        psArrayAdd(job->args, 1, options);
        psArrayAdd(job->args, 1, config);
        PS_ARRAY_ADD_SCALAR(job->args, safe, PS_TYPE_U8);
        PS_ARRAY_ADD_SCALAR(job->args, normalise, PS_TYPE_U8);
        if (!psThreadJobAddPending(job)) {
            psFree(job);
            psFree(reject);
            return false;
        }
        psFree(job);
    }

    if (!psThreadPoolWait(true)) {
        psError(psErrorCodeLast(), false, "Unable to do final combination.");
        psFree(reject);
        return false;
    }

    psFree(reject);

    // Sum covariance matrices
    if (covariances) {
        outRO->covariance = psImageCovarianceAverageWeighted(covariances, options->weightings);
    } else {
        outRO->covariance = psImageCovarianceNone();
    }

    // Propagate WCS
    bool wcsDone = false;           // Have we done the WCS?
    for (int i = 0; i < options->num && !wcsDone; i++) {
        if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            continue;
        }

        ppStackThread *thread = stack->threads->data[0]; // Representative stack
        pmReadout *inRO = thread->readouts->data[i]; // Template readout
        if (inRO && !wcsDone) {
            // Copy astrometry over
            wcsDone = true;
            pmHDU *inHDU = pmHDUFromCell(inRO->parent); // Template HDU
            pmHDU *outHDU = pmHDUFromCell(outRO->parent); // Output HDU
            pmChip *outChip = outRO->parent->parent; // Output chip
            pmFPA *outFPA = outChip->parent; // Output FPA
            if (!outHDU || !inHDU) {
                psWarning("Unable to find HDU at FPA level to copy astrometry.");
            } else {
                if (!pmAstromReadWCS(outFPA, outChip, inHDU->header, 1.0)) {
                    psErrorClear();
                    psWarning("Unable to read WCS astrometry from input FPA.");
                    wcsDone = false;
                } else {
                    if (!outHDU->header) {
                        outHDU->header = psMetadataAlloc();
                    }
                    if (!pmAstromWriteWCS(outHDU->header, outFPA, outChip, WCS_TOLERANCE)) {
                        psErrorClear();
                        psWarning("Unable to write WCS astrometry to output FPA.");
                        wcsDone = false;
                    }
                }
            }
        }
    }

    // Set exposure time correctly
    {
        float exptime = 0.0;            // Summed exposure time
        for (int i = 0; i < options->num; i++) {
            if (options->inputMask->data.U8[i]) {
                continue;
            }
            exptime += options->exposures->data.F32[i];
        }

        {
            psMetadataItem *item = psMetadataLookup(outRO->parent->concepts, "CELL.EXPOSURE");
            item->data.F32 = exptime;
        }
        {
            psMetadataItem *item = psMetadataLookup(outRO->parent->parent->parent->concepts, "FPA.EXPOSURE");
            item->data.F32 = exptime;
        }
    }

    // Put version information into the header
    pmHDU *hdu = pmHDUFromCell(outRO->parent);
    if (!hdu) {
        psError(PPSTACK_ERR_PROG, false, "Unable to find HDU for output.");
        return false;
    }
    if (!hdu->header) {
        hdu->header = psMetadataAlloc();
    }
    ppStackVersionHeader(hdu->header);


#ifdef TESTING
    static int pass = 0;                // Pass through
    psString name = NULL;               // Name of file
    psStringAppend(&name, "combined_image_final_%d.fits", pass);
    pass++;
    ppStackWriteImage(name, NULL, outRO->image, config);
    psStringSubstitute(&name, "mask", "image");
    ppStackWriteImage(name, NULL, outRO->mask, config);
    psStringSubstitute(&name, "variance", "mask");
    ppStackWriteImage(name, NULL, outRO->variance, config);
    psFree(name);

    pmStackVisualPlotTestImage(outRO->image, "combined_image_final.fits");
#endif

    return true;
}
