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

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

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmFPAMaskWeight.h"
#include "pmConceptsAverage.h"
#include "pmReadoutStack.h"

#include "pmReadoutCombine.h"

//#define SHOW_BUSY 1                   // Show that the function is busy


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

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

    params->combine = combine;
    params->maskVal = 0;
    params->blank = 0;
    params->nKeep = 0;
    params->fracHigh = 0.0;
    params->fracHigh = 0.0;
    params->iter = 1;
    params->rej = INFINITY;
    params->weights = false;

    return params;
}


// XXX: Maybe add support for S16 and S32 types.  Currently, only F32 supported.
bool pmReadoutCombine(pmReadout *output, const psArray *inputs, const psVector *zero, const psVector *scale,
                      const pmCombineParams *params)
{
    // Check inputs
    PS_ASSERT_PTR_NON_NULL(output, false);
    PS_ASSERT_ARRAY_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;
    }
    for (int i = 0; i < inputs->n; i++) {
        pmReadout *readout = inputs->data[i]; // Readout of interest
        if (params->weights && !readout->weight) {
            psError(PS_ERR_UNEXPECTED_NULL, true,
                    "Rejection based on weights requested, but no weights supplied for image %d.\n", i);
            return false;
        }
    }

    bool first = !output->image;        // First pass through?

    pmHDU *hdu = pmHDUFromReadout(output); // Output HDU
    if (!hdu) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find HDU for readout.\n");
        return false;
    }

    if (first) {
        psString comment = NULL;        // Comment to add to header
        psStringAppend(&comment, "Combining using statistic: %x", params->combine);
        if (!hdu->header) {
            hdu->header = psMetadataAlloc();
        }
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
        psFree(comment);
    }

    psStatsOptions combineStdev = 0; // 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("Should never get here --- checked params->combine before.\n");
    }

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

        if (first) {
            psString comment = NULL;    // Comment to add to header
            psStringAppend(&comment, "Combination clipping: %d iterations, rejection at %f sigma",
                           params->iter, params->rej);
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
            psFree(comment);
        }
    }
    if (params->weights && first) {
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                         "Using input weights to combine images", "");
    }

    int minInputCols, maxInputCols, minInputRows, maxInputRows; // Smallest and largest values to combine
    int xSize, ySize;                   // Size of the output image
    if (!pmReadoutStackValidate(&minInputCols, &maxInputCols, &minInputRows, &maxInputRows, &xSize, &ySize,
                                inputs)) {
        psError(PS_ERR_UNKNOWN, false, "No valid input readouts.");
        return false;
    }

    pmReadoutUpdateSize(output, minInputCols, minInputRows, xSize, ySize, true, params->weights,
                        params->blank);
    psTrace("psModules.imcombine", 7, "Output minimum: %d,%d\n", output->col0, output->row0);

    psImage *counts = pmReadoutAnalysisImage(output, PM_READOUT_STACK_ANALYSIS_COUNT, xSize, ySize,
                                             PS_TYPE_U16, 0);
    if (!counts) {
        return false;
    }
    psImage *sigma = pmReadoutAnalysisImage(output, PM_READOUT_STACK_ANALYSIS_SIGMA, xSize, ySize,
                                            PS_TYPE_F32, NAN);
    if (!sigma) {
        psFree(counts);
        return false;
    }

    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
    psF32 *pixelsData = pixels->data.F32; // Dereference pixels

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

    psVector *weights = NULL;           // Stack of weights
    psVector *errors = NULL;            // Stack of errors (sqrt of variance/weights), for psVectorStats
    psF32 *weightsData = NULL;          // Dereference weights
    if (params->weights) {
        weights = psVectorAlloc(inputs->n, PS_TYPE_F32); // Stack of weights
        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
    if (keepFrac != 1.0 && first) {
        psString comment = NULL;        // Comment to add to header
        psStringAppend(&comment, "Min/max rejection: %f high, %f low, keep %d",
                       params->fracHigh, params->fracLow, params->nKeep);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
        psFree(comment);
    }

    psMaskType maskVal = params->maskVal; // The mask value
    if (maskVal && first) {
        psString comment = NULL;        // Comment to add to header
        psStringAppend(&comment, "Mask for combination: %x", maskVal);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
        psFree(comment);
    }

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

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

    psVector *invScale = NULL;          // Inverse scale; pre-calculated for efficiency
    if (scale) {
        invScale = (psVector*)psBinaryOp(NULL, psScalarAlloc(1.0, PS_TYPE_F32), "/", (const psPtr)scale);
    }

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

        if (psTraceGetLevel("psModules.imcombine") > 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

                #if 0 // This should have been taken care of already:
                // Check bounds
                if (xIn < 0 || xIn >= image->numCols || yIn < 0 || yIn >= image->numRows) {
                    continue;
                }
                #endif

                pixelsData[r] = image->data.F32[yIn][xIn];
                if (!isfinite(pixelsData[r])) {
                    maskData[r] = 1;
                    continue;
                }

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

                if (params->weights) {
                    weightsData[r] = readout->weight->data.F32[yIn][xIn];
                }

                if (zero) {
                    pixelsData[r] -= zero->data.F32[r];
                }
                if (scale) {
                    pixelsData[r] *= invScale->data.F32[r];
                    if (params->weights) {
                        weightsData[r] *= invScale->data.F32[r] * invScale->data.F32[r];
                    }
                }

                numValid++;
            }

            if (numValid == 0) {
                outputMask[yOut][xOut] = params->blank;
                outputImage[yOut][xOut] = NAN;
                counts->data.U16[yOut][xOut] = 0;
                sigma->data.F32[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++;
                        numValid--;
                    }
                }
                // 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++;
                        numValid--;
                    }
                }
            }
            counts->data.U16[yOut][xOut] = numValid;

            // XXXXX this step probably is very expensive : convert errors to variance everywhere?
            if (params->weights) {
                errors = (psVector*)psUnaryOp(errors, weights, "sqrt");
            }

            // Combination
            if (!psVectorStats(stats, pixels, errors, mask, 1)) {
                // Can't do much about it, but it's not worth worrying about
                psErrorClear();
                outputImage[yOut][xOut] = NAN;
                outputMask[yOut][xOut] = params->blank;
                if (params->weights) {
                    outputWeight[yOut][xOut] = NAN;
                }
                sigma->data.F32[yOut][xOut] = NAN;
            } else {
                outputImage[yOut][xOut] = psStatsGetValue(stats, params->combine);
                outputMask[yOut][xOut] = isfinite(outputImage[yOut][xOut]) ? 0 : params->blank;
                if (params->weights) {
                    float stdev = psStatsGetValue(stats, combineStdev);
                    outputWeight[yOut][xOut] = PS_SQR(stdev); // Variance
                    // XXXX this is not the correct formal error.
                    // also, the weighted mean is not obviously the correct thing here
                }
                sigma->data.F32[yOut][xOut] = psStatsGetValue(stats, combineStdev);
            }
        }
    }
    #ifdef SHOW_BUSY
    if (psTraceGetLevel("psModules.imcombine") > 9) {
        printf("\n");
    }
    #endif
    psFree(index);
    psFree(pixels);
    psFree(mask);
    psFree(weights);
    psFree(errors);
    psFree(stats);
    psFree(invScale);
    psFree(counts);
    psFree(sigma);

    // Update the "concepts"
    psList *inputCells = psListAlloc(NULL); // List of cells
    for (long i = 0; i < inputs->n; i++) {
        pmReadout *readout = inputs->data[i]; // Readout of interest
        psListAdd(inputCells, PS_LIST_TAIL, readout->parent);
    }
    bool success = pmConceptsAverageCells(output->parent, inputCells, NULL, NULL, true);
    psFree(inputCells);

    output->data_exists = true;
    output->parent->data_exists = true;
    output->parent->parent->data_exists = true;

    return success;
}

