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

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

#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmFPAMaskWeight.h"
#include "pmMaskBadPixels.h"

bool pmMaskBadPixels(pmReadout *input, const pmReadout *mask, psMaskType maskVal)
{
    PS_ASSERT_PTR_NON_NULL(input, false);
    PS_ASSERT_PTR_NON_NULL(input->mask, false);
    PS_ASSERT_IMAGE_TYPE(input->mask, PS_TYPE_MASK, false);

    PS_ASSERT_PTR_NON_NULL(mask, false);
    PS_ASSERT_PTR_NON_NULL(mask->mask, false);
    PS_ASSERT_IMAGE_TYPE(mask->mask, PS_TYPE_MASK, false);

    psImage *inMask = input->mask;
    psImage *exMask = mask->mask;

    // Add mask MD5 to header
    pmHDU *hdu = pmHDUFromReadout(input);  // HDU of interest
    psVector *md5 = psImageMD5(mask->mask); // md5 hash
    psString md5string = psMD5toString(md5); // String
    psFree(md5);
    psStringPrepend(&md5string, "MASK image MD5: ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     md5string, "");
    psFree(md5string);

    int rowMax = input->row0 + inMask->numRows;
    int colMax = input->col0 + inMask->numCols;

    if (mask->row0 > input->row0 || mask->col0 > input->col0 ||
            mask->row0 + exMask->numRows < rowMax || mask->col0 + exMask->numCols < colMax) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                "Input image size exceeds that of mask image: (%d, %d) vs (%d, %d)",
                inMask->numRows, inMask->numCols, exMask->numRows, exMask->numCols);
        return false;
    }

    // Determine total offset based on image offset with chip offset
    // XXX if we choose to correct for the readout location, apply input->col0,row0 here
    int offCol = input->col0 - mask->col0;
    int offRow = input->row0 - mask->row0;

    // masks are both of type PS_TYPE_MASK
    psMaskType **exVal = exMask->data.U8;
    psMaskType **inVal = inMask->data.U8;

    // apply exMask values
    if (maskVal) {
        // set raised pixels in exMask which are selected by maskVal
        for (int j = 0; j < inMask->numRows; j++) {
            int xJ = j - offRow;
            for (int i = 0; i < inMask->numCols; i++) {
                int xI = i - offCol;
                inVal[j][i] |= (maskVal & exVal[xJ][xI]);
            }
        }
    } else {
        // set raised pixels in exMask which are selected by maskVal
        for (int j = 0; j < inMask->numRows; j++) {
            int xJ = j - offRow;
            for (int i = 0; i < inMask->numCols; i++) {
                int xI = i - offCol;
                if (exVal[xJ][xI] == 0) {
                    inVal[j][i] |= PM_MASK_BAD;
                }
            }
        }
    }

    psTime *time = psTimeGetNow(PS_TIME_TAI); // The time now, used for reporting
    psString timeString = psTimeToISO(time); // String with time
    psFree(time);
    psStringPrepend(&timeString, "Static mask applied at ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     timeString, "");
    psFree(timeString);

    return true;
}


