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



psArray *stacTransform(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
		       psArray **outErrors, // Transformed error images for output
		       const psArray *masks, // Masks of input images
		       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 NULL;
    }
    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 NULL;
    }

    if (errors && (*outErrors == NULL)) {
	// Allocate the output error images, if required
	*outErrors = psArrayAlloc(errors->n);
    } 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 NULL;
    }

    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 NULL;
	}
	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 NULL;
	    }
	}
    }

    // Arrays of images for return
    psArray *outImages = psArrayAlloc(nImages);	// Output images

    // 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

	// Initialise the output images
	psImage *outImage = psImageAlloc(nx, ny, PS_TYPE_F32);
	psImage *outError = psImageAlloc(nx, ny, PS_TYPE_F32);
	psTrace("stac.transform", 5, "Allocating space for transformed image, %dx%d\n", nx, ny);
	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;
	    }
	}

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

#if 0
	// Calculate scales; could be adapted for higher order than linear by moving this into the
	// pixel-dependent section and calculating derivatives
	double xscale = 0.5*sqrt(SQUARE(map->x->coeff[0][1]) + SQUARE(map->x->coeff[1][0]));
	double yscale = 0.5*sqrt(SQUARE(map->y->coeff[0][1]) + SQUARE(map->y->coeff[1][0]));
	double areascale = 0.25 / xscale / yscale;
#endif

	// Iterate over the output image pixels
	for (int y = 0; y < ny; y++) {
	    for (int x = 0; x < nx; x++) {
		// Transform!
		// Could no doubt make this quicker by calculating derivatives and incrementing by those.
		sky->x = (double)x + 0.5;
		sky->y = (double)y + 0.5;
		(void)psPlaneTransformApply(detector, map, sky);

		outImage->data.F32[y][x] = (psF32)psImagePixelInterpolate(image, detector->x, detector->y,
									  mask, 1, 0.0,
									  PS_INTERPOLATE_BILINEAR);
		outError->data.F32[y][x] = (psF32)p_psImageErrorInterpolateBILINEAR_F32(error, detector->x,
											detector->y, mask, 1,
											0.0);
		outError->data.F32[y][x] = sqrtf(outError->data.F32[y][x]);
	    }
	} // 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

	// Plug the new images into arrays
	outImages->data[n] = outImage;
	(*outErrors)->data[n] = outError;

    } // Iterating over images

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

    return outImages;
}

