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

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

#include "pmSubtraction.h"
#include "pmSubtractionStamps.h"

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

// Free function for pmSubtractionStamp
void subtractionStampFree(pmSubtractionStamp *stamp // Stamp to free
    )
{
    psFree(stamp->matrix);
    psFree(stamp->vector);
    psFree(stamp->reference);
    psFree(stamp->input);
    psFree(stamp->weight);
    psFree(stamp->convolutions);
}

// Is this region OK?
static bool checkStampRegion(int x, int y, // Coordinates of stamp
                             const psRegion *region // Region of interest
    )
{
    if (!region) {
        return true;
    }
    return (x < region->x0 || x > region->x1 || y < region->y0 || y > region->y1) ?
        false : true;
}

// Is this position unmasked?
static bool checkStampMask(int x, int y, // Coordinates of stamp
                           const psImage *mask
    )
{
    if (!mask) {
        return true;
    }
    return (mask->data.PS_TYPE_MASK_DATA[y][x] & (PM_SUBTRACTION_MASK_BORDER |
                                                  PM_SUBTRACTION_MASK_FOOTPRINT | PM_SUBTRACTION_MASK_REJ)) ?
        false : true;
}

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

pmSubtractionStamp *pmSubtractionStampAlloc(pmSubtractionStampStatus status)
{
    pmSubtractionStamp *stamp = psAlloc(sizeof(pmSubtractionStamp)); // Stamp to return
    psMemSetDeallocator(stamp, (psFreeFunc)subtractionStampFree);

    stamp->x = NAN;
    stamp->y = NAN;
    stamp->flux = NAN;
    stamp->xNorm = NAN;
    stamp->yNorm = NAN;
    stamp->matrix = NULL;
    stamp->vector = NULL;
    stamp->status = status;

    stamp->reference = NULL;
    stamp->input = NULL;
    stamp->weight = NULL;
    stamp->convolutions = NULL;

    return stamp;
}


psArray *pmSubtractionFindStamps(psArray *stamps, const psImage *image, const psImage *subMask,
                                 const psRegion *region, float threshold, float spacing)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, NULL);
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(image, subMask, NULL);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_MASK, NULL);
    }
    PS_ASSERT_FLOAT_LARGER_THAN(spacing, 0.0, NULL);
    if (region) {
        if (psRegionIsNaN(*region)) {
            psString string = psRegionToString(*region);
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Input region (%s) contains NAN values", string);
            psFree(string);
            return false;
        }
        if (region->x0 < 0 || region->x1 > image->numCols ||
            region->y0 < 0 || region->y1 > image->numRows) {
            psString string = psRegionToString(*region);
            psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Input region (%s) does not fit in image (%dx%d)",
                    string, image->numCols, image->numRows);
            psFree(string);
            return false;
        }
    }

    int numRows = image->numRows, numCols = image->numCols; // Size of image

    // Get region in which to find stamps: [xMin:xMax,yMin:yMax]
    int xMin = 0, xMax = numCols, yMin = 0, yMax = numRows;
    if (region) {
        xMin = PS_MAX(region->x0, xMin);
        xMax = PS_MIN(region->x1, xMax);
        yMin = PS_MAX(region->y0, yMin);
        yMax = PS_MIN(region->y1, yMax);
    }
    int xSize = xMax - xMin, ySize = yMax - yMin; // Size of region of interest

    // Number of stamps
    int xNumStamps = xSize / spacing + 1;
    int yNumStamps = ySize / spacing + 1;

    int numFound = 0;                   // Number of stamps found

    if (stamps) {
        PS_ASSERT_INT_EQUAL(stamps->n, xNumStamps * yNumStamps, NULL);
        // Just in case, check for NULL stamps
        for (int i = 0; i < xNumStamps * yNumStamps; i++) {
            if (!stamps->data[i]) {
                stamps->data[i] = pmSubtractionStampAlloc(PM_SUBTRACTION_STAMP_REJECTED);
            }
        }
    } else {
        stamps = psArrayAlloc(xNumStamps * yNumStamps);
        for (int i = 0; i < xNumStamps * yNumStamps ; i++) {
            stamps->data[i] = pmSubtractionStampAlloc(PM_SUBTRACTION_STAMP_INIT);
        }
    }

    for (int j = 0, index = 0; j < yNumStamps; j++) {
        for (int i = 0; i < xNumStamps; i++, index++) {
            pmSubtractionStamp *stamp = stamps->data[index]; // Stamp of interest

            // Only find a new stamp if we need to
            if (stamp->status == PM_SUBTRACTION_STAMP_REJECTED ||
                stamp->status == PM_SUBTRACTION_STAMP_INIT) {

                // Find maximum non-masked value in the image section,
                // but don't include a footprint around the edge
                float fluxBest = threshold; // Flux of best stamp pixel
                int xBest = 0, yBest = 0; // Coordinates of best stamp

                // Bounds of region to search for stamp
                int yStart = yMin + j * ySize / (yNumStamps + 1);
                int yStop = yMin + (j + 1) * ySize / (yNumStamps + 1) - 1;
                int xStart = xMin + i * xSize / (xNumStamps + 1);
                int xStop = xMin + (i + 1) * xSize / (xNumStamps + 1) - 1;
                assert(yStop < numRows && xStop < numCols && yStart >= 0 && xStart >= 0);

                for (int y = yStart; y <= yStop ; y++) {
                    for (int x = xStart; x <= xStop ; x++) {
                        if (checkStampMask(x, y, subMask) && image->data.F32[y][x] > fluxBest) {
                            fluxBest = image->data.F32[y][x];
                            xBest = x;
                            yBest = y;
                        }
                    }
                }

                stamp->x = xBest;
                stamp->y = yBest;
                stamp->flux = fluxBest;

                // Reset the postage stamps since we're making a new stamp
                psFree(stamp->reference);
                psFree(stamp->input);
                psFree(stamp->weight);
                stamp->reference = stamp->input = stamp->weight = NULL;

                if (fluxBest > threshold) {
                    stamp->status = PM_SUBTRACTION_STAMP_CALCULATE;
                    numFound++;
                    psTrace("psModules.imcombine", 5, "Found stamp in region (%d,%d): %d,%d\n",
                            i, j, (int)stamp->x, (int)stamp->y);
                } else {
                    stamp->status = PM_SUBTRACTION_STAMP_NONE;
                }
            }
        }
    }

    psLogMsg("psModules.imcombine", PS_LOG_INFO, "Found %d stamps in %dx%d regions",
             numFound, xNumStamps, yNumStamps);

    return stamps;
}


