/** @file  pmReadoutCombine.c
 *
 *  This file will contain a module which will combine multiple readout images.
 *
 *  @author GLG, MHPCC
 *
 *  @version $Revision: 1.18 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-06-29 00:30:36 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

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

#include "pmFPA.h"
#include "pmFPAMaskWeight.h"

#include "pmReadoutCombine.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// File-static (private) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Return the statistic of interest
static inline double getStat(const psStats *stats, // Statistics structure
                             psStatsOptions option // Statistics option
                            )
{
    assert(stats);

    switch (option) {
    case PS_STAT_SAMPLE_MEAN:
        return stats->sampleMean;
    case PS_STAT_SAMPLE_MEDIAN:
        return stats->sampleMedian;
    case PS_STAT_SAMPLE_STDEV:
        return stats->sampleStdev;
    case PS_STAT_ROBUST_MEDIAN:
        return stats->robustMedian;
    case PS_STAT_ROBUST_STDEV:
        return stats->robustStdev;
    case PS_STAT_FITTED_MEAN:
        return stats->fittedMean;
    case PS_STAT_FITTED_STDEV:
        return stats->fittedStdev;
    case PS_STAT_CLIPPED_MEAN:
        return stats->clippedMean;
    case PS_STAT_CLIPPED_STDEV:
        return stats->clippedStdev;
    case PS_STAT_MAX:
        return stats->max;
    case PS_STAT_MIN:
        return stats->min;
    default:
        psAbort(__func__, "Bad option: %x\n", option);
    }
    return NAN;
}

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

// Allocator for pmCombineParams
pmCombineParams *pmCombineParamsAlloc(psStatsOptions combine)
{
    pmCombineParams *params = psAlloc(sizeof(pmCombineParams));

    params->combine = combine;
    params->maskVal = 0;
    params->nKeep = 0;
    params->fracHigh = 0.0;
    params->fracHigh = 0.0;
    params->iter = 1;
    params->rej = 3.0;

    return params;
}


/******************************************************************************
XXX: Maybe add support for S16 and S32 types.  Currently, only F32 supported.
 *****************************************************************************/
