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

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

#include "pmSubtraction.h"
#include "pmSubtractionKernels.h"

#define PIXEL_LIST_BUFFER 100           // Number of pixels to add to list at a time

//#define TESTING                         // Testing output


psPixels *pmStackReject(const psPixels *in, const psRegion *valid, float threshold,
                        const psArray *subRegions, const psArray *kernels)
{
    PS_ASSERT_PIXELS_NON_NULL(in, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN_OR_EQUAL(threshold, 0.0, NULL);
    PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(threshold, 1.0, NULL);
    PS_ASSERT_ARRAY_NON_NULL(subRegions, NULL);
    PS_ASSERT_ARRAY_NON_NULL(kernels, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(subRegions, kernels, NULL);

    // Trivial case
    if (in->n == 0) {
        return psPixelsAllocEmpty(0);
    }

    // Get the original image size
    int numRegions = subRegions->n;        // Number of regions
    int numCols = 0, numRows = 0;       // Size of original image
    int minCols = INT_MAX, minRows = INT_MAX; // Minimum coordinate for image
    int size = 0;                       // Size of kernel
    for (int i = 0; i < numRegions; i++) {
        psRegion *subRegion = subRegions->data[i]; // Region of interest
        if (subRegion->x0 < minCols) {
            minCols = subRegion->x0;
        }
        if (subRegion->y0 < minRows) {
            minRows = subRegion->y0;
        }
        if (subRegion->x1 > numCols) {
            numCols = subRegion->x1;
        }
        if (subRegion->y1 > numRows) {
            numRows = subRegion->y1;
        }

        pmSubtractionKernels *kernel = kernels->data[i]; // Kernel of interest
        if (size == 0) {
            size = kernel->size;
        } else if (kernel->size != size) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Kernel sizes are not identical: %d vs %d",
                    size, kernel->size);
            return NULL;
        }
    }

    // Adjust the size for the size of the subimage
    if (valid) {
        minCols = PS_MAX(valid->x0, minCols);
        minRows = PS_MAX(valid->y0, minRows);
        numCols = PS_MIN(valid->x1, numCols);
        numRows = PS_MIN(valid->y1, numRows);
    }
    psTrace("psModules.imcombine", 1, "Rejecting [%d:%d,%d:%d]\n", minCols, numCols, minRows, numRows);

    psImage *mask = psPixelsToMask(NULL, in, psRegionSet(minCols, numCols - 1, minRows, numRows - 1),
                                   0x01); // Mask
    psImage *image = psImageCopy(NULL, mask, PS_TYPE_F32); // Floating-point version, so we can convolve
    psFree(mask);

    // Convolve the image with the kernel --- we're basically applying a matched filter and then thresholding
    pmReadout *convRO = pmReadoutAlloc(NULL); // Readout with convolved image
    pmReadout *inRO = pmReadoutAlloc(NULL); // Readout with input image
    inRO->image = image;
    inRO->col0 = minCols;
    inRO->row0 = minRows;
    for (int i = 0; i < numRegions; i++) {
        psRegion *region = subRegions->data[i]; // Region of interest
        if (valid && (region->x0 > valid->x1 || region->x1 < valid->x0 ||
                      region->y0 > valid->y1 || region->y1 < valid->y0)) {
            // Region is outside of our sub-image
            continue;
        }
        pmSubtractionKernels *kernel = kernels->data[i]; // Kernel of interest
        if (!pmSubtractionConvolve(convRO, NULL, inRO, NULL, NULL, 0, region, kernel, false, true)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to convolve mask image in region %d.", i);
            psFree(convRO);
            psFree(inRO);
            return NULL;
        }

        // Need to adjust the thresholding level for the normalisation of the kernel --- the application of
        // the kernel may scale the unit level that we've inserted.

        // Image of the kernel at the centre of the region
        float xNorm = (region->x0 + 0.5 * (region->x1 - region->x0) - kernel->numCols/2.0) /
            (float)kernel->numCols;
        float yNorm = (region->y0 + 0.5 * (region->y1 - region->y0) - kernel->numRows/2.0) /
            (float)kernel->numRows;
        psImage *image = pmSubtractionKernelImage(kernel, xNorm, yNorm, false);
        if (!image) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate kernel image.");
            psFree(convRO);
            psFree(inRO);
            return NULL;
        }
        float sum = 0.0;
        for (int y = 0; y < image->numRows; y++) {
            for (int x = 0; x < image->numCols; x++) {
                sum += image->data.F32[y][x];
            }
        }
        psFree(image);

        // Range for normalisation
        int yMin = PS_MAX(minRows, region->y0) - inRO->row0;
        int yMax = PS_MIN(numRows - 1, region->y1) - inRO->row0;
        int xMin = PS_MAX(minCols, region->x0) - inRO->col0;
        int xMax = PS_MIN(numCols - 1, region->x1) - inRO->col0;
        psTrace("psModules.imcombine", 2, "Normalising convolved mask image by %f over %d:%d,%d:%d\n",
                sum, xMin, xMax, yMin, yMax);
        for (int y = yMin; y <= yMax; y++) {
            for (int x = xMin; x <= xMax; x++) {
                convRO->image->data.F32[y][x] /= sum;
            }
        }
    }
    psFree(inRO);
    psImage *convolved = psMemIncrRefCounter(convRO->image);
    psFree(convRO);

