/** @file  pmStack.c
 *
 *  This file will perform image combination of several images of the
 *  same field, produce a list of questionable pixels, then tag some
 *  of those pixels as cosmic rays.
 *
 *  @author Paul Price, IfA
 *  @author GLG, MHPCC
 *
 *  @version $Revision: 1.5 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2007-05-25 03:43:22 $
 *
 *  Copyright 2004-2007 Institute for Astronomy, University of Hawaii
 *
 */

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

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

#include "pmStack.h"

#define PIXEL_LIST_BUFFER 100           // Number of entries to add to pixel list at a time
#define PIXEL_MAP_BUFFER 2              // Number of entries to add to pixel map at a time

// Internal values for masks
#define MASK_BAD 0x01                   // Mask value for pixels that have formerly been masked as bad
#define MASK_SUSPECT 0x02               // Mask value for pixels that are suspected bad (by pmStackCombine)


// Data structure for use as a buffer when combining pixels
// Use of this structure means we don't have to do an allocation in the combination function for each pixel
typedef struct {
    psVector *pixels;                   // Pixel values
    psVector *masks;                    // Pixel masks
    psStats *stats;                     // Statistics
} combineBuffer;

static void combineBufferFree(combineBuffer *buffer)
{
    psFree(buffer->pixels);
    psFree(buffer->masks);
    psFree(buffer->stats);
    return;
}

static combineBuffer *combineBufferAlloc(long numImages, // Number of images that will be combined
                                         psStatsOptions stat // Statistic to use
    )
{
    combineBuffer *buffer = psAlloc(sizeof(combineBuffer));
    psMemSetDeallocator(buffer, (psFreeFunc)combineBufferFree);

    buffer->pixels = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->masks = psVectorAlloc(numImages, PS_TYPE_MASK);
    buffer->stats = psStatsAlloc(stat);

    return buffer;
}


// Deallocator for the stack data
static void stackDataFree(pmStackData *data)
{
    psFree(data->detector);
    psFree(data->sky);
    psFree(data->pixels);
    return;
}