psArray *pmSubtractionSetStamps(const psVector *x, const psVector *y, const psVector *flux,
                                const psImage *subMask, const psRegion *region)
{
    PS_ASSERT_VECTOR_NON_NULL(x, false);
    PS_ASSERT_VECTOR_TYPE(x, PS_TYPE_F32, false);
    PS_ASSERT_VECTOR_NON_NULL(y, false);
    PS_ASSERT_VECTOR_TYPE(y, PS_TYPE_F32, false);
    PS_ASSERT_VECTORS_SIZE_EQUAL(y, x, false);
    if (flux) {
        PS_ASSERT_VECTOR_NON_NULL(flux, false);
        PS_ASSERT_VECTOR_TYPE(flux, PS_TYPE_F32, false);
        PS_ASSERT_VECTORS_SIZE_EQUAL(flux, x, false);
    }
    if (subMask) {
        PS_ASSERT_IMAGE_NON_NULL(subMask, false);
        PS_ASSERT_IMAGE_TYPE(subMask, PS_TYPE_MASK, false);
    }

    int num = x->n;                     // Number of stamps
    psArray *stamps = psArrayAllocEmpty(num); // Stamps, to return

    for (int i = 0; i < num; i++) {
        int xStamp = x->data.F32[i] + 0.5, yStamp = y->data.F32[i] + 0.5; // Pixel coords of stamp
        if (!checkStampRegion(xStamp, yStamp, region)) {
            continue;
        }
        if (!checkStampMask(xStamp, yStamp, subMask)) {
            continue;
        }

        pmSubtractionStamp *stamp = pmSubtractionStampAlloc(PM_SUBTRACTION_STAMP_CALCULATE); // Stamp
        stamp->x = x->data.F32[i];
        stamp->y = y->data.F32[i];
        if (flux) {
            stamp->flux = flux->data.F32[i];
        }
        stamps->data[stamps->n++] = stamp;
    }

    return stamps;
}


