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

#include <stdio.h>
#include <assert.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, 1 = don't use, 0 = use
    )
{
#if 0
    // Would like to use psVectorStats, but it doesn't have errors built in yet
    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN);
    (void)psVectorStats(stats, values, NULL, masks, 1);
    float mean = stats->sampleMean;
    psFree(stats);
    return mean;
#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]) {
            // "error" here is the variance
            sum += values->data.F32[i] / errors->data.F32[i];
            weights += 1.0 / 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, NULL, masks, 1);
    float median = stats->sampleMedian;
    psFree(stats);
    return median;
}



bool stacCombine(psImage **combined,    // The combined image for output
                 psArray **rejected,    // Array of rejection masks
                 psArray *images,       // Array of transformed images
                 psArray *errors,       // Array of transformed error images
                 int nReject,           // Number of rejection iterations
                 psImage *region,       // Region to combine
                 psVector *saturated,   // Saturation limits for each image
                 psVector *bad,         // Bad pixel limits for each image
                 float reject           // Rejection (k-sigma)
    )
{
    assert(combined);
    assert(images);
    assert(errors);
    assert(images->n == errors->n);
    assert(saturated);
    assert(bad);
    assert(saturated->n == images->n);
    assert(bad->n == images->n);

    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

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

        assert(image->numCols == numCols && image->numRows == numRows);
        assert(error->numCols == numCols && error->numRows == numRows);
    }

    // Check combined image
    assert(!*combined || ((*combined)->numRows == numRows && (*combined)->numCols == numCols));
    if (*combined == NULL) {
        *combined = psImageAlloc(numCols, numRows, PS_TYPE_F32); // Combined image
    }

    // Check area of interest
    assert(!region || (region->numRows == numRows && region->numCols == numCols));

    psTrace("stac.combine", 1, "Combining images....\n");

    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); // Will hold the mask in the statistics step

    // Set up rejection masks
    if (nReject > 0) {
        if (*rejected == NULL) {
            // Allocate the rejection masks, if required
            *rejected = psArrayAlloc(nImages);
        } else {
            assert((*rejected)->n != nImages);
        }

        // Create and initialise rejection masks
        for (int i = 0; i < nImages; i++) {
            (*rejected)->data[i] = psImageAlloc(numCols, numRows, PS_TYPE_U8);
            for (int r = 0; r < numRows; r++) {
                for (int c = 0; c < numCols; c++) {
                    ((psImage*)((*rejected)->data[i]))->data.U8[r][c] = 0;
                }
            }
        }
    }

#ifdef TESTING
    // chi^2 image
    psImage *chi2 = psImageAlloc(numCols, numRows, PS_TYPE_F32);
    static int iteration = 0;           // Number of times function has been called
#endif

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {

            // Only combine those pixels requested
            if (!region || (region && region->data.U8[y][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]; // This is the variance!
                    pixels->data.F32[i] = pixel;
                    deltas->data.F32[i] = delta;
                    if ((pixel >= saturated->data.F32[i]) || (pixel <= bad->data.F32[i]) ||
                        (! isfinite(pixel)) || (! isfinite(delta))) {
                        mask->data.U8[i] = (psU8)1; // Don't use!
                    } else {
                        mask->data.U8[i] = (psU8)0; // Use.
                    }
                }

                float average = stacCombineMean(pixels, deltas, mask); // Combined value

                // We set the value BEFORE the rejection iteration because not all pixels that we reject
                // here will be rejected in the final cut.
                (*combined)->data.F32[y][x] = average;

#ifdef TESTING
                // Calculate chi^2
                chi2->data.F32[y][x] = 0.0;
                int numGoodPix = 0;
                for (int i = 0; i < nImages; i++) {
                    if (! mask->data.U8[i]) {
                        chi2->data.F32[y][x] += SQUARE(pixels->data.F32[i] - average) / deltas->data.F32[i];
                        numGoodPix++;
                    }
                }
                chi2->data.F32[y][x] /= (float)numGoodPix;
#endif

                // Rejection iterations
                bool keepGoing = true;  // Keep going with rejection?
                for (int rejNum = 0; (rejNum < nReject) && keepGoing; rejNum++) {
                    float max = 0.0;    // Maximum deviation
                    int maxIndex = -1;  // Index of the maximum deviation
                    for (int i = 0; i < nImages; i++) {
                        if (!mask->data.U8[i] &&
                            ((pixels->data.F32[i] - average) / sqrtf(deltas->data.F32[i]) > max)) {
                            max = (pixels->data.F32[i] - average) / sqrtf(deltas->data.F32[i]);
                            maxIndex = i;
                        }
                    }
                    // Reject the pixel with the maximum deviation
                    if (max > reject) {
                        mask->data.U8[maxIndex] = 1;
                        ((psImage*)((*rejected)->data[maxIndex]))->data.U8[y][x] += 1;
                        // Re-do combination following rejection
                        average = stacCombineMean(pixels, deltas, mask);
                    } else {
                        keepGoing = false;
                    }
                }

            } // Pixels of interest

        }
    } // Iterating over output pixels

#ifdef TESTING
    // Write chi^2 image
    iteration++;
    char chiName[MAXCHAR];              // Filename of chi^2 image
    sprintf(chiName,"chi2_%d.fits",iteration);
    psFits *chiFile = psFitsOpen(chiName, "w");
    if (!psFitsWriteImage(chiFile, NULL, chi2 , 0, NULL)) {
        psErrorStackPrint(stderr, "Unable to write image: %s\n", chiName);
    }
    psTrace("stac.combine", 1, "Chi^2 image written to %s\n", chiName);
    psFitsClose(chiFile);
    psFree(chi2);
#endif

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

    return combined;
}