// Given a stack of images, combine with optional rejection.
// Pixels in the stack that are rejected are marked for subsequent
static bool combinePixels(psImage *image, // Combined image, for output
                          psImage *mask, // Combined mask, for output
                          const psArray *inputs, // Stack data
                          const psVector *weights, // Weights for data, or NULL
                          const psVector *reject, // Indices of pixels to reject, or NULL
                          int x, int y, // Coordinates of interest
                          psMaskType maskVal, // Value to mask
                          psMaskType bad, // Value to give bad pixels
                          int numIter, // Number of rejection iterations
                          float rej, // Number of standard deviations at which to reject
                          combineBuffer *buffer // Buffer for combination; to avoid multiple allocations
                         )
{
    // Rudimentary error checking
    assert(image);
    assert(mask);
    assert(inputs);
    assert(numIter >= 0);
    assert(rej > 0);
    assert(buffer);

    int num = inputs->n;          // Number of images to combine

    psVector *pixelData = buffer->pixels; // Values for the pixel of interest
    psVector *pixelMasks = buffer->masks; // Masks for the pixel of interest
    psStats *stats = buffer->stats;     // Statistics


    // Extract the pixel and mask data
    int numBad = 0;                     // Number of bad pixels
    for (int i = 0; i < num; i++) {
        pmStackData *data = inputs->data[i]; // Stack data of interest
        psImage *image = data->sky->image; // Image of interest
        psImage *mask = data->sky->mask; // Mask of interest
        pixelData->data.F32[i] = image->data.F32[y][x];
        pixelMasks->data.PS_TYPE_MASK_DATA[i] = mask->data.PS_TYPE_MASK_DATA[y][x];
        if (pixelMasks->data.PS_TYPE_MASK_DATA[i] & maskVal) {
            numBad++;
        }
    }
    if (numBad == num) {
        // Catch the case where everything is masked
        image->data.F32[y][x] = NAN;
        mask->data.PS_TYPE_MASK_DATA[y][x] = bad;
        return false;
    }

    if (reject) {
        for (int i = 0; i < reject->n; i++) {
            pixelMasks->data.PS_TYPE_MASK_DATA[reject->data.U16[i]] |= maskVal;
        }
    }

    // Record the value derived with no clipping, because pixels rejected using the harsh clipping applied in
    // the first pass might later be accepted.
    stats->options |= PS_STAT_SAMPLE_MEAN;
    if (!psVectorStats(stats, pixelData, weights, pixelMasks, maskVal)) {
        // Can't do anything about an error except give it a NAN and mask
        psErrorClear();
        image->data.F32[y][x] = NAN;
        mask->data.PS_TYPE_MASK_DATA[y][x] = bad;
        return false;
    }
    image->data.F32[y][x] = stats->sampleMean;
    mask->data.PS_TYPE_MASK_DATA[y][x] = 0;

    stats->options &= ~ PS_STAT_SAMPLE_MEAN;

    long numClipped = LONG_MAX;         // Number of pixels clipped
    psMaskType ignore = maskVal | bad;  // Ignore values with this mask value
    for (int i = 0; i < numIter && numClipped > 0; i++) {
#if 1
        float mean = stats->sampleMedian;
        float stdev = 0.74 * (stats->sampleUQ - stats->sampleLQ); // Rough estimate of the standard deviation
#else
        float mean = stats->robustMedian;
        float stdev = stats->robustStdev;
#endif
        float limit = rej * stdev; // Rejection limit
        numClipped = 0;
        for (int j = 0; j < num; j++) {
            if (pixelMasks->data.PS_TYPE_MASK_DATA[j] & ignore) {
                continue;
            }
            float diff = pixelData->data.F32[j] - mean;
            if (fabsf(diff) > limit) {
                pmStackData *data = inputs->data[j]; // Stack data of interest
                data->pixels = psPixelsAdd(data->pixels, PIXEL_LIST_BUFFER, x, y);
                pixelMasks->data.PS_TYPE_MASK_DATA[j] |= bad;
                numClipped++;
            }
        }
        if (!psVectorStats(stats, pixelData, weights, pixelMasks, maskVal)) {
            // Can't do anything about an error except give it a NAN and mask
            psErrorClear();
            image->data.F32[y][x] = NAN;
            mask->data.PS_TYPE_MASK_DATA[y][x] = bad;
            return false;
        }
    }

    return true;
}


