#include <stdio.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))

#define MAXCHAR 80

static int numTransforms = 0;		// Number of transformations performed


// 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; 
    psF32 V01; 
    psF32 V10; 
    psF32 V11; 
    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; 
    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; 
    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 stacConfig *config // Configuration
    )
{
    int nImages = images->n;		// Number of images
    int nx = config->outnx, ny = config->outny; // Size of output images
    numTransforms++;

    // Check input sizes
    if (images->n != maps->n) {
	psError("stac.transform", "Number of maps (%d) does not match number of images (%d).\n",
		maps->n, images->n);
	return false;
    }
    if (errors && (images->n != errors->n)) {
	psError("stac.transform", "Number of error images (%d) does not match number of images (%d).\n",
		errors->n, images->n);
	return false;
    }

    // Allocate the output images if required, otherwise check the number
    if (*outputs == NULL) {
	*outputs = psArrayAlloc(nImages);
	psTrace("stac.transform", 5, "Allocating space for transformed images, %dx%d\n", nx, ny);
	for (int i = 0; i < nImages; i++) {
	    (*outputs)->data[i] = psImageAlloc(nx, ny, PS_TYPE_F32);
	}
    } else if ((*outputs)->n != nImages) {
	psError("stac.transform", "Number of output images (%d) does not match number of input images "
		"(%d).\n", (*outputs)->n, nImages);
	return false;
    }

    // Allocate the output error images, if required, otherwise check the number
    if (errors && (*outErrors == NULL)) {
	*outErrors = psArrayAlloc(errors->n);
	psTrace("stac.transform", 5, "Allocating space for transformed error images, %dx%d\n", nx, ny);
	for (int i = 0; i < nImages; i++) {
	    (*outErrors)->data[i] = psImageAlloc(nx, ny, PS_TYPE_F32);
	}
    } else if (errors->n != (*outErrors)->n) {
	psError("stac.transform", "Number of error output images (%d) does not match number of input"
		"images (%d).\n", (*outErrors)->n, errors->n);
	return false;
    }

    // Check the masks, if specified
    if (masks != NULL) {
	if (masks->n != nImages) {
	    psError("stac.transform", "Number of masks (%d) does not match number of input images (%d).\n",
		    masks->n, nImages);
	    return false;
	}
	for (int i = 0; i < nImages; i++) {
	    psImage *image = images->data[i];
	    psImage *mask = masks->data[i];
	    if ((mask->numRows != image->numRows) || (mask->numCols != image->numCols)) {
		psError ("stac.transform",
			 "Size of input mask (%dx%d) does not match that of input image (%dx%d).\n",
			 mask->numCols, mask->numRows, image->numCols, image->numRows);
		return false;
	    }
	}
    }


    // 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 *error = errors->data[n]; // The error image
	psImage *outImage = (*outputs)->data[n]; // The output image
	psImage *outError = (*outErrors)->data[n]; // The output error image

#if 0
	// No need for initialisation, since we iterate over the entire output image.
	for (int y = 0; y < ny; y++) {
	    for (int x = 0; x < nx; x++) {
		outImage->data.F32[y][x] = 0.0;
		outError->data.F32[y][x] = 0.0;
	    }
	}
#endif

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

	// Iterate over the output image pixels
	for (int y = 0; y < ny; y++) {
	    for (int x = 0; x < nx; 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.
		    outImage->data.F32[y][x] = (psF32)psImagePixelInterpolate(image, detector->x + 0.5,
									      detector->y + 0.5, mask, 1, 0.0,
									      PS_INTERPOLATE_BILINEAR);
		    // Error is actually the variance
		    outError->data.F32[y][x] = (psF32)p_psImageErrorInterpolateBILINEAR_F32(error,
											    detector->x + 0.5,
											    detector->y + 0.5,
											    mask, 1, 0.0);
		} // Pixels of interest

	    }
	} // Iterating over output pixels

#ifdef TESTING
	// Write error image out to check
	char shiftfile[MAXCHAR];	// Filename of shift image
	char errfile[MAXCHAR];		// Filename of error image
	sprintf(shiftfile,"%s.shift.%d",config->inputs->data[n],numTransforms);
	sprintf(errfile,"%s.shifterr.%d",config->inputs->data[n],numTransforms);
	psTrace("stac.transform.test", 6,
		"Output files have size: %dx%d\n",outImage->numCols,outImage->numRows);
	psImageWriteSection(outImage,0,0,0,NULL,0,shiftfile);
	psImageWriteSection(outError,0,0,0,NULL,0,errfile);
#endif

    } // Iterating over images

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

    return true;
}