#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);
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// 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 *subMask,
                                 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);

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

    // Number of stamps
    int xNumStamps = numCols / spacing + 1;
    int yNumStamps = numRows / spacing + 1;

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

                for (int y = yMin; y <= yMax ; y++) {
                    for (int x = xMin; x <= xMax ; x++) {
                        if ((!subMask || !(subMask->data.PS_TYPE_MASK_DATA[y][x] &
                                           (PM_SUBTRACTION_MASK_BORDER | PM_SUBTRACTION_MASK_FOOTPRINT |
                                            PM_SUBTRACTION_MASK_REJ))) &&
                            image->data.F32[y][x] > fluxBest) {
                            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;
}
