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


psArray *stacRejection(psArray *inputs,	// Input images
		       psArray *transformed, // Transformed images
		       psImage *combined, // Combined image
		       psArray *maps,	// Maps from input to transformed image
		       psArray *rejected, // Rejected images
		       stacConfig *config // Configuration
    )
{
    int nImages = inputs->n;		// Number of input images

    // Check inputs
    if (inputs->n != transformed->n) {
	psError("stac.rejection", "Number of input (%d) and transformed (%d) images does not match.\n",
		inputs->n, transformed->n);
	return NULL;
    }
    for (int i = 0; i < nImages; i++) {
	if ((((psImage*)(transformed->data[i]))->numRows != combined->numRows) ||
	    (((psImage*)(transformed->data[i]))->numCols != combined->numCols)) {
	    psError("stac.rejection",
		    "Sizes of transformed (%dx%d) and combined (%dx%d) images do not match.\n",
		    ((psImage*)(transformed->data[i]))->numCols, ((psImage*)(transformed->data[i]))->numRows,
		    combined->numCols, combined->numRows);
	    return NULL;
	}
	if ((((psImage*)(transformed->data[i]))->numRows != ((psImage*)(rejected->data[i]))->numRows) ||
	    (((psImage*)(transformed->data[i]))->numCols != ((psImage*)(rejected->data[i]))->numCols)) {
	    psError("stac.rejection",
		    "Sizes of transformed (%dx%d) and rejected (%dx%d) images do not match.\n",
		    ((psImage*)(transformed->data[i]))->numCols, ((psImage*)(transformed->data[i]))->numRows,
		    ((psImage*)(rejected->data[i]))->numCols, ((psImage*)(rejected->data[i]))->numRows);
	    return NULL;
	}
    }

    // Size of output images
    int nxOutput = combined->numCols;
    int nyOutput = combined->numRows;

    // 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");

    // 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_F32); // The mask in the source frame
	psImage *reject = rejected->data[i]; // Pull out the mask in the output frame
	psPlaneTransform *map = maps->data[i]; // The map from input to output

	psTrace("stac.rejection", 3, "Transforming rejection mask %d....\n", i);

	// 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.
	for (int y = 0; y < nyInput; y++) {
	    for (int x = 0; x < nxInput; x++) {
    		inCoords->x = (double)x + 0.5;
		inCoords->y = (double)y + 0.5;
		(void)psPlaneTransformApply(outCoords, map, inCoords);
		mask->data.F32[y][x] = psImagePixelInterpolate(reject, outCoords->x, outCoords->y,
							       NULL, 0, 0.0, PS_INTERPOLATE_BILINEAR);

	    }
	}

#ifdef TESTING
	// Write error image out to check
	char maskfile[MAXCHAR];	// Filename of mask image
	sprintf(maskfile,"%s.mask",config->inputs->data[i]);
	psImageWriteSection(mask,0,0,0,NULL,0,maskfile);
#endif

	// Clip the image, and convert to suitable mask format
	(void)psImageClip(mask, config->frac, 0.0, config->frac, 1.0);
	psImage *maskU8 = psImageCopy(NULL, mask, PS_TYPE_U8);
	psFree(mask);
    
	// Stuff into the array
	inputRej->data[i] = maskU8;

    }

    psFree(inCoords);
    psFree(outCoords);

    return inputRej;
}
