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

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

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

#include "pmSubtractionMask.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Private (file-static) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Mark a pixel as blank in the image, mask and weight
static inline void markBlank(psImage *image, // Image to mark as blank
                             psImage *mask, // Mask to mark as blank (or NULL)
                             psImage *weight, // Weight map to mark as blank (or NULL)
                             int x, int y, // Coordinates to mark blank
                             psImageMaskType blank // Blank mask value
    )
{
    image->data.F32[y][x] = NAN;
    if (mask) {
        mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] |= blank;
    }
    if (weight) {
        weight->data.F32[y][x] = NAN;
    }
    return;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

psImage *pmSubtractionMask(const psImage *mask1, const psImage *mask2, psImageMaskType maskVal,
                           int size, int footprint, float badFrac, bool useFFT)
{
    PS_ASSERT_IMAGE_NON_NULL(mask1, NULL);
    PS_ASSERT_IMAGE_TYPE(mask1, PS_TYPE_IMAGE_MASK, NULL);
    if (mask2) {
        PS_ASSERT_IMAGE_NON_NULL(mask2, NULL);
        PS_ASSERT_IMAGE_TYPE(mask2, PS_TYPE_IMAGE_MASK, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(mask2, mask1, NULL);
    }
    PS_ASSERT_INT_NONNEGATIVE(size, NULL);
    PS_ASSERT_INT_NONNEGATIVE(footprint, NULL);
    if (isfinite(badFrac)) {
        PS_ASSERT_FLOAT_LARGER_THAN(badFrac, 0.0, NULL);
        PS_ASSERT_FLOAT_LESS_THAN_OR_EQUAL(badFrac, 1.0, NULL);
    }

    int numCols = mask1->numCols, numRows = mask1->numRows; // Size of the images

    // Dereference inputs for convenience
    psImageMaskType **data1 = mask1->data.PS_TYPE_IMAGE_MASK_DATA;
    psImageMaskType **data2 = NULL;
    if (mask2) {
        data2 = mask2->data.PS_TYPE_IMAGE_MASK_DATA;
    }

    // First, a pass through to determine the fraction of bad pixels
    if (isfinite(badFrac) && badFrac != 1.0) {
        int numBad = 0;                 // Number of bad pixels
        for (int y = 0; y < numRows; y++) {
            for (int x = 0; x < numCols; x++) {
                if (data1[y][x] & maskVal) {
                    numBad++;
                    continue;
                }
                if (data2 && data2[y][x] & maskVal) {
                    numBad++;
                }
            }
        }
        if (numBad > badFrac * numCols * numRows) {
            psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                    "Fraction of bad pixels (%d/%d=%f) exceeds limit (%f)\n",
                    numBad, numCols * numRows, (float)numBad/(float)(numCols * numRows), badFrac);
            return NULL;
        }
    }

    // Worried about the masks for bad pixels and bad stamps colliding, so make our own mask
    psImage *mask = psImageAlloc(numCols, numRows, PS_TYPE_IMAGE_MASK); // The global mask
    psImageInit(mask, 0);
    psImageMaskType **maskData = mask->data.PS_TYPE_IMAGE_MASK_DATA; // Dereference for convenience

    // Block out a border around the edge of the image

    // Bottom stripe
    for (int y = 0; y < PS_MIN(size + footprint, numRows); y++) {
        for (int x = 0; x < numCols; x++) {
            maskData[y][x] |= PM_SUBTRACTION_MASK_BORDER;
        }
    }
    // Either side
    for (int y = PS_MIN(size + footprint, numRows); y < numRows - size - footprint; y++) {
        for (int x = 0; x < PS_MIN(size + footprint, numCols); x++) {
            maskData[y][x] |= PM_SUBTRACTION_MASK_BORDER;
        }
        for (int x = PS_MAX(numCols - size - footprint, 0); x < numCols; x++) {
            maskData[y][x] |= PM_SUBTRACTION_MASK_BORDER;
        }
    }
    // Top stripe
    for (int y = PS_MAX(numRows - size - footprint, 0); y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            maskData[y][x] |= PM_SUBTRACTION_MASK_BORDER;
        }
    }

    // XXX Could do something smarter here --- we will get images that are predominantly masked (where the
    // skycell isn't overlapped by a large fraction by the observation), so that convolving around every bad
    // pixel is wasting time.  As a first cut, I've put in a check on the fraction of bad pixels, but we could
    // imagine looking for the edge of big regions and convolving just at the edge.  As a second cut, allow
    // use of FFT convolution.

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (data1[y][x] & maskVal) {
                maskData[y][x] |= PM_SUBTRACTION_MASK_BAD_1;
            }
            if (data2 && data2[y][x] & maskVal) {
                maskData[y][x] |= PM_SUBTRACTION_MASK_BAD_2;
            }
        }
    }

    // Block out the entire stamp footprint around bad input pixels.

    // We want to block out with the CONVOLVE mask anything that would be bad if we convolved with a bad
    // reference pixel (within 'size').  Then we want to block out with the FOOTPRINT mask everything within a
    // footprint's distance of those (within 'footprint').

    bool oldThreads = psImageConvolveSetThreads(true); // Old value of threading for psImageConvolve

    if (!psImageConvolveMask(mask, mask, PM_SUBTRACTION_MASK_BAD_1,
                             PM_SUBTRACTION_MASK_CONVOLVE_1,
                             -size, size, -size, size)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to convolve bad pixels from mask 1.");
        psFree(mask);
        return NULL;
    }
    if (!psImageConvolveMask(mask, mask, PM_SUBTRACTION_MASK_BAD_2,
                             PM_SUBTRACTION_MASK_CONVOLVE_2,
                             -size, size, -size, size)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to convolve bad pixels from mask 2.");
        psFree(mask);
        return NULL;
    }

    psImageConvolveSetThreads(oldThreads);

    return mask;
}