bool pmReadoutCombine(pmReadout *output,
                      psArray *inputs,
                      const psVector *zero,
                      const psVector *scale,
                      pmCombineParams *params
                     )
{
    // Check inputs
    PS_ASSERT_PTR_NON_NULL(output, false);
    PS_ASSERT_PTR_NON_NULL(inputs, false);
    PS_ASSERT_PTR_NON_NULL(params, false);
    if (zero) {
        PS_ASSERT_VECTOR_TYPE(zero, PS_TYPE_F32, false);
        PS_ASSERT_VECTOR_SIZE(zero, inputs->n, false);
    }
    if (scale) {
        PS_ASSERT_VECTOR_TYPE(scale, PS_TYPE_F32, false);
        PS_ASSERT_VECTOR_SIZE(scale, inputs->n, false);
    }
    PS_ASSERT_FLOAT_WITHIN_RANGE(params->fracLow, 0.0, 1.0, false);
    PS_ASSERT_FLOAT_WITHIN_RANGE(params->fracHigh, 0.0, 1.0, false);
    if (params->combine != PS_STAT_SAMPLE_MEAN && params->combine != PS_STAT_SAMPLE_MEDIAN &&
            params->combine != PS_STAT_ROBUST_MEDIAN && params->combine != PS_STAT_FITTED_MEAN &&
            params->combine != PS_STAT_CLIPPED_MEAN) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Combination method is not SAMPLE_MEAN, SAMPLE_MEDIAN, "
                "ROBUST_MEDIAN, FITTED_MEAN or CLIPPED_MEAN.\n");
        return false;
    }

    psStats *stats = psStatsAlloc(params->combine); // The statistics to use in the combination
    if (params->combine == PS_STAT_CLIPPED_MEAN) {
        stats->clipSigma = params->rej;
        stats->clipIter = params->iter;
    }

    // Step through each readout in the input image list to determine how big of an output image is needed to
    // combine these input images.
    long maxInputCols = 0;              // The largest input column value
    long maxInputRows = 0;              // The largest input row value
    long minInputCols = LONG_MAX;       // The smallest input column value
    long minInputRows = LONG_MAX;       // The smallest input row value

    bool haveWeights = false;           // Do we have weight images?
    bool valid = false;                 // Do we have a single valid input?
    for (int i = 0; i < inputs->n; i++) {
        pmReadout *readout = inputs->data[i]; // Readout of interest
        if (!readout || !readout->image) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Input readout %d is NULL or has NULL image.\n", i);
            return false;
        }

        if (readout->weight) {
            if (valid && !haveWeights) {
                psLogMsg(__func__, PS_LOG_WARN, "Readout %d has a weight map, but others don't --- "
                         "weights ignored.\n", i);
            } else {
                haveWeights = true;
            }
        } else if (haveWeights) {
            psLogMsg(__func__, PS_LOG_WARN, "Readout %d doesn't have a weight map, but others do --- "
                     "weights ignored.\n", i);
            haveWeights = false;
        }
        valid = true;

        // Size of output image
        minInputRows = PS_MIN(minInputRows, readout->row0);
        maxInputRows = PS_MAX(maxInputRows, readout->row0 + readout->image->numRows);
        minInputCols = PS_MIN(minInputCols, readout->col0);
        maxInputCols = PS_MAX(maxInputCols, readout->col0 + readout->image->numCols);
        psTrace(__func__, 7, "Readout %d: offset %d,%d; size %dx%d\n", i,
                readout->col0, readout->row0, readout->image->numCols, readout->image->numRows);
    }
    if (!valid) {
        psError(PS_ERR_IO, true, "No valid inputs.\n");
        return NULL;
    }

    // Apply scale and zero
    if (scale || zero) {
        for (int i = 0; i < inputs->n; i++) {
            pmReadout *readout = inputs->data[i]; // The particular readout
            if (zero) {
                psBinaryOp(readout->image, readout->image, "-",
                           psScalarAlloc(zero->data.F32[i], PS_TYPE_F32));
            }
            if (scale) {
                float scaleFactor = 1.0 / scale->data.F32[i]; // Scaling
                psBinaryOp(readout->image, readout->image, "*", psScalarAlloc(scaleFactor, PS_TYPE_F32));
                if (haveWeights) {
                    psBinaryOp(readout->weight, readout->weight, "*",
                               psScalarAlloc(scaleFactor * scaleFactor, PS_TYPE_F32));
                }
            }
        }
    }

    // If there's existing images, we need to extend them
    long minOutputRows = (output->image || output->mask || output->weight) ?
                         PS_MIN(minInputRows, output->row0) : minInputRows;  // Smallest output row value
    long minOutputCols = (output->image || output->mask || output->weight) ?
                         PS_MIN(minInputCols, output->col0) : minInputCols; // Smallest output column value
    psTrace(__func__, 7, "Output minimum: %d,%d\n", minOutputCols, minOutputRows);

    // Update the origin
    *(psS32 *) &(output->col0) = minOutputCols;
    *(psS32 *) &(output->row0) = minOutputRows;

    // Generate the new output image by extending the current one, or making a whole new one
    psImage *newImage = psImageAlloc(maxInputCols - minOutputCols, maxInputRows - minOutputRows, PS_TYPE_F32);
    psImageInit(newImage, 0.0);
    if (output->image) {
        psImageOverlaySection(newImage, output->image, output->col0, output->row0, "=");
        psFree(output->image);
    }
    output->image = newImage;

    psImage *newMask = psImageAlloc(maxInputCols - minOutputCols, maxInputRows - minOutputRows, PS_TYPE_U8);
    psImageInit(newMask, 0);
    if (output->mask) {
        psImageOverlaySection(newMask, output->mask, output->col0, output->row0, "=");
        psFree(output->mask);
    }
    output->mask = newMask;

    psStatsOptions combineStdev = 0; // Statistics option for weights
    if (haveWeights) {
        psImage *newWeight = psImageAlloc(maxInputCols - minOutputCols, maxInputRows - minOutputRows,
                                          PS_TYPE_F32);
        psImageInit(newWeight, 0.0);
        if (output->weight) {
            psImageOverlaySection(newWeight, output->weight, output->col0, output->row0, "=");
            psFree(output->weight);
        }
        output->weight = newWeight;

        // Get the correct statistics option for weights
        switch (params->combine) {
        case PS_STAT_SAMPLE_MEAN:
        case PS_STAT_SAMPLE_MEDIAN:
            combineStdev = PS_STAT_SAMPLE_STDEV;
            break;
        case PS_STAT_ROBUST_MEDIAN:
            combineStdev = PS_STAT_ROBUST_STDEV;
            break;
        case PS_STAT_FITTED_MEAN:
            combineStdev = PS_STAT_FITTED_STDEV;
            break;
        case PS_STAT_CLIPPED_MEAN:
            combineStdev = PS_STAT_CLIPPED_STDEV;
            break;
        default:
            psAbort(__func__, "Should never get here --- checked params->combine before.\n");
        }
        stats->options |= combineStdev;
    }

    // We loop through each pixel in the output image.  We loop through each input readout.  We determine if
    // that output pixel is contained in the image from that readout.  If so, we save it in psVector pixels.
    // If not, we set a mask for that element in pixels.  Then, we mask off pixels not between fracLow and
    // fracHigh.  Then we call the vector stats routine on those pixels/mask.  Then we set the output pixel
    // value to the result of the stats call.

    psVector *pixels = psVectorAlloc(inputs->n, PS_TYPE_F32); // Stack of pixels
    pixels->n = inputs->n;
    psF32 *pixelsData = pixels->data.F32; // Dereference pixels

    psVector *mask   = psVectorAlloc(inputs->n, PS_TYPE_U8); // Mask for stack
    mask->n = inputs->n;
    psU8 *maskData = mask->data.U8;     // Dereference mask

    psVector *weights = NULL;           // Stack of weights
    psF32 *weightsData = NULL;          // Dereference weights
    if (haveWeights) {
        weights = psVectorAlloc(inputs->n, PS_TYPE_F32); // Stack of weights
        weights->n = inputs->n;
        weightsData = weights->data.F32;
    }
    psVector *index = NULL;             // The indices to sort the pixels

    float keepFrac = 1.0 - params->fracLow - params->fracHigh; // Fraction of pixels to keep
    psMaskType maskVal = params->maskVal; // The mask value

    psTrace(__func__, 3, "Iterating output: %d --> %d, %d --> %d\n",
            minInputCols - output->col0, maxInputCols - output->col0,
            minInputRows - output->row0, maxInputRows - output->row0);
    if (psTraceGetLevel(__func__) >= 3) {
        for (int r = 0; r < inputs->n; r++) {
            pmReadout *readout = inputs->data[r]; // Input readout
            psTrace(__func__, 3, "Iterating input %d: %d --> %d, %d --> %d\n", r,
                    minInputCols - readout->col0, maxInputCols - readout->col0,
                    minInputRows - readout->row0, maxInputRows - readout->row0);
        }
    }

    // Dereference output products
    psF32 **outputImage  = output->image->data.F32; // Output image
    psU8  **outputMask   = output->mask->data.U8; // Output mask
    psF32 **outputWeight = NULL; // Output weight map
    if (output->weight) {
        outputWeight = output->weight->data.F32;
    }

    for (int i = minInputRows; i < maxInputRows; i++) {
        int yOut = i - output->row0; // y position on output readout
        #if 0

        if (psTraceGetLevel(__func__) > 9) {
            printf("Processing row %d\r", i);
            fflush(stdout);
        }
        #endif
        for (int j = minInputCols; j < maxInputCols; j++) {
            int xOut = j - output->col0; // x position on output readout

            int numValid = 0;           // Number of valid pixels in the stack
            memset(maskData, 0, mask->n * sizeof(psU8)); // Reset the mask
            for (int r = 0; r < inputs->n; r++) {
                pmReadout *readout = inputs->data[r]; // Input readout
                int yIn = i - readout->row0; // y position on input readout
                int xIn = j - readout->col0; // x position on input readout

                psImage *image = readout->image; // The readout image

                // Check bounds
                if (xIn < 0 || xIn >= image->numCols || yIn < 0 || yIn >= image->numRows) {
                    continue;
                }

                pixelsData[r] = image->data.F32[yIn][xIn];

                // Check mask
                psImage *roMask = readout->mask; // The mask image
                if (roMask && roMask->data.U8[yIn][xIn] & maskVal) {
                    maskData[r] = 1;
                    continue;
                }

                numValid++;
            }

            if (numValid == 0) {
                outputMask[yOut][xOut] = PM_MASK_FLAT;
                outputImage[yOut][xOut] = NAN;
                continue;
            }

            // Apply fracLow,fracHigh if there are enough pixels
            if (numValid * keepFrac >= params->nKeep && keepFrac != 1.0) {
                index = psVectorSortIndex(index, pixels);
                int numLow = numValid * params->fracLow; // Number of low pixels to clip
                int numHigh = numValid * params->fracHigh; // Number of high pixels to clip
                // Low pixels
                psS32 *indexData = index->data.S32; // Dereference index
                for (int k = 0, numMasked = 0; numMasked < numLow && k < index->n; k++) {
                    // Don't count the ones that are already masked
                    if (!maskData[indexData[k]]) {
                        maskData[indexData[k]] = 1;
                        numMasked++;
                    }
                }
                // High pixels
                for (int k = pixels->n - 1, numMasked = 0; numMasked < numHigh && k >= 0; k--) {
                    // Don't count the ones that are already masked
                    if (! maskData[indexData[k]]) {
                        maskData[indexData[k]] = 1;
                        numMasked++;
                    }
                }
            }

            // Combination
            psVectorStats(stats, pixels, weights, mask, 1);
            outputImage[yOut][xOut] = getStat(stats, params->combine);
            if (haveWeights) {
                float stdev = getStat(stats, combineStdev);
                outputWeight[yOut][xOut] = PS_SQR(stdev); // Variance
            }
        }
    }
    #if 0
    if (psTraceGetLevel(__func__) > 9) {
        printf("\n");
    }
    #endif
    psFree(index);
    psFree(pixels);
    psFree(mask);
    psFree(weights);

    psFree(stats);

    return true;
}