#ifdef TESTING
    {
        static int seqNum = 0;          // Sequence number
        psString name = NULL;           // Name of image
        psStringAppend(&name, "inspect_conv_%02d.fits", seqNum);
        seqNum++;
        psFits *fits = psFitsOpen(name, "w"); // FITS file pointer
        psFree(name);
        psFitsWriteImage(fits, NULL, convolved, 0, NULL);
        psFitsClose(fits);
    }
#endif

    // Threshold the convolved image
    psPixels *bad = psPixelsAllocEmpty(PIXEL_LIST_BUFFER); // List of pixels that should be masked
    for (int y = size; y < convolved->numRows - size; y++) {
        for (int x = size; x < convolved->numCols - size; x++) {
            if (convolved->data.F32[y][x] > threshold) {
                // Pixel coordinates in "bad" correspond to the full image
                bad = psPixelsAdd(bad, PIXEL_LIST_BUFFER, x + minCols, y + minRows);
            }
        }
    }
    psFree(convolved);

    // Now, grow the mask to include everything that touches a bad pixel in the convolution
    int x0 = minCols, y0 = minRows;     // Offset for mask image
    mask = psPixelsToMask(NULL, bad, psRegionSet(x0, numCols - 1, y0, numRows - 1), 0xff);
    for (int i = 0; i < bad->n; i++) {
        int xPix = bad->data[i].x - x0, yPix = bad->data[i].y - y0; // Coordinates in frame of mask image
        // Convolution limits
        int xMin = PS_MAX(xPix - size, 0);
        int xMax = PS_MIN(xPix + size, mask->numCols - 1);
        int yMin = PS_MAX(yPix - size, 0);
        int yMax = PS_MIN(yPix + size, mask->numRows - 1);
        for (int y = yMin; y <= yMax; y++) {
            for (int x = xMin; x <= xMax; x++) {
                assert(x < mask->numCols && y < mask->numRows);
                mask->data.PS_TYPE_MASK_DATA[y][x] = 0xff;
            }
        }
    }
    bad = psPixelsFromMask(bad, mask, 0xff);
    psFree(mask);

    // Convert coordinates to frame of original image
    for (int i = 0; i < bad->n; i++) {
        int x = bad->data[i].x + x0;
        int y = bad->data[i].y + y0;
        if (x < 0 || x >= numCols || y < 0 || y >= numRows) {
            psWarning("Bad pixel coordinate %d: %d,%d --- ignored.",
                      i, x, y);
            continue;
        }
        bad->data[i].x = x;
        bad->data[i].y = y;
    }

    return bad;
}
