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

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

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

/// Print a summary of the inputs, and return the number of good inputs
static int stackSummary(const ppStackOptions *options, // Stack options, with input mask
                        const char *place              // Place in code
    )
{
    int numGood = 0;                // Number of good inputs
    psString summary = NULL;        // Summary of images
    for (int i = 0; i < options->num; i++) {
        psString reason = NULL;         // Reason for rejecting
        if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] == 0) {
            psStringAppend(&reason, " Good.");
            numGood++;
        } else {
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_CAL) {
                psStringAppend(&reason, " Calibration failed.");
            }
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_PSF) {
                psStringAppend(&reason, " PSF measurement failed.");
            }
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_MATCH) {
                psStringAppend(&reason, " PSF matching failed.");
            }
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_CHI2) {
                psStringAppend(&reason, " PSF matching chi^2 deviant.");
            }
            if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PPSTACK_MASK_REJECT) {
                psStringAppend(&reason, " Rejection exceeded threshold.");
            }
        }
        psStringAppend(&summary, "Image %d: %s\n", i, reason);
        psFree(reason);
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Summary of images for %s:\n%s", place, summary);
    psFree(summary);

    return numGood;
}



bool ppStackLoop(pmConfig *config, ppStackOptions *options)
{
    assert(config);

    psTimerStart("PPSTACK_TOTAL");
    psTimerStart("PPSTACK_STEPS");

    // Setup
    psTrace("ppStack", 1, "Setup....\n");
    if (!ppStackSetup(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to setup.");
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 0: Setup: %f sec", psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("setup");


    // Preparation for stacking
    psTrace("ppStack", 1, "Preparation for stacking: merging sources, determining target PSF....\n");
    if (!ppStackPrepare(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to prepare for stacking.");
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 1: Load Sources and Generate Target PSF: %f sec",
             psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("prepare");
    if (options->quality) {
        // Can't do anything else
        return true;
    }

    // Convolve inputs
    psTrace("ppStack", 1, "Convolving inputs to target PSF....\n");
    if (!ppStackConvolve(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to convolve images.");
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 2: Generate Convolutions and Save: %f sec",
             psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("convolve");

    if (options->quality) {
        // Can't do anything else
        return true;
    }

    // Ensure sufficient inputs
    {
        int numGood = stackSummary(options, "initial combination");
        psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
        bool safe = psMetadataLookupBool(NULL, recipe, "SAFE"); // Be safe when combining
        if (safe && numGood <= 1) {
            options->quality = PPSTACK_ERR_REJECTED;
            psErrorStackPrint(stderr, "Insufficient inputs for combination with safety on");
            psErrorClear();
            psWarning("Insufficient inputs for combination with safety on");
            return true;
        }
    }

    // Start threading
    ppStackThreadInit();
    ppStackThreadData *stack = ppStackThreadDataSetup(options, config, true);
    if (!stack) {
        psError(psErrorCodeLast(), false, "Unable to initialise stack threads.");
        return false;
    }

    // Prepare for combination
    if (!ppStackCombinePrepare("PPSTACK.OUTPUT", "PPSTACK.OUTPUT.EXP", PPSTACK_FILES_STACK,
                               stack, options, config)) {
        psError(psErrorCodeLast(), false, "Unable to prepare for combination.");
        psFree(stack);
        return false;
    }

    // Initial combination
    psTrace("ppStack", 1, "Initial stack of convolved images....\n");
    if (!ppStackCombineInitial(stack, options, config)) {
        psError(psErrorCodeLast(), false, "Unable to perform initial combination.");
        psFree(stack);
        return false;
    }
    ppStackMemDump("initial");
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 3: Make Initial Stack: %f sec", psTimerClear("PPSTACK_STEPS"));

    // Done with stack inputs for now
    for (int i = 0; i < options->num; i++) {
        pmCellFreeData(options->cells->data[i]);
    }
    psFree(stack);

    // Pixel rejection
    psTrace("ppStack", 1, "Reject pixels....\n");
    if (!ppStackReject(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to reject pixels.");
        psFree(stack);
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 4: Pixel Rejection: %f sec", psTimerClear("PPSTACK_STEPS"));

    // Check inputs
    {
        int numGood = stackSummary(options, "final combination");
        if (numGood <= 0) {
            options->quality = PPSTACK_ERR_REJECTED;
            psErrorStackPrint(stderr, "Insufficient inputs survived rejection stage");
            psErrorClear();
            psWarning("Insufficient inputs survived rejection stage");
            return true;
        }
    }

    stack = ppStackThreadDataSetup(options, config, true);
    if (!stack) {
        psError(psErrorCodeLast(), false, "Unable to initialise stack threads.");
        return false;
    }

    // Final combination
    psTrace("ppStack", 2, "Final stack of convolved images....\n");
    if (!ppStackCombineFinal(stack, options->convCovars, options, config, false, false, true)) {
        psError(psErrorCodeLast(), false, "Unable to perform final combination.");
        psFree(stack);
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 5: Final Stack: %f sec", psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("final");

    // Photometry
    psTrace("ppStack", 1, "Photometering stacked image....\n");
    if (!ppStackPhotometry(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to perform photometry.");
        return false;
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 6: Photometry Analysis: %f sec", psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("photometry");

    // Clean up
    psTrace("ppStack", 2, "Cleaning up after combination....\n");
    if (!ppStackCleanup(stack, options, config)) {
        psError(psErrorCodeLast(), false, "Unable to clean up.");
        psFree(stack);
        return false;
    }
    for (int i = 0; i < options->num; i++) {
        pmCellFreeData(options->cells->data[i]);
    }
    psFree(stack);
    psLogMsg("ppStack", PS_LOG_INFO, "Stage 7: Cleanup, WCS & JPEGS: %f sec", psTimerClear("PPSTACK_STEPS"));
    ppStackMemDump("cleanup");

#if 1
    // Unconvolved stack --- it's cheap to calculate, compared to everything else!
    if (options->convolve) {
        // Start threading
        ppStackThreadData *stack = ppStackThreadDataSetup(options, config, false);
        if (!stack) {
            psError(psErrorCodeLast(), false, "Unable to initialise stack threads.");
            return false;
        }

        // Prepare for combination
        if (!ppStackCombinePrepare("PPSTACK.UNCONV", "PPSTACK.UNCONV.EXP", PPSTACK_FILES_UNCONV,
                                   stack, options, config)) {
            psError(psErrorCodeLast(), false, "Unable to prepare for combination.");
            psFree(stack);
            return false;
        }

        psTrace("ppStack", 2, "Stack of unconvolved images....\n");
        if (!ppStackCombineFinal(stack, options->origCovars, options, config,
                                 false, true, false)) {
            psError(psErrorCodeLast(), false, "Unable to perform unconvolved combination.");
            psFree(stack);
            return false;
        }
        psLogMsg("ppStack", PS_LOG_INFO, "Stage 8: Unconvolved Stack: %f sec", psTimerClear("PPSTACK_STEPS"));
        ppStackMemDump("unconv");

        if (!ppStackFilesIterateUp(config)) {
            psError(psErrorCodeLast(), false, "Unable to close files.");
            psFree(stack);
            return false;
        }
        ppStackFileActivation(config, PPSTACK_FILES_UNCONV, false);
        options->outRO->data_exists = false;
        options->outRO->parent->data_exists = false;
        options->outRO->parent->parent->data_exists = false;
        psFree(options->outRO);
        options->outRO = NULL;
        options->expRO->data_exists = false;
        options->expRO->parent->data_exists = false;
        options->expRO->parent->parent->data_exists = false;
        psFree(options->expRO);
        options->expRO = NULL;

        for (int i = 0; i < options->num; i++) {
            pmCellFreeData(options->cells->data[i]);
        }
        psFree(stack);
    }
    psFree(options->cells); options->cells = NULL;
#endif

    // Finish up
    psTrace("ppStack", 1, "Finishing up....\n");
    if (!ppStackFinish(options, config)) {
        psError(psErrorCodeLast(), false, "Unable to finish up.");
        return false;
    }
    ppStackMemDump("finish");

    return true;
}