// Ensure the input array of pmStackData is valid, and get some details out of it
static bool validateInputData(bool *haveDetector, // Do we have the detector images?
                           bool *haveSky, // Do we have the sky images?
                           bool *havePixels, // Do we have lists of pixels?
                           int *num,    // Number of inputs
                           int *numCols, int *numRows, // Size of (sky) images
                           psArray *input // Input array of pmStackData to validate
    )
{
    PS_ASSERT_ARRAY_NON_NULL(input, false);
    *num = input->n;
    {
        pmStackData *data = input->data[0]; // First image off the rank
        PS_ASSERT_PTR_NON_NULL(data, false);
        assert(psMemGetDeallocator(data) == (psFreeFunc)stackDataFree); // Ensure it's the right type
        *haveDetector = (data->detector != NULL);
        if (*haveDetector) {
            PS_ASSERT_IMAGE_NON_NULL(data->detector->image, false);
            PS_ASSERT_IMAGE_NON_NULL(data->detector->mask, false);
        }
        *haveSky = (data->sky != NULL);
        if (*haveSky) {
            PS_ASSERT_IMAGE_NON_NULL(data->sky->image, false);
            PS_ASSERT_IMAGE_NON_NULL(data->sky->mask, false);
            *numCols = data->sky->image->numCols;
            *numRows = data->sky->image->numRows;
        }
        *havePixels = (data->pixels != NULL);
    }
    for (int i = 1; i < *num; i++) {
        pmStackData *data = input->data[i]; // Stack data for this input
        assert(psMemGetDeallocator(data) == (psFreeFunc)stackDataFree); // Ensure it's the right type
        if ((*haveDetector && !data->detector) || (data->detector && !*haveDetector)) {
            psError(PS_ERR_UNEXPECTED_NULL, true,
                    "The detector readout is specified in some but not all inputs.");
            return false;
        }
        if ((*haveSky && !data->sky) || (data->sky && !*haveSky)) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "The sky cell is specified in some but not all inputs.");
            return false;
        }
        if ((*havePixels && !data->pixels) || (data->pixels && !*havePixels)) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "The pixels are specified in some but not all inputs.");
            return false;
        }
        if (*haveDetector) {
            PS_ASSERT_IMAGE_NON_NULL(data->detector->image, false);
            PS_ASSERT_IMAGE_NON_NULL(data->detector->mask, false);
            PS_ASSERT_IMAGE_TYPE(data->detector->image, PS_TYPE_F32, false);
            PS_ASSERT_IMAGE_TYPE(data->detector->mask, PS_TYPE_MASK, false);
            PS_ASSERT_IMAGES_SIZE_EQUAL(data->detector->image, data->detector->mask, false);
        }
        if (*haveSky) {
            PS_ASSERT_IMAGE_NON_NULL(data->sky->image, false);
            PS_ASSERT_IMAGE_NON_NULL(data->sky->mask, false);
            PS_ASSERT_IMAGE_TYPE(data->sky->image, PS_TYPE_F32, false);
            PS_ASSERT_IMAGE_TYPE(data->sky->mask, PS_TYPE_MASK, false);
            PS_ASSERT_IMAGE_SIZE(data->sky->image, *numCols, *numRows, false);
            PS_ASSERT_IMAGES_SIZE_EQUAL(data->sky->image, data->sky->mask, false);
        }
    }

    return true;
}



// Examine a pixel carefully to see if it should be rejected in the stack by convolving all the input images
// to the same seeing, and clipping there.
static bool convolveTest(psVector *reject, // Image indices to reject for this pixel (output)
                         psVector *inspect, // Images indices to inspect for this pixel (input)
                         psArray *input, // The input pmStackData
                         int xPix, int yPix,   // Pixel coordinates of interest
                         psVector *seeing, // Seeing convolution for each image
                         psMaskType maskVal, // Value to mask
                         float extent,   // Extent of convolution (Gaussian sigmas)
                         float threshold// Rejection threshold (standard deviations)
    )
{
    assert(reject && reject->type.type == PS_TYPE_U16);
    assert(input);
    assert(seeing);
    assert(seeing->n == input->n);

    int num = input->n;                // Number of input images
    psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV); // Statistics
    psVector *values = psVectorAlloc(num, PS_TYPE_F32); // Values from convolving the image
    psVector *valuesMask = psVectorAlloc(num, PS_TYPE_MASK); // Mask for convolution values

    for (int i = 0; i < num; i++) {
        pmStackData *data = input->data[i]; // Stacking data
        int radius = extent * seeing->data.F32[i]; // How much to convolve
        psImage *image = data->sky->image; // Image to convolve
        psImage *mask = data->sky->mask; // Image mask

        int xMin = PS_MAX(xPix - radius, 0);
        int xMax = PS_MIN(xPix + radius, image->numCols - 1);
        int yMin = PS_MAX(yPix - radius, 0);
        int yMax = PS_MIN(yPix + radius, image->numRows - 1);

        float sum = 0.0;            // Sum of pixels in convolution
        float sumKernel = 0.0;      // Sum of kernel
        for (int y = yMin; y <= yMax; y++) {
            int v2 = PS_SQR(y - yPix);
            for (int x = xMin; x <= xMax; x++) {
                int u2 = PS_SQR(x - xPix);
                float kernel = expf(-(u2 + v2) / seeing->data.F32[i]);
                if (!(mask->data.PS_TYPE_MASK_DATA[y][x] & maskVal)) {
                    sum += image->data.F32[y][x] * kernel;
                    sumKernel += kernel;
                }
            }
        }

        if (sumKernel != 0.0) {
            values->data.F32[i] = sum / sumKernel;
            valuesMask->data.PS_TYPE_MASK_DATA[i] = 0;
        } else {
            // Everything's masked out; don't include this pixel in the analysis
            values->data.F32[i] = NAN;
            valuesMask->data.PS_TYPE_MASK_DATA[i] = MASK_BAD;
        }
    }

    // Mask the ones we're inspecting
    for (int i = 0; i < inspect->n; i++) {
        valuesMask->data.PS_TYPE_MASK_DATA[inspect->data.U16[i]] |= MASK_SUSPECT;
    }

    if (!psVectorStats(stats, values, NULL, valuesMask, MASK_BAD | MASK_SUSPECT)) {
        psFree(stats);
        psFree(values);
        psFree(valuesMask);
        return false;
    }

    // Inspect the pixels of interest
    float mean = stats->sampleMean;       // Mean value
    float limit = threshold * stats->sampleStdev; // Rejection threshold
    for (int j = 0; j < inspect->n; j++) {
        int index = inspect->data.U16[j]; // Index of interest
        if (fabsf(values->data.F32[index] - mean) > limit) {
            reject->data.U16[reject->n++] = index;
        }
    }

    psFree(stats);
    psFree(values);
    psFree(valuesMask);

    return true;
}



