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

#include <stdio.h>
#include <assert.h>
#include "pslib.h"
#include "stac.h"


#define MAX(x,y) ((x) > (y) ? (x) : (y))
#define MIN(x,y) ((x) < (y) ? (x) : (y))

float stacGradient(psImage *image,      // Input for which to measure the gradient
                   int x, int y         // Coordinates at which to measure the gradient
    )
{
    int num = 0;
    psVector *pixels = psVectorAllocEmpty(8, PS_TYPE_F32); // Array of pixels
    psVector *mask = psVectorAllocEmpty(8, PS_TYPE_U8); // Corresponding mask

    // Get limits
    int xMin = MAX(x - 1, 0);
    int xMax = MIN(x + 1, image->numCols - 1);
    int yMin = MAX(y - 1, 0);
    int yMax = MIN(y + 1, image->numRows - 1);
    for (int j = yMin; j <= yMax; j++) {
        for (int i = xMin; i <= xMax; i++) {
            if ((i != x) && (j != y)) {
                pixels->data.F32[num] = image->data.F32[j][i];
                mask->data.U8[num] = 0;
                num++;
            }
        }
    }
    pixels->n = num;
    mask->n = num;

    // Get the median
    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN);
    (void)psVectorStats(stats, pixels, NULL, mask, 1);
    float median = stats->sampleMedian;
    psFree(stats);
    psFree(pixels);
    psFree(mask);
    return median / image->data.F32[y][x];
}

