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

#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <pslib.h>

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmSubtractionParams.h"
#include "pmSubtractionKernels.h"
#include "pmSubtractionStamps.h"
#include "pmSubtraction.h"
#include "pmSubtractionMatch.h"


//#define TESTING

// Output memory usage information
static void memCheck(const char *where)
{
#ifdef TESTING
    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));
#endif
    return;
}


static bool getStamps(pmSubtractionStampList **stamps, // Stamps to read
                      const pmReadout *reference, // Reference readout
                      const pmReadout *input, // Input readout, or NULL to generate fake stamps
                      const psImage *subMask, // Mask for subtraction, or NULL
                      psImage *weight,  // Weight map
                      const psRegion *region, // Region of interest, or NULL
                      float threshold,  // Threshold for stamp finding
                      float stampSpacing, // Spacing between stamps
                      float targetWidth,// Target width for fake stamps
                      int size,         // Kernel half-size
                      int footprint     // Convolution footprint for stamps
    )
{
    psTrace("psModules.imcombine", 3, "Finding stamps...\n");
    *stamps = pmSubtractionStampsFind(*stamps, reference->image, subMask, region, threshold, stampSpacing);
    if (!*stamps) {
        psError(PS_ERR_UNKNOWN, false, "Unable to find stamps.");
        return false;
    }

    memCheck("  find stamps");

    if (!input && !pmSubtractionStampsGenerate(*stamps, targetWidth, footprint, size)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to generate target stamps.");
        return false;
    }

    memCheck("   generate stamps");

    psTrace("psModules.imcombine", 3, "Extracting stamps...\n");
    if (!pmSubtractionStampsExtract(*stamps, reference->image, input ? input->image : NULL,
                                    weight, footprint, size)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to extract stamps.");
        return false;
    }

    memCheck("   extract stamps");

    return true;
}





