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

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

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

bool ppStackCombineFinal(ppStackThreadData *stack, ppStackOptions *options, pmConfig *config)
{
    psAssert(stack, "Require stack");
    psAssert(options, "Require options");
    psAssert(config, "Require configuration");

    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(PS_ERR_UNKNOWN, false, "Unable to read chunk %d", numChunk);
            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, options);
        psArrayAdd(job->args, 1, config);
        if (!psThreadJobAddPending(job)) {
            psFree(job);
            return false;
        }
        psFree(job);
    }

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

    // Sum covariance matrices
    double sumWeights = 0.0;            // Sum of weights
    for (int i = 0; i < options->num; i++) {
        if (options->inputMask->data.U8[i]) {
            psFree(options->covariances->data[i]);
            options->covariances->data[i] = NULL;
            continue;
        }
        psKernel *covar = options->covariances->data[i]; // Covariance matrix
        if (!covar) {
            continue;
        }
        float weight = options->weightings->data.F32[i]; // Weight to apply
        psBinaryOp(covar->image, covar->image, "*", psScalarAlloc(weight, PS_TYPE_F32));
        sumWeights += weight;
    }
    if (sumWeights > 0.0) {
        pmReadout *outRO = options->outRO;  // Output readout
        outRO->covariance = psImageCovarianceSum(options->covariances);
        psBinaryOp(outRO->covariance->image, outRO->covariance->image, "/",
                   psScalarAlloc(sumWeights, PS_TYPE_F32));
        psFree(options->covariances); options->covariances = NULL;
        psImageCovarianceTransfer(outRO->variance, outRO->covariance);
    }

#ifdef TESTING
    pmStackVisualPlotTestImage(outRO->image, "combined_initial.fits");
    ppStackWriteImage("combined_final.fits", NULL, outRO->image, config);
#endif

    return true;
}
