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

#include <stdio.h>
#include <string.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, int numCols, int numRows, float threshold, float poorFrac,
                        const psArray *subRegions, const psArray *subKernels)
{
    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(subKernels, NULL);
    PS_ASSERT_ARRAYS_SIZE_EQUAL(subRegions, subKernels, NULL);

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

    // Check consistency of kernels
    int numRegions = subRegions->n;     // Number of regions
    int size = 0;                       // Size of kernel
    for (int i = 0; i < numRegions; i++) {
        pmSubtractionKernels *kernels = subKernels->data[i]; // Kernel of interest
        if (size == 0) {
            size = kernels->size;
        } else if (kernels->size != size) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Kernel sizes are not identical: %d vs %d",
                    size, kernels->size);
            return NULL;
        }
    }

    psImage *mask = psPixelsToMask(NULL, in, psRegionSet(0, numCols - 1, 0, 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;
    for (int i = 0; i < numRegions; i++) {
        psRegion *region = subRegions->data[i]; // Region of interest
        pmSubtractionKernels *kernels = subKernels->data[i]; // Kernel of interest
        if (!pmSubtractionConvolve(convRO, NULL, inRO, NULL, NULL, 0, 0, 1.0, region, kernels, 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) - kernels->numCols/2.0) /
            (float)kernels->numCols;
        float yNorm = (region->y0 + 0.5 * (region->y1 - region->y0) - kernels->numRows/2.0) /
            (float)kernels->numRows;
        psImage *image = pmSubtractionKernelImage(kernels, 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 xMin = PS_MAX(0, region->x0), xMax = PS_MIN(numCols - 1, region->x1);
        int yMin = PS_MAX(0, region->y0), yMax = PS_MIN(numRows - 1, region->y1);
        psTrace("psModules.imcombine", 2, "Normalising convolved mask image by %f over %d:%d,%d:%d\n",
                sum, xMin, xMax, yMin, yMax);
        sum = 1.0 / sum;
        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) {
                bad = psPixelsAdd(bad, bad->nalloc, x, y);
            }
        }
    }
    psFree(convolved);

    // Now, grow the mask to include everything that touches a bad pixel in the convolution
    // Mask values:
    // 0x0f: we think this is bad
    // 0xf0: this is within the convolution kernel of a bad pixel
    mask = psPixelsToMask(NULL, bad, psRegionSet(0, numCols - 1, 0, numRows - 1), 0x0f);
    for (int i = 0; i < subRegions->n; i++) {
        psRegion *region = subRegions->data[i]; // Subtraction region
        pmSubtractionKernels *kernels = subKernels->data[i]; // Subtraction kernel

        int size = kernels->size;           // Half-size of kernel
        int fullSize = 2 * size + 1;        // Full size of kernel

        // Get region for convolution: [xMin:xMax,yMin:yMax]
        int xMin = PS_MAX(region->x0, size), xMax = PS_MIN(region->x1, numCols - size);
        int yMin = PS_MAX(region->y0, size), yMax = PS_MIN(region->y1, numRows - size);

        psImage *polyValues = NULL;     // Pre-calculated polynomial values
        for (int j = yMin; j < yMax; j += fullSize) {
            int ySubMax = PS_MIN(j + fullSize, yMax); // Range for subregion of interest
            for (int i = xMin; i < xMax; i += fullSize) {
                int xSubMax = PS_MIN(i + fullSize, xMax); // Range for subregion of interest

                polyValues = p_pmSubtractionPolynomialFromCoords(polyValues, kernels, numCols, numRows,
                                                                 i + size + 1, j + size + 1);
                int box = p_pmSubtractionBadRadius(NULL, kernels, polyValues,
                                                   false, poorFrac); // Radius of bad box
                if (box > 0) {
                    // Convolve a subimage, then stick it in the original
                    psImage *subMask = psImageSubset(mask, psRegionSet(i - box, xSubMax + box,
                                                                       j - box, ySubMax + box)); // Subimage
                    psImage *convolved = psImageConvolveMask(NULL, subMask, 0x0f, 0xf0,
                                                             -box, box, -box, box); // Convolved mask
                    psFree(subMask);

                    int numBytes = (xSubMax - i) * PSELEMTYPE_SIZEOF(PS_TYPE_MASK); // Number of bytes to copy
                    psAssert(convolved->numCols - 2 * box == xSubMax - i, "Bad number of columns");
                    psAssert(convolved->numRows - 2 * box == ySubMax - j, "Bad number of rows");

                    for (int yTarget = j, ySource = box; yTarget < ySubMax; yTarget++, ySource++) {
                        memcpy(&mask->data.PS_TYPE_MASK_DATA[yTarget][i],
                               &convolved->data.PS_TYPE_MASK_DATA[ySource][box], numBytes);
                    }
                    psFree(convolved);
                }
            }
        }
        psFree(polyValues);
    }
    bad = psPixelsFromMask(bad, mask, 0xff);

    return bad;
}
