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

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

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

bool ppStackReject(ppStackOptions *options, pmConfig *config)
{
    psAssert(options, "Require options");
    psAssert(config, "Require configuration");

    int num = options->num;             // Number of inputs
    options->rejected = psArrayAlloc(num);

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

    psMetadata *ppsub = psMetadataLookupMetadata(NULL, config->recipes, "PPSUB"); // PPSUB recipe
    psAssert(ppsub, "We've thrown an error on this before.");

    float threshold = psMetadataLookupF32(NULL, recipe, "THRESHOLD.MASK"); // Threshold for mask deconvolution
    float poorFrac = psMetadataLookupF32(NULL, recipe, "POOR.FRACTION"); // Fraction for "poor"
    float imageRej = psMetadataLookupF32(NULL, recipe, "IMAGE.REJ"); // Maximum fraction of image to reject
                                                                     // before rejecting entire image
    int stride = psMetadataLookupS32(NULL, ppsub, "STRIDE"); // Size of convolution patches

    // Count images rejected out of hand
    int numRejected = 0;        // Number of inputs rejected completely
    for (int i = 0; i < num; i++) {
        if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            numRejected++;
        }
    }

    // Concatenate inspection lists
    for (int i = 0; i < num; i++) {
        if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            continue;
        }

        psThreadJob *job = psThreadJobAlloc("PPSTACK_INSPECT"); // Job to start
        psArrayAdd(job->args, 1, options->inspect);
        PS_ARRAY_ADD_SCALAR(job->args, i, PS_TYPE_S32);
        if (!psThreadJobAddPending(job)) {
            psFree(job);
            return false;
        }
        psFree(job);
    }
    if (!psThreadPoolWait(true)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to concatenate inspection lists.");
        return false;
    }


    // Reject bad pixels
    if (psMetadataLookupS32(NULL, config->arguments, "-threads") > 0) {
        pmStackRejectThreadsInit();
    }
    for (int i = 0; i < num; i++) {
        if (options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            continue;
        }
        psTimerStart("PPSTACK_REJECT");

#ifdef TESTING
        {
            psImage *mask = psPixelsToMask(NULL, options->inspect->data[i],
                                           psRegionSet(0, options->numCols - 1, 0, options->numRows - 1),
                                           0xff); // Mask image
            psString name = NULL;           // Name of image
            psStringAppend(&name, "inspect_%03d.fits", i);
            pmStackVisualPlotTestImage(mask, name);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, mask, 0, NULL);
            psFree(mask);
            psFitsClose(fits);
        }
#endif

        psPixels *reject = pmStackReject(options->inspect->data[i], options->numCols, options->numRows,
                                         threshold, poorFrac, stride, options->subRegions->data[i],
                                         options->subKernels->data[i]); // Rejected pixels

#ifdef TESTING
        {
            psImage *mask = psPixelsToMask(NULL, reject,
                                           psRegionSet(0, options->numCols - 1, 0, options->numRows - 1),
                                           0xff); // Mask image
            psString name = NULL;           // Name of image
            psStringAppend(&name, "reject_%03d.fits", i);
            pmStackVisualPlotTestImage(mask, name);
            psFits *fits = psFitsOpen(name, "w");
            psFree(name);
            psFitsWriteImage(fits, NULL, mask, 0, NULL);
            psFree(mask);
            psFitsClose(fits);
        }
#endif

        psFree(options->inspect->data[i]);
        options->inspect->data[i] = NULL;

        if (!reject) {
            psWarning("Rejection on image %d didn't work --- reject entire image.", i);
            numRejected++;
            options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] = PPSTACK_MASK_REJECT;
        } else {
            float frac = reject->n / (float)(options->numCols * options->numRows); // Pixel fraction
            psLogMsg("ppStack", PS_LOG_INFO, "%ld pixels rejected from image %d (%.1f%%)",
                     reject->n, i, frac * 100.0);
            if (frac > imageRej) {
                psWarning("Image %d rejected completely because rejection fraction (%.3f) "
                          "exceeds limit (%.3f)", i, frac, imageRej);
                psFree(reject);
                // reject == NULL means reject image completely
                reject = NULL;
                options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PPSTACK_MASK_BAD;
                numRejected++;
            }
        }

        // Images without a list of rejected pixels (the list may be empty) are rejected completely
        options->rejected->data[i] = reject;

        if (options->stats) {
            psMetadataAddF32(options->stats, PS_LIST_TAIL, "TIME_REJECT", PS_META_DUPLICATE_OK,
                             "Time to perform rejection", psTimerMark("PPSTACK_REJECT"));
            psMetadataAddS32(options->stats, PS_LIST_TAIL, "REJECT_PIXELS", PS_META_DUPLICATE_OK,
                             "Number of pixels rejected", reject ? reject->n : 0);
        }
        psLogMsg("ppStack", PS_LOG_INFO, "Time to perform rejection on image %d: %f sec", i,
                 psTimerClear("PPSTACK_REJECT"));
    }

    psFree(options->inspect); options->inspect = NULL;
    psFree(options->subKernels); options->subKernels = NULL;
    psFree(options->subRegions); options->subRegions = NULL;

    if (numRejected >= num - 1) {
        psError(PS_ERR_UNKNOWN, true, "All inputs completely rejected; unable to proceed.");
        return false;
    }

    if (options->stats) {
        psMetadataAddS32(options->stats, PS_LIST_TAIL, "REJECT_IMAGES", 0,
                         "Number of images rejected completely", numRejected);
    }

    return true;
}
