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

#include <stdio.h>
#include <pslib.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);
}


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

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

    stamp->x = 0;
    stamp->y = 0;
    stamp->matrix = NULL;
    stamp->vector = NULL;
    stamp->status = status;

    return stamp;
}


psArray *pmSubtractionFindStamps(psArray *stamps, const psImage *image, const psImage *mask,
                                 psMaskType maskVal, psMaskType used, float threshold,
                                 float spacing, int border)
{
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, NULL);
    if (mask) {
        PS_ASSERT_IMAGE_NON_NULL(mask, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(image, mask, NULL);
        PS_ASSERT_IMAGE_TYPE(mask, PS_TYPE_MASK, NULL);
    }
    if (used == 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "The mask value for used stamps cannot be zero.");
        return NULL;
    }
    PS_ASSERT_FLOAT_LARGER_THAN(spacing, 0.0, NULL);
    PS_ASSERT_INT_NONNEGATIVE(border, NULL);
    PS_ASSERT_INT_LARGER_THAN(image->numCols, 2 * border, NULL);
    PS_ASSERT_INT_LARGER_THAN(image->numRows, 2 * border, NULL);

    psMaskType badCentral = maskVal | used; // For central pixel of stamp, don't get a stamp already used
    int xNumStamps = (image->numCols - 2 * border) / spacing + 1; // Number of stamps in x
    int yNumStamps = (image->numRows - 2 * border) / spacing + 1; // Number of stamps in y

    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);
        }
    }
    // Footprint of image
    int numRows = image->numRows;
    int numCols = image->numCols;

    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 yMin = border + j * (numRows - border) / (yNumStamps + 1);
                int yMax = border + (j + 1) * (numRows - border) / (yNumStamps + 1) - 1;
                int xMin = border + i * (numCols - border) / (xNumStamps + 1);
                int xMax = border + (i + 1) * (numCols - border) / (xNumStamps + 1) - 1;
                assert(yMax < image->numRows - border && xMax < image->numCols - border &&
                       yMin >= border && xMin >= border);

                for (int y = yMin; y <= yMax ; y++) {
                    for (int x = xMin; x <= xMax ; x++) {
                        if (image->data.F32[y][x] > fluxBest) {
                            bool ok = true;
                            // Check kernel footprint for bad pixels
                            if (mask && (mask->data.PS_TYPE_MASK_DATA[y][x] & badCentral)) {
                                ok = false;
                            } else {
                                for (int v = -border; v <= border && ok; v++) {
                                    for (int u = -border; u <= border && ok; u++) {
                                        if ((mask &&
                                             (mask->data.PS_TYPE_MASK_DATA[y + v][x + u] & maskVal)) ||
                                            !isfinite(image->data.F32[y + v][x + u])) {
                                            ok = false;
                                            if (mask) {
                                                // Mark it so we don't have to look so hard next time
                                                mask->data.PS_TYPE_MASK_DATA[y][x] |= used;
                                            }
                                        }
                                    }
                                }
                            }

                            if (ok) {
                                fluxBest = image->data.F32[y][x];
                                xBest = x;
                                yBest = y;
                            }
                        }
                    }
                }

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

    return stamps;
}