bool pmSubtractionExtractStamps(psArray *stamps, psImage *reference, psImage *input, psImage *weight,
                                int footprint, int kernelSize)
{
    PS_ASSERT_ARRAY_NON_NULL(stamps, false);
    PS_ASSERT_IMAGE_NON_NULL(reference, false);
    PS_ASSERT_IMAGE_TYPE(reference, PS_TYPE_F32, false);
    if (input) {
        PS_ASSERT_IMAGE_NON_NULL(input, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(input, reference, false);
        PS_ASSERT_IMAGE_TYPE(input, PS_TYPE_F32, false);
    }
    if (weight) {
        PS_ASSERT_IMAGE_NON_NULL(weight, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(weight, reference, false);
        PS_ASSERT_IMAGE_TYPE(weight, PS_TYPE_F32, false);
    }
    PS_ASSERT_INT_NONNEGATIVE(footprint, false);
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);

    int numCols = reference->numCols, numRows = reference->numRows; // Size of images
    int size = kernelSize + footprint; // Size of postage stamps

    if (!weight) {
        // Use the input as a rough approximation to the variance map, and HOPE that it's positive.
        weight = input ? input : reference;
    }

    for (int i = 0; i < stamps->n; i++) {
        pmSubtractionStamp *stamp = stamps->data[i]; // Stamp of interest
        if (!stamp || stamp->status != PM_SUBTRACTION_STAMP_CALCULATE) {
            continue;
        }

        if (isnan(stamp->xNorm)) {
            stamp->xNorm = 2.0 * (stamp->x - (float)numCols/2.0) / (float)numCols;
        }
        if (isnan(stamp->yNorm)) {
            stamp->yNorm = 2.0 * (stamp->y - (float)numRows/2.0) / (float)numRows;
        }

        int x = stamp->x + 0.5, y = stamp->y + 0.5; // Stamp coordinates
        if (x < size || x > numCols - size || y < size || y > numRows - size) {
            psError(PS_ERR_UNKNOWN, false, "Stamp %d (%d,%d) is within the image border.\n", i, x, y);
            return false;
        }

        psRegion region = psRegionSet(x - size, x + size + 1, y - size, y + size + 1); // Region of interest

        psImage *refSub = psImageSubset(reference, region); // Subimage with stamp
        stamp->reference = psKernelAllocFromImage(refSub, size, size);
        psFree(refSub);                 // Drop reference

        if (input) {
            psImage *inSub = psImageSubset(input, region); // Subimage with stamp
            stamp->input = psKernelAllocFromImage(inSub, size, size);
            psFree(inSub);              // Drop reference
        }

        psImage *wtSub = psImageSubset(weight, region); // Subimage with stamp
        stamp->weight = psKernelAllocFromImage(wtSub, size, size);
        psFree(wtSub);                 // Drop reference
    }

    return true;
}


bool pmSubtractionGenerateStamps(psArray *stamps, float fwhm, int footprint, int kernelSize)
{
    PS_ASSERT_ARRAY_NON_NULL(stamps, false);
    PS_ASSERT_FLOAT_LARGER_THAN(fwhm, 0.0, false);
    PS_ASSERT_INT_NONNEGATIVE(footprint, false);

    int size = kernelSize + footprint; // Size of postage stamps
    int num = stamps->n;                // Number of stamps
    float sigma = fwhm / (2.0 * sqrtf(2.0 * log(2.0))); // Gaussian sigma

    for (int i = 0; i < num; i++) {
        pmSubtractionStamp *stamp = stamps->data[i]; // Stamp of interest
        if (!(stamp->status & PM_SUBTRACTION_STAMP_CALCULATE)) {
            continue;
        }

        float x = stamp->x, y = stamp->y; // Coordinates of stamp
        float flux = stamp->flux; // Flux of star
        if (!isfinite(flux)) {
            psWarning("Unable to generate PSF for stamp %d --- bad flux.", i);
            continue;
        }

        float xStamp = x - (int)(x + 0.5); // x coordinate of star in stamp frame
        float yStamp = y - (int)(y + 0.5); // y coordinate of star in stamp frame

        psFree(stamp->input);
        stamp->input = psKernelAlloc(-size, size, -size, size);
        psKernel *input = stamp->input; // Target stamp

        // Put in a Gaussian, just for fun!
        for (int v = -size; v <= size; v++) {
            for (int u = -size; u <= size; u++) {
                input->kernel[v][u] = flux / sigma * 0.5 * M_2_SQRTPI * M_SQRT1_2 *
                    expf(-0.5 * (PS_SQR(u + xStamp) + PS_SQR(v + yStamp)) / PS_SQR(sigma));
            }
        }

    }

    return stamps;
}
