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

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

#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmFPAMaskWeight.h"
#include "pmConceptsAverage.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->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;
    }

    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);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK, comment, "");
        psFree(comment);
    }

    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;

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

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

    bool valid = false;                 // Do we have a single valid input?
    for (long i = 0; i < inputs->n; i++) {
        pmReadout *readout = inputs->data[i]; // Readout of interest
        pmCell *cell = readout->parent; // The parent cell

        if (!readout || !readout->image) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "Input readout %ld is NULL or has NULL image.\n", i);
            return false;
        }

        bool mdok = true;       // Status of MD lookup
        psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
        if (!mdok || !trimsec || psRegionIsNaN(*trimsec)) {
            psLogMsg(__func__, PS_LOG_WARN, "CELL.TRIMSEC is not set for readout %ld --- ignored.\n", i);
        } else {
            xSize = PS_MAX(xSize, trimsec->x1 - trimsec->x0);
            ySize = PS_MAX(ySize, trimsec->y1 - trimsec->y0);
        }

        if (params->weights && !readout->weight) {
            psError(PS_ERR_UNEXPECTED_NULL, true,
                    "Rejection based on weights requested, but no weights supplied for image %ld.\n", i);
            return false;
        }

        valid = true;

        // Range of pixels on 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("psModules.imcombine", 7, "Readout %ld: 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;
    }

    // Update the origin
    if (output->image) {
        *(psS32 *) &(output->col0) = PS_MIN(minInputCols, output->col0);
        *(psS32 *) &(output->row0) = PS_MIN(minInputRows, output->row0);
    } else {
        *(psS32 *) &(output->col0) = minInputCols;
        *(psS32 *) &(output->row0) = minInputRows;
    }
    psTrace("psModules.imcombine", 7, "Output minimum: %d,%d\n", output->col0, output->row0);

    // Allocate the output products

    if (!output->image) {
        output->image = psImageAlloc(xSize, ySize, PS_TYPE_F32);
    }
    if (output->image->numCols < xSize || output->image->numRows < ySize) {
        // Generate the new output image by extending the current one, or making a whole new one
        psImage *newImage = psImageAlloc(xSize, ySize, PS_TYPE_F32);
        psImageInit(newImage, 0.0);
        psImageOverlaySection(newImage, output->image, output->col0, output->row0, "=");
        psFree(output->image);
        output->image = newImage;
    }

    if (!output->mask) {
        output->mask = psImageAlloc(xSize, ySize, PS_TYPE_U8);
    }
    if (output->mask->numCols < xSize || output->mask->numRows < ySize) {
        psImage *newMask = psImageAlloc(xSize, ySize, PS_TYPE_U8);
        psImageInit(newMask, 0);
        psImageOverlaySection(newMask, output->mask, output->col0, output->row0, "=");
        psFree(output->mask);
        output->mask = newMask;
    }

    psStatsOptions combineStdev = 0; // Statistics option for weights
    if (params->weights) {

        if (!output->weight) {
            output->weight = psImageAlloc(xSize, ySize, PS_TYPE_F32);
        }
        if (output->weight->numCols < xSize || output->weight->numRows < ySize) {
            psImage *newWeight = psImageAlloc(xSize, ySize, PS_TYPE_F32);
            psImageInit(newWeight, 0.0);
            psImageOverlaySection(newWeight, output->weight, output->col0, output->row0, "=");
            psFree(output->weight);
            output->weight = newWeight;
        }

        if (first) {
            psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                             "Using input weights to combine images", "");
        }

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

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

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

            // Combination
            psVectorStats(stats, pixels, errors, mask, 1);
            outputImage[yOut][xOut] = psStatsGetValue(stats, params->combine);
            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
            }
        }
    }
    #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);

    // 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);
    }
    pmConceptsAverageCells(output->parent, inputCells, NULL, NULL);
    psFree(inputCells);

    return true;
}

