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

#define SQUARE(x) ((x)*(x))
#define ABS(x) ((x) >= 0 ? (x) : -(x))

float stacCombineMean(psVector *values,	// Values for which to take the mean
		      psVector *errors,	// Errors in the values
		      psVector *masks	// Masks for the values, 0 = don't use, 1 = use
    )
{
#if 0
    // Would like to use psVectorStats, but it doesn't have errors built in yet
    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);
    (void)psVectorStats(stats, values, masks, 1);
    psFree(stats);
    return stats->sampleMean;
#else
    // Instead, do it ourselves
    double sum = 0.0;
    double weights = 0.0;
    int num = values->n;
    for (int i = 0; i < num; i++) {
	if (masks->data.U8[i]) {
	    sum += values->data.F32[i] / SQUARE(errors->data.F32[i]);
	    weights += 1.0 / SQUARE(errors->data.F32[i]);
	}
    }
    if (weights > 0.0) {
	return (float)(sum/weights);
    } else {
	return NAN;
    }
#endif

}


float stacCombineMedian(psVector *values, // Values for which to take the median
			psVector *errors, // Errors in the values
			psVector *masks	// Masks for the values, 0 = don't use, 1 = use
    )
{
    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEDIAN);
    (void)psVectorStats(stats, values, masks, 1);
    psFree(stats);
    return stats->sampleMedian;
}


psImage *stacCombine(psArray *images,	// Array of transformed images
		     psArray *errors,	// Array of transformed error images
		     int nReject,	// Number of rejection iterations
		     psArray **rejected, // Array of rejection masks
		     stacConfig *config	// Configuration
    )
{
    int nImages = images->n;		// Number of images
    int numRows = ((psImage*)images->data[0])->numRows;	// Image size
    int numCols = ((psImage*)images->data[0])->numCols; // Image size
    float saturated = config->saturated;// Saturation limit
    float bad = config->bad;		// Bad pixel limit
    float reject = config->reject;	// Rejection (k-sigma)
    
    // Check dimensions for consistency
    for (int i = 0; i < nImages; i++) {
	psImage *image = (psImage *)images->data[i]; // The image
	psImage *error = (psImage *)errors->data[i]; // The error image

	if ((image->numCols != numCols) || (image->numRows != numRows)) {
	    psError("stac.combine",
		    "Image size mismatch on error image %d\nExpected %dx%d, got %dx%d\n",
		    i, numCols, numRows, image->numCols, image->numRows);
	    return NULL;
	}
	if ((error->numCols != numCols) || (error->numRows != numRows)) {
	    psError("stac.combine",
		    "Image size mismatch on error image %d\nExpected %dx%d, got %dx%d\n",
		    i, numCols, numRows, error->numCols, error->numRows);
	    return NULL;
	}
    }

    psTrace("stac.combine", 1, "Combining images....\n");
    psTrace("stac.combine", 3, "Saturation: %f Bad: %f\n", saturated, bad);

    psVector *pixels = psVectorAlloc(nImages, PS_TYPE_F32); // Will hold the pixels in the statistics step
    psVector *deltas = psVectorAlloc(nImages, PS_TYPE_F32); // Will hold the errors in the statistics step
    psVector *mask = psVectorAlloc(nImages, PS_TYPE_U8); // Mask bad pixels
    psImage *combined = psImageAlloc(numCols, numRows, PS_TYPE_F32); // Combined image


    // Set up rejection masks
    if (nReject > 0) {
	if (*rejected == NULL) {
	    // Allocate the rejection masks, if required
	    *rejected = psArrayAlloc(nImages);
	} else if ((*rejected)->n != nImages) {
	    psError("stac.combine", "Number of rejection masks (%d) does not match number of input"
		    "images (%d).\n", (*rejected)->n, nImages);
	    return NULL;
	}
	for (int i = 0; i < nImages; i++) {
	    (*rejected)->data[i] = psImageAlloc(numCols, numRows, PS_TYPE_F32); // Create rejection mask
	    // Zero out the mask
	    for (int r = 0; r < numRows; r++) {
		for (int c = 0; c < numCols; c++) {
		    ((psImage*)((*rejected)->data[i]))->data.F32[r][c] = 0.0;
		}
	    }
	}
    }
    
    
    for (int y = 0; y < numRows; y++) {
	for (int x = 0; x < numCols; x++) {
	    
	    // Export pixels into the vector and get stats
	    for (int i = 0; i < nImages; i++) {
		float pixel = ((psImage*)images->data[i])->data.F32[y][x];
		float delta = ((psImage*)errors->data[i])->data.F32[y][x];
		pixels->data.F32[i] = pixel;
		deltas->data.F32[i] = delta;
		if ((pixel >= saturated) || (pixel <= bad) || (! isfinite(pixel)) || (! isfinite(delta))) {
		    mask->data.U8[i] = (psU8)0; // Don't use!
		} else {
		    mask->data.U8[i] = (psU8)1; // Use.
		}
	    }
	    
	    float average = stacCombineMean(pixels, deltas, mask); // Combined value
	    
	    // Rejection iterations
	    for (int rejNum = 0; rejNum < nReject; rejNum++) {
		float max = 0.0;
		int maxIndex = 0;
		for (int i = 0; i < nImages; i++) {
		    if (mask->data.U8[i] && (ABS(pixels->data.F32[i] - average) / deltas->data.F32[i] > max)) {
			max = ABS(pixels->data.F32[i] - average) / deltas->data.F32[i];
			maxIndex = i;
		    }
		}
		// Reject the pixel with the maximum deviation
		if (max > reject) {    
		    mask->data.U8[maxIndex] = 0;
		    ((psImage*)((*rejected)->data[maxIndex]))->data.F32[y][x] += 1.0;
		    // Re-do combination following rejection
		    average = stacCombineMean(pixels, deltas, mask);
		}
	    } // Rejection iterations
	    
	    combined->data.F32[y][x] = average;
	}
    } // Iterating over output pixels

#ifdef TESTING
    // Write rejection images out to check
    if (nReject > 0) {
	for (int i = 0; i < nImages; i++) {
	    char rejfile[MAXCHAR];		// Filename of rejection image
	    sprintf(rejfile,"%s.shiftrej",config->inputs->data[i]);
	    psImageWriteSection((psImage*)((*rejected)->data[i]),0,0,0,NULL,0,rejfile);
	}
    }
#endif


    // Clean up
    psFree(pixels);
    psFree(deltas);
    psFree(mask);

    return combined;
}