bool pmSubtractionMatch(pmReadout *convolved, const pmReadout *reference, const pmReadout *input,
                        int footprint, float regionSize, float stampSpacing, float threshold,
                        const psArray *sources, const char *stampsName, float targetWidth,
                        pmSubtractionKernelsType type, int size, int spatialOrder,
                        const psVector *isisWidths, const psVector *isisOrders,
                        int inner, int ringsOrder, int binning, bool optimum, const psVector *optFWHMs,
                        int optOrder, float optThreshold, int iter, float rej, psMaskType maskBad,
                        psMaskType maskBlank)
{
    PS_ASSERT_PTR_NON_NULL(convolved, false);
    PS_ASSERT_PTR_NON_NULL(reference, false);
    PS_ASSERT_IMAGE_NON_NULL(reference->image, false);
    PS_ASSERT_IMAGE_TYPE(reference->image, PS_TYPE_F32, false);
    if (reference->mask) {
        PS_ASSERT_IMAGE_NON_NULL(reference->mask, false);
        PS_ASSERT_IMAGE_TYPE(reference->mask, PS_TYPE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(reference->mask, reference->image, false);
    }
    if (reference->weight) {
        PS_ASSERT_IMAGE_NON_NULL(reference->weight, false);
        PS_ASSERT_IMAGE_TYPE(reference->weight, PS_TYPE_F32, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(reference->weight, reference->image, false);
    }
    if (input) {
        PS_ASSERT_IMAGE_NON_NULL(input->image, false);
        PS_ASSERT_IMAGE_TYPE(input->image, PS_TYPE_F32, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(input->image, reference->image, false);
        if (input->mask) {
            PS_ASSERT_IMAGE_NON_NULL(input->mask, false);
            PS_ASSERT_IMAGE_TYPE(input->mask, PS_TYPE_MASK, false);
            PS_ASSERT_IMAGES_SIZE_EQUAL(input->mask, reference->image, false);
        }
        if (input->weight) {
            PS_ASSERT_IMAGE_NON_NULL(input->weight, false);
            PS_ASSERT_IMAGE_TYPE(input->weight, PS_TYPE_F32, false);
            PS_ASSERT_IMAGES_SIZE_EQUAL(input->weight, reference->image, false);
        }
    } else if (!stampsName && !sources) {
        psError(PS_ERR_UNEXPECTED_NULL, true,
                "A list of sources is required for convolving to a target PSF.");
        return false;
    }
    PS_ASSERT_INT_NONNEGATIVE(footprint, false);
    // regionSize can be just about anything (except maybe negative, but it can be NAN)
    PS_ASSERT_FLOAT_LARGER_THAN(stampSpacing, 0.0, false);
    // Don't care what threshold is
    if (sources) {
        PS_ASSERT_ARRAY_NON_NULL(sources, false);
    }
    // stampsName may be anything
    // targetWidth can be just about anything (except maybe negative, but it can be NAN)
    // We'll check kernel type when we allocate the kernels
    PS_ASSERT_INT_POSITIVE(size, false);
    PS_ASSERT_INT_NONNEGATIVE(spatialOrder, false);
    if (isisWidths || isisOrders) {
        PS_ASSERT_VECTOR_NON_NULL(isisWidths, false);
        PS_ASSERT_VECTOR_TYPE(isisWidths, PS_TYPE_F32, false);
        PS_ASSERT_VECTOR_NON_NULL(isisOrders, false);
        PS_ASSERT_VECTOR_TYPE(isisOrders, PS_TYPE_S32, false);
        PS_ASSERT_VECTORS_SIZE_EQUAL(isisWidths, isisOrders, false);
    }
    PS_ASSERT_INT_NONNEGATIVE(inner, false);
    PS_ASSERT_INT_NONNEGATIVE(ringsOrder, false);
    PS_ASSERT_INT_POSITIVE(binning, false);
    if (optimum) {
        PS_ASSERT_VECTOR_NON_NULL(optFWHMs, false);
        PS_ASSERT_INT_NONNEGATIVE(optOrder, false);
        PS_ASSERT_FLOAT_LARGER_THAN(optThreshold, 0.0, false);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(optThreshold, 1.0, false);
    }
    PS_ASSERT_INT_POSITIVE(iter, false);
    PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    // Don't care about maskBad
    // Don't care about maskBlank

    // Reset the output readout, just in case
    if (convolved->image) {
        psFree(convolved->image);
        convolved->image = NULL;
    }
    if (convolved->mask) {
        psFree(convolved->mask);
        convolved->mask = NULL;
    }
    if (convolved->weight) {
        psFree(convolved->weight);
        convolved->weight = NULL;
    }

    psImage *inImage = NULL, *inMask = NULL; // Input image, mask, weight
    if (input) {
        inImage = input->image;
        inMask = input->mask;
    }

    // Where does our weight map come from?
    psImage *weight = NULL;             // Weight image to use
    if (input && input->weight) {
        weight = input->weight;
    } else if (reference->weight) {
        weight = reference->weight;
    } else if (input) {
        weight = input->image;
    } else {
        weight = reference->image;
    }

    // Putting important variable declarations here, since they are freed after a "goto" if there is an error.
    psImage *subMask = NULL;            // Mask for subtraction
    psRegion *region = NULL;            // Iso-kernel region
    psString regionString = NULL;       // String for region
    pmSubtractionStampList *stamps = NULL; // Stamps for matching PSF
    psVector *solution = NULL;          // Solution to match PSF
    pmSubtractionKernels *kernels = NULL; // Kernel basis functions

    int numCols = reference->image->numCols, numRows = reference->image->numRows; // Image dimensions

    memCheck("start");

    subMask = pmSubtractionMask(reference->mask, inMask, maskBad, size, footprint);

    memCheck("mask");

    // Get region of interest
    int xRegions = 1, yRegions = 1;     // Number of iso-kernel regions
    float xRegionSize = 0, yRegionSize = 0; // Size of iso-kernel regions
    if (isfinite(regionSize) && regionSize != 0.0) {
        xRegions = numCols / regionSize + 1;
        yRegions = numRows / regionSize + 1;
        xRegionSize = (float)numCols / (float)xRegions;
        yRegionSize = (float)numRows / (float)yRegions;
        region = psRegionAlloc(NAN, NAN, NAN, NAN);
    }

    // Iterate over iso-kernel regions
    for (int j = 0; j < yRegions; j++) {
        for (int i = 0; i < xRegions; i++) {
            psTrace("psModules.imcombine", 1, "Subtracting region %d of %d...\n",
                    j * xRegions + i + 1, xRegions * yRegions);
            if (region) {
                *region = psRegionSet((int)(i * xRegionSize), (int)((i + 1) * xRegionSize),
                                      (int)(j * yRegionSize), (int)((j + 1) * yRegionSize));
                psFree(regionString);
                regionString = psRegionToString(*region);
                psTrace("psModules.imcombine", 3, "Iso-kernel region: %s out of %d,%d\n",
                        regionString, numCols, numRows);
            }

            if (sources) {
                stamps = pmSubtractionStampsSetFromSources(sources, subMask, region, stampSpacing,
                                                           input ? 0 : 2 * footprint);
            } else if (stampsName && strlen(stampsName) > 0) {
                // Read stamps from file
                psTrace("psModules.imcombine", 3, "Reading stamps from %s...\n", stampsName);
                const char *stampFormat = input ? "%f %f" : "%f %f %f"; // Format for reading stamp file
                psArray *stampsData = stampsData = psVectorsReadFromFile(stampsName, stampFormat);
                psVector *xStamp = NULL, *yStamp = NULL, *fluxStamp = NULL; // Stamp positions and fluxes
                if (!stampsData) {
                    psError(PS_ERR_IO, false, "Unable to read stamps file %s", stampsName);
                    goto ERROR;
                }
                xStamp = stampsData->data[0];
                yStamp = stampsData->data[1];
                if (!input) {
                    fluxStamp = stampsData->data[2];
                }

                // Correct for IRAF/FITS (unit-offset) positions to C (zero-offset) positions
                psBinaryOp(xStamp, xStamp, "-", psScalarAlloc(1.0, PS_TYPE_F32));
                psBinaryOp(yStamp, yStamp, "-", psScalarAlloc(1.0, PS_TYPE_F32));

                stamps = pmSubtractionStampsSet(xStamp, yStamp, fluxStamp, reference->image, subMask,
                                                region, stampSpacing, input ? 0 : 2 * footprint);
                psFree(stampsData);
            }

            // Define kernel basis functions
            if (optimum && (type == PM_SUBTRACTION_KERNEL_ISIS || type == PM_SUBTRACTION_KERNEL_GUNK)) {
                if (!getStamps(&stamps, reference, input, subMask, weight, NULL,
                               threshold, stampSpacing, targetWidth, size, footprint)) {
                    goto ERROR;
                }
                kernels = pmSubtractionKernelsOptimumISIS(type, size, inner, spatialOrder, optFWHMs, optOrder,
                                                          stamps, footprint, optThreshold);
                if (!kernels) {
                    psErrorClear();
                    psWarning("Unable to derive optimum ISIS kernel --- switching to default.");
                }
            }
            if (kernels == NULL) {
                // Not an ISIS/GUNK kernel, or the optimum kernel search failed
                kernels = pmSubtractionKernelsGenerate(type, size, spatialOrder, isisWidths, isisOrders,
                                                       inner, binning, ringsOrder);
            }
            psMetadataAddPtr(convolved->analysis, PS_LIST_TAIL, "SUBTRACTION.KERNEL",
                             PS_DATA_UNKNOWN | PS_META_DUPLICATE_OK, "Subtraction kernels", kernels);

            memCheck("kernels");

            int numRejected = -1;               // Number of rejected stamps in each iteration
            for (int k = 0; k < iter && numRejected != 0; k++) {
                psTrace("psModules.imcombine", 2, "Iteration %d...\n", k);

                if (!getStamps(&stamps, reference, input, subMask, weight, region,
                               threshold, stampSpacing, targetWidth, size, footprint)) {
                    goto ERROR;
                }

                psTrace("psModules.imcombine", 3, "Calculating equation...\n");
                if (!pmSubtractionCalculateEquation(stamps, kernels, footprint)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
                    goto ERROR;
                }

                memCheck("  calculate equation");

                psTrace("psModules.imcombine", 3, "Solving equation...\n");
                solution = pmSubtractionSolveEquation(solution, stamps);
                if (!solution) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to calculate least-squares equation.");
                    goto ERROR;
                }

                memCheck("  solve equation");

                psTrace("psModules.imcombine", 3, "Rejecting stamps...\n");
                numRejected = pmSubtractionRejectStamps(stamps, subMask, solution, footprint, rej, kernels);
                if (numRejected < 0) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to reject stamps.");
                    goto ERROR;
                }
                psLogMsg("psModules.imcombine", PS_LOG_INFO, "%d stamps rejected on iteration %d.",
                         numRejected, k);

                memCheck("  reject stamps");
            }

            if (numRejected > 0) {
                psTrace("psModules.imcombine", 3, "Solving equation...\n");
                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");

            {
                psTrace("psModules.imcombine", 2, "Generating diagnostics...\n");
                // 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);
                    }
                }

                psString comment = NULL; // Comment for metadata
                psStringAppend(&comment, "Subtraction kernel for region %s", regionString);
                psMetadataAddImage(convolved->analysis, PS_LIST_TAIL, "SUBTRACTION.KERNEL.IMAGE",
                                   PS_META_DUPLICATE_OK, comment, convKernels);
                psFree(comment);
                psFree(convKernels);
            }

#ifdef TESTING
            {
                // 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");

            psTrace("psModules.imcombine", 2, "Convolving...\n");
            if (!pmSubtractionConvolve(&convolved->image, &convolved->weight, &convolved->mask,
                                       reference->image, reference->weight, subMask, maskBlank, region,
                                       solution, kernels, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to convolve reference image.");
                goto ERROR;
            }
            psFree(kernels);
            kernels = NULL;

            // Put the solution on the metadata
            {
                psString comment = NULL; // Comment for metadata
                psStringAppend(&comment, "Subtraction solution for region %s", regionString);
                psMetadataAddVector(convolved->analysis, PS_LIST_TAIL, "SUBTRACTION.SOLUTION",
                                    PS_META_DUPLICATE_OK, comment, solution);
                psFree(comment);
                if (region) {
                    psMetadataAddPtr(convolved->analysis, PS_LIST_TAIL, "SUBTRACTION.REGION",
                                     PS_META_DUPLICATE_OK | PS_DATA_REGION, comment, region);
                } else {
                    region = psRegionAlloc(0, numCols, 0, numRows);
                    psMetadataAddPtr(convolved->analysis, PS_LIST_TAIL, "SUBTRACTION.REGION",
                                     PS_META_DUPLICATE_OK | PS_DATA_REGION, comment, region);
                    psFree(region);
                    region = NULL;
                }
            }

            psFree(solution);
            solution = NULL;

            // There is data in the readout now
            convolved->data_exists = true;
            if (convolved->parent) {
                convolved->parent->data_exists = true;
                convolved->parent->parent->data_exists = true;
            }
        }
    }
    psFree(region);
    region = NULL;
    psFree(regionString);
    regionString = NULL;
    psFree(subMask);
    subMask = NULL;

    if (!pmSubtractionBorder(convolved->image, convolved->weight, convolved->mask, size, maskBlank)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to set border of convolved image.");
        goto ERROR;
    }

    memCheck("convolution");


    return true;

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