psImage *pmMaskFlagSuspectPixels(psImage *out, const pmReadout *readout, float rej,
                                 psMaskType maskVal, float frac, psRandom *rng)
{
    PS_ASSERT_PTR_NON_NULL(readout, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, NULL);
    PS_ASSERT_FLOAT_WITHIN_RANGE(frac, 0.0, 1.0, NULL);
    PS_ASSERT_IMAGE_NON_NULL(readout->image, NULL);
    PS_ASSERT_IMAGE_NON_EMPTY(readout->image, NULL);
    PS_ASSERT_IMAGE_TYPE(readout->image, PS_TYPE_F32, NULL);
    if (readout->mask) {
        PS_ASSERT_IMAGE_NON_EMPTY(readout->mask, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(readout->image, readout->mask, NULL);
        PS_ASSERT_IMAGE_TYPE(readout->mask, PS_TYPE_MASK, NULL);
    }
    if (out) {
        PS_ASSERT_IMAGE_NON_EMPTY(out, NULL);
        PS_ASSERT_IMAGE_TYPE(out, PS_TYPE_S32, NULL);
        PS_ASSERT_IMAGES_SIZE_EQUAL(readout->image, out, NULL);
    }

    psImage *image = readout->image;    // Image of interest
    psImage *mask = readout->mask;      // Corresponding mask

    if (rng) {
        psMemIncrRefCounter(rng);
    } else {
        rng = psRandomAlloc(PS_RANDOM_TAUS, 0);
    }

    // XXX note that this now will accept any of several stats options
    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    stats->nSubsample = frac * image->numCols * image->numRows;
    if (!psImageBackground(stats, image, mask, maskVal, rng) ||
            !isfinite(stats->robustMedian) || !isfinite(stats->robustUQ) || !isfinite(stats->robustLQ)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to measure image statistics.\n");
        psFree(stats);
        psFree(rng);
        return NULL;
    }
    psFree(rng);

    float median = stats->robustMedian; // Median value
    float stdev = stats->robustStdev; // Estimate of the standard deviation
    psFree(stats);

    if (!out) {
        out = psImageAlloc(image->numCols, image->numRows, PS_TYPE_S32);
        psImageInit(out, 0);
    }

    for (int y = 0; y < image->numRows; y++) {
        for (int x = 0; x < image->numCols; x++) {
            if (fabs((image->data.F32[y][x] - median) / stdev) >= rej &&
                    (!mask || !(mask->data.PS_TYPE_MASK_DATA[y][x] & maskVal))) {
                out->data.S32[y][x]++;
            }
        }
    }

    return out;
}


psImage *pmMaskIdentifyBadPixels(const psImage *suspects, float thresh, psMaskType maskVal)
{
    PS_ASSERT_IMAGE_NON_NULL(suspects, NULL);
    PS_ASSERT_IMAGE_NON_EMPTY(suspects, NULL);
    PS_ASSERT_IMAGE_TYPE(suspects, PS_TYPE_S32, NULL);

    float limit = NAN;                  // Limit for masking
    if (thresh > 0) {
        psStats *stats = psStatsAlloc(PS_STAT_CLIPPED_STDEV); // Statistics
        stats->clipSigma = 5.0;
        stats->clipIter = 1;
        if (!psImageStats(stats, suspects, NULL, 0) || !isfinite(stats->clippedStdev)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to perform statistics.\n");
            psFree(stats);
            return NULL;
        }
        limit = thresh * stats->clippedStdev;
        psFree(stats);
    } else if (thresh < 0) {
        psStats *stats = psStatsAlloc(PS_STAT_MAX); // Statistics
        if (!psImageStats(stats, suspects, NULL, 0)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to perform statistics.\n");
            psFree(stats);
            return NULL;
        }
        psHistogram *histo = psHistogramAlloc(-0.5, stats->max + 0.5, stats->max + 1);
        psFree(stats);
        if (!psImageHistogram(histo, suspects, NULL, 0)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate histogram.\n");
            psFree(histo);
            return NULL;
        }

        // Find the mode.  Since this is a Poisson distribution (more or less), this should also be the mean
        // and variance.
        int max = 0;                    // Index of the mode
        for (int i = 0; i < histo->nums->n; i++) {
            if (histo->nums->data.F32[i] > histo->nums->data.F32[max]) {
                max = i;
            }
        }

        // Since the mode is most likely zero, we add one to get something realistic.  Then "thresh" is
        // negative, so we subtract instead of add.
        limit = max + 1.0 - thresh * sqrtf((float)max + 1.0);

    } else {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                "Threshold (%f) must be a positive or negative real number.\n", thresh);
        return NULL;
    }

    psImage *badpix = psImageAlloc(suspects->numCols, suspects->numRows, PS_TYPE_MASK); // Bad pixel mask
    psImageInit(badpix, 0);

    if (psTraceGetLevel("psModules.detrend") > 9) {
        psStats *stats = psStatsAlloc(PS_STAT_MIN | PS_STAT_MAX); // Statistics
        psImageStats(stats, suspects, NULL, 0);
        psHistogram *histo = psHistogramAlloc(-0.5, stats->max + 0.5, stats->max + 1);
        psImageHistogram(histo, suspects, NULL, 0);
        for (int i = 0; i < histo->nums->n; i++) {
            printf("%f --> %f : %f\n", histo->bounds->data.F32[i], histo->bounds->data.F32[i + 1],
                   histo->nums->data.F32[i]);
        }
        psFree(stats);
        psFree(histo);
        printf("Threshold: %f\n", limit);
    }

    for (int y = 0; y < suspects->numRows; y++) {
        for (int x = 0; x < suspects->numCols; x++) {
            if (suspects->data.S32[y][x] >= limit) {
                badpix->data.PS_TYPE_MASK_DATA[y][x] = maskVal;
            }
        }
    }

    return badpix;
}
