#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

psPixels *pmStackReject(const psPixels *in, float threshold, const psArray *regions,
                        const psArray *solutions, const pmSubtractionKernels *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(regions, NULL);
    PS_ASSERT_ARRAY_NON_NULL(solutions, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(regions, solutions, NULL);
    PS_ASSERT_PTR_NON_NULL(kernels, NULL);

    // Get the original image size
    int numRegions = regions->n;        // Number of regions
    int numCols = 0, numRows = 0;       // Size of original image
    int minCols = INT_MAX, minRows = INT_MAX; // Minimum coordinate for image --- should be 0,0
    for (int i = 0; i < numRegions; i++) {
        psRegion *region = regions->data[i]; // Region of interest
        if (region->x0 < minCols) {
            minCols = region->x0;
        }
        if (region->y0 < minRows) {
            minRows = region->y0;
        }
        if (region->x1 > numCols) {
            numCols = region->x1;
        }
        if (region->y1 > numRows) {
            numRows = region->y1;
        }
    }
    if (minCols != 0 || minRows != 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "Some error with image regions --- minimum coordinate is not 0,0");
        return NULL;
    }

    psImage *mask = psPixelsToMask(NULL, in, psRegionSet(0, numCols, 0, numRows), 0x01); // Mask image
    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;
    for (int i = 0; i < numRegions; i++) {
        psRegion *region = regions->data[i]; // Region of interest
        if (!pmSubtractionConvolve(convRO, NULL, inRO, NULL, NULL, 0, region, kernels, 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) - numCols/2.0) / (float)numCols;
        float yNorm = (region->y0 + 0.5 * (region->y1 - region->y0) - numRows/2.0) / (float)numRows;
        psImage *kernel = pmSubtractionKernelImage(kernels, xNorm, yNorm, false);
        if (!kernel) {
            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 < kernel->numRows; y++) {
            for (int x = 0; x < kernel->numCols; x++) {
                sum += kernel->data.F32[y][x];
            }
        }
        psFree(kernel);

        psImage *subConv = psImageSubset(convRO->image, *region); // Sub-image of convolved image
        psBinaryOp(subConv, subConv, "*", psScalarAlloc(1.0 / sum, PS_TYPE_F32));
    }
    psFree(inRO);
    psImage *convolved = psMemIncrRefCounter(convRO->image);
    psFree(convRO);

    // Threshold the convolved image
    psPixels *bad = psPixelsAllocEmpty(PIXEL_LIST_BUFFER); // List of pixels that should be masked
    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (convolved->data.F32[y][x] > threshold) {
                bad = psPixelsAdd(bad, PIXEL_LIST_BUFFER, x, y);
            }
        }
    }
    psFree(convolved);

    // Now, we want to convolve the original pixels properly
    mask = psPixelsToMask(NULL, bad, psRegionSet(0, numCols, 0, numRows), 0xff);
    int size = kernels->size;           // Size of kernels
    for (int i = 0; i < bad->n; i++) {
        int xPix = bad->data[i].x, yPix = bad->data[i].y; // Coordinates of interest
        // Convolution limits
        int xMin = PS_MAX(xPix - size, 0);
        int xMax = PS_MIN(xPix + size, numCols - 1);
        int yMin = PS_MAX(yPix - size, 0);
        int yMax = PS_MIN(yPix + size, numRows - 1);
        for (int y = yMin; y <= yMax; y++) {
            for (int x = xMin; x <= xMax; x++) {
                mask->data.PS_TYPE_MASK_DATA[y][x] = 0xff;
            }
        }
    }

    bad = psPixelsFromMask(bad, mask, 0xff);
    psFree(mask);

    return bad;
}