bool pmSubtractionBorder(psImage *image, psImage *weight, psImage *mask,
                         int size, psImageMaskType blank)
{
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, false);
    if (mask) {
        PS_ASSERT_IMAGE_NON_NULL(mask, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(mask, image, false);
        PS_ASSERT_IMAGE_TYPE(mask, PS_TYPE_IMAGE_MASK, false);
    }
    if (weight) {
        PS_ASSERT_IMAGE_NON_NULL(weight, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(weight, image, false);
        PS_ASSERT_IMAGE_TYPE(weight, PS_TYPE_F32, false);
    }

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

    for (int y = size; y < numRows - size; y++) {
        for (int x = 0; x < size; x++) {
            markBlank(image, mask, weight, x, y, blank);
        }
        for (int x = numCols - size; x < numCols; x++) {
            markBlank(image, mask, weight, x, y, blank);
        }
    }
    for (int y = 0; y < size; y++) {
        for (int x = 0; x < numCols; x++) {
            markBlank(image, mask, weight, x, y, blank);
        }
    }
    for (int y = numRows - size; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            markBlank(image, mask, weight, x, y, blank);
        }
    }

    return true;
}


bool pmSubtractionMaskApply(psImage *image, psImage *weight, const psImage *mask, pmSubtractionMode mode)
{
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, false);
    if (weight) {
        PS_ASSERT_IMAGE_NON_NULL(weight, false);
        PS_ASSERT_IMAGE_TYPE(weight, PS_TYPE_F32, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(weight, image, false);
    }
    PS_ASSERT_IMAGE_NON_NULL(mask, false);
    PS_ASSERT_IMAGE_TYPE(mask, PS_TYPE_IMAGE_MASK, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(mask, image, false);

    bool maskVal = PM_SUBTRACTION_MASK_BORDER; // Value to mask
    switch (mode) {
      case PM_SUBTRACTION_MODE_1:
        maskVal |= PM_SUBTRACTION_MASK_CONVOLVE_1;
        break;
      case PM_SUBTRACTION_MODE_2:
        maskVal |= PM_SUBTRACTION_MASK_CONVOLVE_2;
        break;
      case PM_SUBTRACTION_MODE_DUAL:
        maskVal |= PM_SUBTRACTION_MASK_CONVOLVE_2 | PM_SUBTRACTION_MASK_CONVOLVE_2;
        break;
      case PM_SUBTRACTION_MODE_ERR:
      case PM_SUBTRACTION_MODE_UNSURE:
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unsuppored subtraction mode: %x", mode);
        return false;
    }

    int numCols = image->numCols, numRows = image->numRows; // Size of image
    psImageMaskType **maskData = mask->data.PS_TYPE_IMAGE_MASK_DATA; // Dereference mask

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (maskData[y][x] & maskVal) {
                image->data.F32[y][x] = NAN;
                if (weight) {
                    weight->data.F32[y][x] = NAN;
                }
            }
        }
    }

    return true;
}