// Generate a "pixel map".
//
// A "pixel map" is an image-like structure containing a vector that contains the indices of images.  The idea
// is to provide a reverse lookup for an array of pixel lists, so that the image for which a pixel is flagged
// can be identified easily.
static psArray *pixelMapGenerate(const psArray *input, // Data to stack
                                 int numCols, int numRows // Size of (sky) images
    )
{
    psArray *map = psArrayAlloc(numRows); // The pixel map
    for (int y = 0; y < numRows; y++) {
        map->data[y] = psArrayAlloc(numCols);
    }

    for (int i = 0; i < input->n; i++) {
        pmStackData *data = input->data[i];
        assert(data && data->pixels);
        psPixels *pixels = data->pixels;// The pixels of interest
        for (int j = 0; j < pixels->n; j++) {
            int x = pixels->data[j].x, y = pixels->data[j].y; // Coordinates of interest
            psArray *columns = map->data[y]; // The columns for that row
            psVector *images = columns->data[x]; // The images for that column
            if (!images) {
                images = columns->data[x] = psVectorAllocEmpty(PIXEL_MAP_BUFFER, PS_TYPE_U16);
            }
            int size = images->n;       // Element number at which to add
            columns->data[x] = psVectorExtend(images, PIXEL_MAP_BUFFER, 1);
            images->data.U16[size] = i;
        }
    }

    return map;
}

