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

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

#include "ppSub.h"

bool ppSubReadout(pmConfig *config, const pmFPAview *view)
{
    pmReadout *inRO = pmFPAfileThisReadout(config->files, view, "PPSUB.INPUT"); // Input readout
    pmReadout *refRO = pmFPAfileThisReadout(config->files, view, "PPSUB.REF"); // Reference readout
    pmCell *outCell = pmFPAfileThisCell(config->files, view, "PPSUB.OUTPUT"); // Output cell
    pmReadout *outRO = pmReadoutAlloc(outCell); // Output readout

    psImage *input = inRO->image;       // Input image
    psImage *reference = refRO->image;  // Reference image
    PS_ASSERT_IMAGES_SIZE_EQUAL(input, reference, false);

    // Look up appropriate values
    int size = psMetadataLookupS32(NULL, config->arguments, "KERNEL.SIZE"); // Kernel half-size
    int order = psMetadataLookupS32(NULL, config->arguments, "SPATIAL.ORDER"); // Spatial polynomial order
    float spacing = psMetadataLookupF32(NULL, config->arguments, "STAMP.SPACING"); // Typical stamp spacing
    int footprint = psMetadataLookupS32(NULL, config->arguments, "STAMP.FOOTPRINT"); // Stamp half-size
    float threshold = psMetadataLookupF32(NULL, config->arguments, "STAMP.THRESHOLD"); // Threshold for stmps
    int iter = psMetadataLookupS32(NULL, config->arguments, "ITER"); // Rejection iterations
    float rej = psMetadataLookupF32(NULL, config->arguments, "REJ"); // Rejection threshold
    pmSubtractionKernelsType type = psMetadataLookupBool(NULL, config->arguments,
                                                         "KERNEL.TYPE"); // Kernel type
    bool reverse = psMetadataLookupBool(NULL, config->arguments, "REVERSE"); // Reverse sense of subtraction?
    psVector *widths = psMetadataLookupPtr(NULL, config->arguments, "ISIS.WIDTHS"); // ISIS Gaussian widths
    psVector *orders = psMetadataLookupPtr(NULL, config->arguments, "ISIS.ORDERS"); // ISIS Polynomial orders
    psMaskType maskBad = psMetadataLookupU8(NULL, config->arguments, "MASK.BAD"); // Value to mask
    psMaskType maskBlank = psMetadataLookupU8(NULL, config->arguments, "MASK.BLANK"); // Mask for blank reg.

    int numCols = input->numCols, numRows = input->numRows; // Image dimensions
    if (!inRO->mask) {
        inRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
        psImageInit(inRO->mask, 0);
    }
    if (!refRO->mask) {
        refRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
        psImageInit(refRO->mask, 0);
    }

    // Mask for subtraction
    psImage *subMask = pmSubtractionMask(inRO->mask, refRO->mask, maskBad, size, footprint);

#if 0
    if (!inRO->weight) {
        // Need to ensure input image is completely above zero, so we're not weighting by negative numbers
        psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
        if (!psImageBackground(stats, inRO->image, inRO->mask, maskBad, NULL)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to determine statistics of input image.");
            return false;
        }
        float diff = stats->robustMedian - 5.0 * stats->robustStdev;
        if (diff < 0.0) {
            psLogMsg("ppSub", PS_LOG_INFO, "Adding %f to image background to get positive pixels.");
            (void)psBinaryOp(in->image, inRO->image, "+", psScalarAlloc(diff, PS_TYPE_F32));
        }
    }
#endif

    pmSubtractionKernels *kernels = pmSubtractionKernelsGenerate(type, size, order,
                                                                 widths, orders); // Kernel basis functions
    psArray *stamps = NULL;             // Stamps for matching PSF
    psVector *solution = NULL;          // Solution to match PSF

    int numRejected = -1;               // Number of rejected stamps in each iteration
    for (int i = 0; i < iter && numRejected != 0; i++) {
        stamps = pmSubtractionFindStamps(stamps, refRO->image, subMask, threshold, spacing);
        if (!stamps) {
            psError(PS_ERR_UNKNOWN, false, "Unable to find stamps on reference image.");
            goto ERROR;
        }

        if (!pmSubtractionCalculateEquation(stamps, refRO->image, inRO->image, inRO->weight,
                                            kernels, footprint)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
            goto ERROR;
         }

        solution = pmSubtractionSolveEquation(solution, stamps);
        if (!solution) {
            psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
            goto ERROR;
        }

        numRejected = pmSubtractionRejectStamps(stamps, refRO->image, inRO->image, subMask,
                                                solution, footprint, rej, kernels);
        if (numRejected < 0) {
            psError(PS_ERR_UNKNOWN, false, "Unable to reject stamps.");
            goto ERROR;
        }
        psLogMsg("ppSub", PS_LOG_INFO, "%d stamps rejected on iteration %d.", numRejected, i);
    }

    if (numRejected > 0) {
        solution = pmSubtractionSolveEquation(solution, stamps);
        if (!solution) {
            psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
            goto ERROR;
        }
    }

    psImage *convImage = NULL, *convWeight = NULL, *convMask = NULL; // Convolved images
    if (!pmSubtractionConvolve(&convImage, &convWeight, &convMask, refRO->image, refRO->weight, subMask,
                               maskBlank, solution, kernels)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to convolve reference image.");
        goto ERROR;
    }
    psFree(subMask);

    // Do the subtraction
    if (reverse) {
        outRO->image = (psImage*)psBinaryOp(NULL, convImage, "-", inRO->image);
    } else {
        outRO->image = (psImage*)psBinaryOp(NULL, inRO->image, "-", convImage);
    }
    outRO->mask = (psImage*)psBinaryOp(NULL, convMask, "|", inRO->mask);
    if (convWeight) {
        outRO->weight = (psImage*)psBinaryOp(NULL, convWeight, "+", inRO->weight);
    }

    psFree(convImage);
    psFree(convMask);
    psFree(convWeight);

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

    if (!pmFPACopyConcepts(outCell->parent->parent, inRO->parent->parent->parent)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to copy concepts from input to output.");
        return false;
    }

#if 0
    // Generate image with convolution kernels
    int fullSize = 2 * size + 1 + 1;    // Full size of kernel
    psImage *convKernels = psImageAlloc(5 * fullSize - 1, 5 * fullSize - 1, PS_TYPE_F32);
    psImageInit(convKernels, NAN);
    for (int j = -2; j <= 2; j++) {
        for (int i = -2; i <= 2; i++) {
            psImage *kernel = pmSubtractionKernelImage(solution, kernels, (float)i / 2.0,
                                                       (float)j / 2.0); // Image of the kernel
            if (!kernel) {
                psError(PS_ERR_UNKNOWN, false, "Unable to generate kernel image.");
                psFree(convKernels);
                goto ERROR;
            }
            if (psImageOverlaySection(convKernels, kernel, (i + 2) * fullSize,
                                      (j + 2) * fullSize, "=") == 0) {
                psError(PS_ERR_UNKNOWN, false, "Unable to overlay kernel image.");
                psFree(kernel);
                psFree(convKernels);
                goto ERROR;
            }
        }
    }

    // XXX What do we do with this image?

    psFits *kernelFile = psFitsOpen("kernel.fits", "w");
    (void)psFitsWriteImage(kernelFile, NULL, convKernels, 0, NULL);
    psFitsClose(kernelFile);

    psFree(convKernels);
#endif

    psFree(kernels);
    psFree(stamps);
    psFree(solution);

    if (psMetadataLookupBool(NULL, config->arguments, "PHOTOMETRY")) {
        pmFPAfile *photFile = psMetadataLookupPtr(NULL, config->files, "PSPHOT.INPUT");
        pmFPACopy(photFile->fpa, inRO->parent->parent->parent);

        psphotReadout(config, view);

        pmFPAfileActivate(config->files, false, "PSPHOT.INPUT");
    }

    return true;

ERROR:
    psFree(kernels);
    psFree(stamps);
    psFree(solution);
    return false;
}
