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

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

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

// Hacked the original ps_ImagePixelInterpolateBILINEAR_F32 to add variances
// i.e., to square the fractions when combining.
inline psF64 p_psImageErrorInterpolateBILINEAR_F32(const psImage* input,
                                                   float x,
                                                   float y,
                                                   const psImage* mask,
                                                   unsigned int maskVal,
                                                   psF64 unexposedValue)
{
    double floorX = floor((psF64)(x) - 0.5);
    double floorY = floor((psF64)(y) - 0.5);
    psF64 fracX = x - 0.5 - floorX;
    psF64 fracY = y - 0.5 - floorY;
    int intFloorX = (int) floorX;
    int intFloorY = (int) floorY;
    int lastX = input->numCols - 1;
    int lastY = input->numRows - 1;
    psF32 V00 = 0.0;
    psF32 V01 = 0.0;
    psF32 V10 = 0.0;
    psF32 V11 = 0.0;
    bool valid00 = false;
    bool valid01 = false;
    bool valid10 = false;
    bool valid11 = false;

    if (intFloorY >= 0 && intFloorY <= lastY) {
        if (intFloorX >= 0 && intFloorX <= lastX) {
            V00 = input->data.F32[intFloorY][intFloorX];
            valid00 = true;
        }
        if (intFloorX >= -1 && intFloorX < lastX) {
            V10 = input->data.F32[intFloorY][intFloorX+1];
            valid10 = true;
        }
    }
    if (intFloorY >= -1 && intFloorY < lastY) {
        if (intFloorX >= 0 && intFloorX <= lastX) {
            V01 = input->data.F32[intFloorY+1][intFloorX];
            valid01 = true;
        }
        if (intFloorX >= -1 && intFloorX < lastX) {
            V11 = input->data.F32[intFloorY+1][intFloorX+1];
            valid11 = true;
        }
    }

    /* cover likely case of all pixels being valid more efficiently */
    if (valid00 && valid10 && valid01 && valid11) {
        /* formula from the ADD */
        return V00*SQUARE((1.0-fracX)*(1.0-fracY)) + V10*SQUARE(fracX*(1.0-fracY)) +
            V01*SQUARE(fracY*(1.0-fracX)) + V11*SQUARE(fracX*fracY);
    }

    /* OK, at least one pixel is not valid - need to do it piecemeal */

    psF64 V0 = 0.0;
    bool valid0 = true;
    if (valid00 && valid10) {
        V0 = V00*SQUARE(1-fracX)+V10*SQUARE(fracX);
    } else if (valid00) {
        V0 = V00;
    } else if (valid10) {
        V0 = V10;
    } else {
        valid0 = false;
    }

    psF64 V1 = 0.0;
    bool valid1 = true;
    if (valid01 && valid11) {
        V1 = V01*SQUARE(1-fracX)+V11*SQUARE(fracX);
    } else if (valid01) {
        V1 = V01;
    } else if (valid11) {
        V1 = V11;
    } else {
        valid1 = false;
    }

    if (valid0 && valid1) {
        return ( V0*SQUARE(1-fracY) + V1*SQUARE(fracY) );
    } else if (valid0) {
        return V0;
    } else if (valid1) {
        return V1;
    }

    return unexposedValue;
}