// Query a "pixel map", by returning the list of image indices for a particular pixel.
static psVector *pixelMapQuery(const psArray *map, // Pixel map
                               int x, int y // Coordinates of interest
    )
{
    assert(y >= 0 && y < map->n);
    psArray *colMap = map->data[y];     // Columns for that row
    assert(x >= 0 && x < colMap->n);
    return colMap->data[x];
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

/// Constructor
pmStackData *pmStackDataAlloc(pmReadout *sky, float seeing, float weight)
{
    pmStackData *data = psAlloc(sizeof(pmStackData)); // Stack data, to return
    psMemSetDeallocator(data, (psFreeFunc)stackDataFree);

    data->detector = NULL;

    data->sky = psMemIncrRefCounter(sky);
    data->pixels = NULL;
    data->seeing = seeing;
    data->weight = weight;

    return data;
}

/// Stack input images
bool pmStackCombine(pmReadout *combined, psArray *input, psMaskType maskVal, psMaskType bad,
                    int numIter, float rej)
{
    PS_ASSERT_PTR_NON_NULL(combined, false);
    bool haveDetector;                  // Do we have the detector images?
    bool haveSky;                       // Do we have the sky images?
    bool havePixels;                    // Do we have lists of pixels?
    int num;                            // Number of inputs
    int numCols, numRows;               // Size of (sky) images
    if (!validateInputData(&haveDetector, &haveSky, &havePixels, &num, &numCols, &numRows, input)) {
        return false;
    }
    PS_ASSERT_INT_POSITIVE(bad, false);
    PS_ASSERT_INT_NONNEGATIVE(numIter, false);
    PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    if (havePixels) {
        // This is a subsequent combination, so expect that the image and mask already exist
        PS_ASSERT_IMAGE_NON_NULL(combined->image, false);
        PS_ASSERT_IMAGE_TYPE(combined->image, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_NON_NULL(combined->mask, false);
        PS_ASSERT_IMAGE_TYPE(combined->mask, PS_TYPE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(combined->image, combined->mask, false);
    }
    if (!haveDetector && !haveSky) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Nothing to combine!");
        return false;
    }


    if (!haveSky) {
        // Need to generate the sky cell images

        // ...

#if 0
        numCols = data->sky->image->numCols;
        numRows = data->sky->image->numRows;
#endif
    }

    // Pull out the weightings
    psVector *weights = psVectorAlloc(num, PS_TYPE_F32);
    for (int i = 0; i < num; i++) {
        pmStackData *data = input->data[i]; // Stack data for this input
        weights->data.F32[i] = data->weight;
    }

    // Buffer for combination
    combineBuffer *buffer = combineBufferAlloc(num, numIter == 0 ? PS_STAT_SAMPLE_MEAN :
                                               PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_MEDIAN);

    if (havePixels) {
        // Only combine select pixels
        psImage *combinedImage = combined->image; // Combined image
        psImage *combinedMask = combined->mask; // Combined mask

        psArray *pixelMap = pixelMapGenerate(input, numCols, numRows); // Map of pixels to source
        psPixels *pixels = NULL;            // Total list of pixels, with no duplicates
        for (int i = 0; i < num; i++) {
            pmStackData *data = input->data[i]; // Stacking data; contains the list of pixels
            pixels = psPixelsConcatenate(pixels, data->pixels);
            data->pixels = psPixelsRealloc(data->pixels, PIXEL_LIST_BUFFER); // Just in case more rejection
        }
        for (int i = 0; i < pixels->n; i++) {
            int x = pixels->data[i].x, y = pixels->data[i].y; // Coordinates of interest
            psVector *reject = pixelMapQuery(pixelMap, x, y); // Inspect these images closely
            combinePixels(combinedImage, combinedMask, input, weights, reject, x, y,
                          maskVal, bad, numIter, rej, buffer);
        }
        psTrace("psModules.imcombine", 5, "Additional %ld pixels fixed.\n", pixels->n);
        psFree(pixels);
        psFree(pixelMap);
    } else {
        // Pull the products out, allocate if necessary
        psImage *combinedImage = combined->image; // Combined image
        if (!combinedImage) {
            combined->image = psImageAlloc(numCols, numRows, PS_TYPE_F32);
            combinedImage = combined->image;
        }
        psImage *combinedMask = combined->mask; // Combined mask
        if (!combinedMask) {
            combined->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
            combinedMask = combined->mask;
        }

        // Generate the pixel lists in which to place the rejected pixels
        if (numIter != 0) {
            for (int i = 0; i < num; i++) {
                pmStackData *data = input->data[i]; // Stack data for this input
                data->pixels = psPixelsAllocEmpty(PIXEL_LIST_BUFFER);
            }
        }

        for (int y = 0; y < numRows; y++) {
            for (int x = 0; x < numCols; x++) {
                combinePixels(combinedImage, combinedMask, input, weights, NULL, x, y,
                              maskVal, bad, numIter, rej, buffer);
            }
        }

        if (psTraceGetLevel("psModules.imcombine") >= 5) {
            for (int i = 0; i < num; i++) {
                pmStackData *data = input->data[i]; // Stack data for this input
                psTrace("psModules.imcombine", 5, "Image %d: %ld pixels to inspect.\n", i, data->pixels->n);
            }
        }
    }

    psFree(weights);
    psFree(buffer);

    return true;
}

bool pmStackReject(psArray *input, psMaskType maskVal, float extent, float threshold)
{
    bool haveDetector;                  // Do we have the detector images?
    bool haveSky;                       // Do we have the sky images?
    bool havePixels;                    // Do we have lists of pixels?
    int num;                            // Number of inputs
    int numCols, numRows;               // Size of (sky) images
    if (!validateInputData(&haveDetector, &haveSky, &havePixels, &num, &numCols, &numRows, input)) {
        return false;
    }
    if (!havePixels) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "No pixels have been flagged for possible rejection.");
        return false;
    }
    PS_ASSERT_FLOAT_LARGER_THAN(extent, 0.0, false);

    // Want to concatenate the list of pixels, but also want to know where each came from
    psArray *pixelMap = pixelMapGenerate(input, numCols, numRows); // Map of pixels to source
    psPixels *pixels = NULL;            // Total list of pixels, with no duplicates
    for (int i = 0; i < num; i++) {
        pmStackData *data = input->data[i]; // Stacking data; contains the list of pixels
        pixels = psPixelsConcatenate(pixels, data->pixels);
        // Throw out the old in preparation for the new
        data->pixels = psPixelsRealloc(data->pixels, 0);
    }

    // Get the convolution widths
    psVector *seeing = psVectorAlloc(num, PS_TYPE_F32); // Seeing for each image
    const float fwhm2sigma = 1.0 / (2.0*sqrt((double)2.0*(log((double)2.0)))); // Conversion FWHM --> sigma
    float seeingMax = -INFINITY;                // Maximum FWHM
    for (int i = 0; i < num; i++) {
        pmStackData *data = input->data[i]; // Stacking data
        seeing->data.F32[i] = data->seeing * fwhm2sigma;
        if (data->seeing > seeingMax) {
            seeingMax = data->seeing;
        }
    }
    (void)psBinaryOp(seeing, seeing, "*", seeing);
    (void)psBinaryOp(seeing, psScalarAlloc(PS_SQR(seeingMax), PS_TYPE_F32), "-", seeing);
    (void)psUnaryOp(seeing, seeing, "sqrt");


    psVector *reject = psVectorAllocEmpty(num, PS_TYPE_U16); // Image indices to reject for a given pixel

    for (int i = 0; i < pixels->n; i++) {
        int x = pixels->data[i].x;      // x coordinate of interest
        int y = pixels->data[i].y;      // y coordinate of interest
        psVector *inspect = pixelMapQuery(pixelMap, x, y); // Inspect these images closely

        // Can daisy-chain multiple tests here
        reject->n = 0;

        if (!convolveTest(reject, inspect, input, x, y, seeing, maskVal, extent, threshold)) {
            // Some sort of error message here; NAN and mask all the pixels
            for (int i = 0; i < num; i++) {
                pmStackData *data = input->data[i]; // Stacking data
                psPixelsAdd(data->pixels, PIXEL_LIST_BUFFER, x, y);
            }
            continue;
        }

        // Add the rejected pixels to the official rejection list
        for (int i = 0; i < reject->n; i++) {
            pmStackData *data = input->data[reject->data.U16[i]]; // Stacking data
            psPixelsAdd(data->pixels, PIXEL_LIST_BUFFER, x, y);
        }
    }
    psFree(reject);
    psFree(seeing);
    psFree(pixels);
    psFree(pixelMap);

    if (psTraceGetLevel("psModules.imcombine") >= 5) {
        for (int i = 0; i < num; i++) {
            pmStackData *data = input->data[i]; // Stack data for this input
            psTrace("psModules.imcombine", 5, "Image %d: %ld pixels rejected.\n", i, data->pixels->n);
        }
    }

    return true;
}

