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

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

#include "ppSub.h"

#define TESTING
#define WCS_TOLERANCE 0.001             // Tolerance for WCS

#include <unistd.h>
static void memCheck(const char *where)
{
    return;
    psMemBlock **leaks = NULL;
    int numLeaks = psMemCheckLeaks(0, &leaks, NULL, true);
    size_t largestSize = 0;
    psMemId largest = 0;
    size_t totalSize = 0;
    for (int i = 0; i < numLeaks; i++) {
        psMemBlock *mb = leaks[i];
        totalSize += mb->userMemorySize;
        if (mb->userMemorySize > largestSize) {
            largestSize = mb->userMemorySize;
            largest = mb->id;
        }
    }
    psFree(leaks);
    fprintf(stderr, "%s:\n", where);
    fprintf(stderr, "    Memory in use: %zd\n", totalSize);
    fprintf(stderr, "    Largest block: %ld\n", largest);
    fprintf(stderr, "    sbrk(): %zd\n", (size_t)sbrk(0));
    return;
}

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 = psMetadataLookupS32(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
    int inner = psMetadataLookupS32(NULL, config->arguments, "INNER"); // Inner radius
    int ringsOrder = psMetadataLookupS32(NULL, config->arguments, "RINGS.ORDER"); // RINGS polynomial order
    int binning = psMetadataLookupS32(NULL, config->arguments, "SPAM.BINNING"); // Binning for SPAM kernel
    psMaskType maskBad = pmConfigMask(psMetadataLookupStr(NULL, config->arguments, "MASK.BAD"),
                                    config); // Value to mask
    psMaskType maskBlank = pmConfigMask(psMetadataLookupStr(NULL, config->arguments, "MASK.BLANK"),
                                        config); // Mask for blank reg.

    int numCols = input->numCols, numRows = input->numRows; // Image dimensions
    if (!inRO->mask) {
        if (psMetadataLookupBool(NULL, config->arguments, "MASK.GENERATE")) {
            pmReadoutSetMask(inRO, pmConfigMask("SAT", config), pmConfigMask("BAD", config));
        } else {
            inRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
            psImageInit(inRO->mask, 0);
        }
    }
    if (!refRO->mask) {
        if (psMetadataLookupBool(NULL, config->arguments, "MASK.GENERATE")) {
            pmReadoutSetMask(refRO, pmConfigMask("SAT", config), pmConfigMask("BAD", config));
        } else {
            refRO->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
            psImageInit(refRO->mask, 0);
        }
    }

    memCheck("start");

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

    memCheck("mask");

#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

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

    memCheck("kernels");

    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;
        }

        memCheck("  find stamps");

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

        memCheck("  calculate equation");

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

        memCheck("  solve equation");

        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);

        memCheck("  reject stamps");
    }

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

    memCheck("solution");

#ifdef TESTING
    {
        // 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;
                }
                psFree(kernel);
            }
        }

        // 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);
    }

    {
        // Generate images of the kernel components
        psMetadata *header = psMetadataAlloc(); // Header
        for (int i = 0; i < solution->n; i++) {
            psString name = NULL;       // Header keyword
            psStringAppend(&name, "SOLN%04d", i);
            psMetadataAddF64(header, PS_LIST_TAIL, name, 0, NULL, solution->data.F64[i]);
            psFree(name);
        }
        psArray *kernelImages = pmSubtractionKernelSolutions(solution, kernels, 0.0, 0.0);
        psFits *kernelFile = psFitsOpen("kernels.fits", "w");
        (void)psFitsWriteImageCube(kernelFile, header, kernelImages, NULL);
        psFitsClose(kernelFile);
        psFree(kernelImages);
        psFree(header);
    }
#endif

    memCheck("diag outputs");

    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);
    subMask = NULL;

    psFree(kernels);
    psFree(solution);

    memCheck("convolution");

    // 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);
    }

    memCheck("subtraction");

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

    memCheck("subtraction freed");

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

#ifdef TESTING
    for (int y = 0; y < outRO->image->numRows; y++) {
        for (int x = 0; x < outRO->image->numCols; x++) {
            if (isnan(outRO->image->data.F32[y][x]) && !(outRO->mask->data.U8[y][x] & maskBlank)) {
                printf("%d %d --> %d\n", x, y, outRO->mask->data.U8[y][x]);
            }
        }
    }
#endif

    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 (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");
    }

    // Copy astrometry over
    pmFPA *refFPA = refRO->parent->parent->parent; // Reference FPA
    pmHDU *refHDU = refFPA->hdu;        // Reference HDU
    pmFPA *outFPA = outCell->parent->parent; // Output FPA
    pmHDU *outHDU = outFPA->hdu; // Output HDU
    if (!outHDU || !refHDU) {
        psWarning("Unable to find HDU at FPA level to copy astrometry.");
    } else {
        if (!pmAstromReadWCS(outFPA, outCell->parent, refHDU->header, 1.0)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to read WCS astrometry from reference FPA.");
            psFree(outRO);
            return false;
        }
        if (!outHDU->header) {
            outHDU->header = psMetadataAlloc();
        }
        if (!pmAstromWriteWCS(outHDU->header, outFPA, outCell->parent, WCS_TOLERANCE)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to write WCS astrometry to output FPA.");
            psFree(outRO);
            return false;
        }
    }

    memCheck("header");

    psFree(outRO);

    return true;

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