bool stacTransform(psArray **outputs,   // Transformed images for output
                   psArray **outErrors, // Transformed error images for output
                   const psArray *images, // Array of images to be transformed
                   const psArray *maps, // Array of polynomials that do the transformation
                   const psArray *errors, // Array of error images to be transformed
                   const psArray *masks, // Masks of input images
                   const psImage *region, // Region of interest for transformation
                   const psVector *scales, // Relative scales
                   const psVector *offsets, // Relative offsets
                   int outnx, int outny // Size of output images
    )
{
    int nImages = images->n;            // Number of images

    // Check input sizes
    assert(images->n == maps->n);
    assert(!errors || (images->n == errors->n));
    assert(!scales || scales->n == images->n);
    assert(!offsets || offsets->n == images->n);
    assert(!scales || scales->type.type == PS_TYPE_F32);
    assert(!offsets || offsets->type.type == PS_TYPE_F32);

    // Allocate the output images if required, otherwise check the number
    assert(!*outputs || (*outputs)->n == nImages);
    if (*outputs == NULL) {
        *outputs = psArrayAlloc(nImages);
        psTrace("stac.transform", 5, "Allocating space for transformed images, %dx%d\n", outnx, outny);
        for (int i = 0; i < nImages; i++) {
            (*outputs)->data[i] = psImageAlloc(outnx, outny, PS_TYPE_F32);
            psImageInit((*outputs)->data[i], 0.0);
        }
    }

    // Allocate the output error images, if required, otherwise check the number
    assert(!errors || ! *outErrors || errors->n == (*outErrors)->n);
    if (errors && (*outErrors == NULL)) {
        *outErrors = psArrayAlloc(errors->n);
        psTrace("stac.transform", 5, "Allocating space for transformed error images, %dx%d\n", outnx, outny);
        for (int i = 0; i < nImages; i++) {
            (*outErrors)->data[i] = psImageAlloc(outnx, outny, PS_TYPE_F32);
        }
    }

    // Check the masks, if specified
    assert(!masks || masks->n == nImages);
    if (masks != NULL) {
        for (int i = 0; i < nImages; i++) {
            psImage *image = images->data[i];
            psImage *mask = masks->data[i];
            assert(mask->numRows == image->numRows && mask->numCols == image->numCols);
        }
    }


    // Stuff for the transformations
    psPlane *detector = psAlloc(sizeof(psPlane)); // Coordinates on the detector
    psPlane *sky = psAlloc(sizeof(psPlane)); // Coordinates on the sky


    // Iterate over the images
    for (int n = 0; n < nImages; n++) {
        psTrace("stac.transform", 1, "Transforming image %d....\n",n);

        // Pull out the various stuff we're working on
        psImage *image = images->data[n]; // The input image
        psPlaneTransform *map = maps->data[n]; // The map
        psImage *outImage = (*outputs)->data[n]; // The output image
        psImage *error = NULL; // The error image
        psImage *outError = NULL; // The output error image
        if (errors) {
            error = errors->data[n];
            outError = (*outErrors)->data[n];
        }
        float offset = 0.0;             // Relative offset
        float scale = 1.0;              // Relative scale
        if (offsets) {
            offset = offsets->data.F32[n];
        }
        if (scales) {
            scale = scales->data.F32[n];
        }

        // Mask
        psImage *mask = NULL;
        if (masks != NULL) {
            mask = masks->data[n];
        }

        psImageInterpolateOptions *interp = psImageInterpolateOptionsAlloc(PS_INTERPOLATE_BILINEAR,
                                                                           image, error, mask, 1, NAN, NAN,
                                                                           1, 0, 0.0);

        // Iterate over the output image pixels
        for (int y = 0; y < outny; y++) {
            for (int x = 0; x < outnx; x++) {
                // Only transform those pixels requested
                if (!region || (region && region->data.U8[y][x])) {
                    // Transform!
                    sky->x = (double)x + 0.5;
                    sky->y = (double)y + 0.5;
                    (void)psPlaneTransformApply(detector, map, sky);

                    // Change PS_INTERPOLATE_BILINEAR to best available technique.
                    double imageValue, varianceValue;
                    if (!psImageInterpolate(&imageValue, &varianceValue, NULL,
                                            detector->x, detector->y, interp)) {
                        psError(PS_ERR_UNKNOWN, false, "Unable to interpolate image.");
                        psFree(interp);
                        psFree(detector);
                        psFree(sky);
                        return false;
                    }
                    outImage->data.F32[y][x] = imageValue;

                    if (error) {
                        // Error is actually the variance
                        outError->data.F32[y][x] = varianceValue;
                    }

                    outImage->data.F32[y][x] = (outImage->data.F32[y][x] - offset) / scale;
                    if (error) {
                        outError->data.F32[y][x] = outError->data.F32[y][x] / SQUARE(scale);
                    }

                } // Pixels of interest

            }
        } // Iterating over output pixels
        psFree(interp);

    } // Iterating over images

    // Done with transformations
    psFree(detector);
    psFree(sky);

    return true;
}