psArray *stacRejection(psArray *inputs, // Input images
                       psArray *rejected, // Rejected images
                       psArray *regions, // Regions of interest
                       psArray *maps,   // Maps from input to transformed image
                       psArray *inverseMaps, // Maps from transformed to input image
                       float frac,      // Fraction of pixel rejected before looking more carefully
                       float grad,      // Gradient limit for rejection
                       psArray *names   // Names of original images (only for writing out when TESTING)
    )
{
    int nImages = inputs->n;            // Number of input images

    // Check inputs
    assert(inputs->n == rejected->n);
    assert(inputs->n == regions->n);
    assert(inputs->n == maps->n);
    assert(inputs->n == inverseMaps->n);
    assert(!names || names->n == inputs->n);

    for (int i = 0; i < nImages; i++) {
        psImage *input = inputs->data[i];
        psImage *region = regions->data[i];
        assert(input->numRows == region->numRows && input->numCols == region->numCols);
    }

    // Stuff for the transformations
    psPlane *inCoords = psAlloc(sizeof(psPlane)); // Coordinates on the input
    psPlane *outCoords = psAlloc(sizeof(psPlane)); // Coordinates on the output

    psTrace("stac.rejection", 1, "Mapping rejection masks back to source....\n");

    // Vectors for calculating the mean gradient
    psVector *grads = psVectorAlloc(nImages, PS_TYPE_F32); // Gradient for each image
    psVector *gradsMask = psVectorAlloc(nImages, PS_TYPE_U8); // Mask for gradient vector

    // Transform rejection masks back to source
    psArray *inputRej = psArrayAlloc(nImages);
    for (int i = 0; i < nImages; i++) {
        // Size of input image
        int nxInput = ((psImage*)(inputs->data[i]))->numCols;
        int nyInput = ((psImage*)(inputs->data[i]))->numRows;
        psImage *mask = psImageAlloc(nxInput, nyInput, PS_TYPE_U8); // The pixel mask for the input
#ifdef TESTING
        psImage *rejmap = psImageAlloc(nxInput, nyInput, PS_TYPE_F32); // The rejections in the source frame
        psImage *gradient = psImageAlloc(nxInput, nyInput, PS_TYPE_F32);        // The gradient image
#endif
        psImage *reject = rejected->data[i]; // Pull out the mask in the output frame
        psPlaneTransform *map = maps->data[i]; // The map from input to output
        psImage *region = regions->data[i]; // The region of interest for this image

        psTrace("stac.rejection", 3, "Transforming rejection mask %d....\n", i);
        int nBad = 0;                   // Number of bad pixels

#ifdef CRFLUX
        // Set up CR output
        FILE *crs = NULL;               // File for outputting details of rejected pixels
        char crfile[MAXCHAR];   // Filename
        sprintf(crfile,"%s.cr",names->data[i]);
        if ((crs = fopen(crfile, "w")) == NULL) {
            fprintf(stderr, "Unable to open file for detailed output, %s\n");
            return NULL;
        }
#endif

        psImageInterpolateOptions *interp = psImageInterpolateOptionsAlloc(PS_INTERPOLATE_BILINEAR,
                                                                           reject, NULL, NULL, 0, 0.0, 0.0,
                                                                           0, 0, 0.0);

        // Transform the mask
        // Optimisation option is to only transform the pixels that have been rejected in the output,
        // calculate derivatives of the map, and use that as a buffer around the transformed position
        // in the input image.
        psStats *median = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN);
        for (int y = 0; y < nyInput; y++) {
            for (int x = 0; x < nxInput; x++) {

                // Only transform pixels of interest
                if (region->data.U8[y][x]) {

                    inCoords->x = (double)x + 0.5;
                    inCoords->y = (double)y + 0.5;
                    (void)psPlaneTransformApply(outCoords, map, inCoords);
                    double maskVal;
                    if (!psImageInterpolate(&maskVal, NULL, NULL, outCoords->x, outCoords->y, interp)) {
                        psError(PS_ERR_UNKNOWN, false, "Unable to interpolate image.");
                        psFree(grads);
                        psFree(gradsMask);
                        psFree(inCoords);
                        psFree(outCoords);
                        return NULL;
                    }

#ifdef TESTING
                    rejmap->data.F32[y][x] = maskVal;
#endif

                    if (maskVal > frac) {
                        // Calculate mean gradient on other images
                        float meanGrads = 0.0;
                        int numGrads = 0;
                        for (int j = 0; j < nImages; j++) {
                            if (i != j) {
                                // Get coordinates for that image
                                (void)psPlaneTransformApply(inCoords, inverseMaps->data[j], outCoords);
                                int xPix = (int)(inCoords->x + 0.5);
                                int yPix = (int)(inCoords->y + 0.5);
                                if ((xPix >= 0) && (xPix <= ((psImage*)(inputs->data[j]))->numCols - 1) &&
                                    (yPix >= 0) && (yPix <= ((psImage*)(inputs->data[j]))->numRows - 1)) {
                                    // Calculate the gradient
                                    grads->data.F32[j] = stacGradient(inputs->data[j], xPix, yPix);
                                    if (isfinite(grads->data.F32[j])) {
                                        gradsMask->data.U8[j] = 0;
                                        numGrads++;
                                    } else {
                                        gradsMask->data.U8[j] = 1;
                                    }
                                } else {
                                    gradsMask->data.U8[j] = 1; // Mask this one
                                }
                            } else {
                                gradsMask->data.U8[j] = 1; // Mask this one
                            }
                        }
                        if (numGrads > 0) {
                            (void)psVectorStats(median, grads, NULL, gradsMask, (psU8)1);
                            meanGrads = median->sampleMedian;
                        } else {
                            meanGrads = 0.0;
                        }

#ifdef TESTING
                        //gradient->data.F32[y][x] = stacGradient(inputs->data[i], x, y) / meanGrads;
                        gradient->data.F32[y][x] = meanGrads;
#endif

                        //if (stacGradient(inputs->data[i], x, y) < grad * meanGrads) {
                        if (meanGrads > grad) {
                            mask->data.U8[y][x] = 1;
                            nBad++;

#ifdef CRFLUX
                            fprintf(crs, "%d %d --> %f %f %f\n", x, y,
                                    ((psImage*)(inputs->data[i]))->data.F32[y][x], maskVal,
                                    stacGradient(inputs->data[i], x, y));
#endif

                        } else {
                            mask->data.U8[y][x] = 0;
                        } // Gradient threshold
                    } else {
                        mask->data.U8[y][x] = 0;
                    } // Rejection threshold

                } else {
                    mask->data.U8[y][x] = 0;
                } // Only touching pixels of interest

            }
        } // Iterating over pixels
        psFree(median);

#ifdef CRFLUX
        // Close file used for detailed output
        fclose(crs);
#endif

        psTrace("stac.rejection", 1, "%d pixels masked in image %d\n", nBad, i);
        // Clip the image, and convert to suitable mask format

#ifdef TESTING
        // Write error image out to check
        char maskName[MAXCHAR];         // Filename of mask image
        char rejmapName[MAXCHAR];       // Filename of rejection image
        char gradName[MAXCHAR];         // Filename of gradient image
        sprintf(maskName, "%s.mask", (char*)names->data[i]);
        sprintf(rejmapName, "%s.rejmap", (char*)names->data[i]);
        sprintf(gradName, "%s.grad", (char*)names->data[i]);

        psFits *maskFile = psFitsOpen(maskName, "w");
        psFits *rejmapFile = psFitsOpen(rejmapName, "w");
        psFits *gradFile = psFitsOpen(gradName, "w");
        if (!psFitsWriteImage(maskFile, NULL, mask, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", maskName);
        }
        psTrace("stac", 1, "Mask image written to %s\n", maskName);
        if (!psFitsWriteImage(rejmapFile, NULL, rejmap, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", rejmapName);
        }
        psTrace("stac", 1, "Rejection map written to %s\n", rejmapName);
        if (!psFitsWriteImage(gradFile, NULL, gradient, 0, NULL)) {
            psErrorStackPrint(stderr, "Unable to write image: %s\n", gradName);
        }
        psTrace("stac", 1, "Gradient image written to %s\n", gradName);
        psFitsClose(maskFile);
        psFitsClose(rejmapFile);
        psFitsClose(gradFile);
        psFree(rejmap);
        psFree(gradient);
#endif

        // Stuff into the array
        inputRej->data[i] = mask;

    }

    psFree(grads);
    psFree(gradsMask);

    psFree(inCoords);
    psFree(outCoords);

    return inputRej;
}
