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

#include "pmFPA.h"

#include "pmPattern.h"


// Mask a row as bad
static void patternMaskRow(pmReadout *ro, // Readout to mask
                           int y,       // Row to mask
                           psImageMaskType bad // Mask value to give
                           )
{
    psImage *image = ro->image;         // Image to mask
    psAssert(image, "Require image");
    psAssert(y < image->numRows, "Row not in image");

    int numCols = image->numCols;       // Number of columns
    for (int x = 0; x < numCols; x++) {
        image->data.F32[y][x] = NAN;
    }
    if (ro->mask) {
        psImage *mask = ro->mask;       // Mask image to mask
        for (int x = 0; x < numCols; x++) {
            mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] |= bad;
        }
    }
    return;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Measurement and application
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

bool pmPatternRow(pmReadout *ro, int order, int iter, float rej, float thresh,
                  psStatsOptions clipMean, psStatsOptions clipStdev,
                  psImageMaskType maskVal, psImageMaskType maskBad)
{
    PM_ASSERT_READOUT_NON_NULL(ro, false);
    PM_ASSERT_READOUT_IMAGE(ro, false);
    PS_ASSERT_INT_NONNEGATIVE(order, false);
    PS_ASSERT_INT_NONNEGATIVE(iter, false);
    PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    PS_ASSERT_FLOAT_LARGER_THAN(thresh, 0.0, false);

    psImage *image = ro->image;         // Image to correct
    psImage *mask = ro->mask;           // Mask for image
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS); // Random number generator
    if (!psImageBackground(stats, NULL, ro->image, ro->mask, maskVal, rng)) {
        psWarning("Unable to calculate statistics on readout; masking entire readout.");
        psErrorClear();
        psFree(stats);
        psFree(rng);
        psImageInit(image, NAN);
        if (mask) {
            psBinaryOp(mask, mask, "|", psScalarAlloc(maskBad, PS_TYPE_IMAGE_MASK));
        }
        if (ro->variance) {
            psImageInit(image, NAN);
        }
        return true;
    }
    float lower = stats->robustMedian - thresh * stats->robustStdev; // Lower bound for data
    float upper = stats->robustMedian + thresh * stats->robustStdev; // Upper bound for data
    float background = stats->robustMedian;
    psFree(stats);
    psFree(rng);

    // Indices are distributed [-1:1)
    psVector *indices = psVectorAlloc(numCols, PS_TYPE_F32); // Indices for fitting
    float norm = 2.0 / (float)numCols;  // Normalisation for indices
    for (int x = 0; x < numCols; x++) {
        indices->data.F32[x] = x * norm - 1.0;
    }

    psStats *clip = psStatsAlloc(clipMean | clipStdev); // Clipping statistics
    clip->clipIter = iter;
    clip->clipSigma = rej;
    psVector *clipMask = psVectorAlloc(numCols, PS_TYPE_VECTOR_MASK); // Mask for clipping
    psPolynomial1D *poly = psPolynomial1DAlloc(PS_POLYNOMIAL_ORD, order); // Polynomial to fit
    psVector *data = psVectorAlloc(numCols, PS_TYPE_F32); // Data to fit

    psImage *corr = psImageAlloc(order + 1, numRows, PS_TYPE_F64); // Corrections applied
    psImageInit(corr, NAN);

    for (int y = 0; y < numRows; y++) {
        psVectorInit(clipMask, 0);
        data = psImageRow(data, image, y);
        int num = 0;                    // Number of good pixels
        for (int x = 0; x < numCols; x++) {
            if ((mask && mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] & maskVal) ||
                data->data.F32[x] < lower || data->data.F32[x] > upper) {
                clipMask->data.PS_TYPE_VECTOR_MASK_DATA[x] = 0xFF;
            } else {
                clipMask->data.PS_TYPE_VECTOR_MASK_DATA[x] = 0;
                num++;
            }
        }
        if (num < order + 1) {
            // Not enough points to fit
            patternMaskRow(ro, y, maskBad);
            continue;
        }
        if (!psVectorClipFitPolynomial1D(poly, clip, clipMask, 0xFF, data, NULL, indices)) {
            psWarning("Unable to fit polynomial to row %d", y);
            psErrorClear();
            patternMaskRow(ro, y, maskBad);
            continue;
        }

        poly->coeff[0] -= background;
        memcpy(corr->data.F64[y], poly->coeff, (order + 1) * PSELEMTYPE_SIZEOF(PS_TYPE_F64));
        psVector *solution = psPolynomial1DEvalVector(poly, indices); // Solution vector
        if (!solution) {
            psWarning("Unable to evaluate polynomial for row %d", y);
            psErrorClear();
            patternMaskRow(ro, y, maskBad);
            continue;
        }

        for (int x = 0; x < numCols; x++) {
            image->data.F32[y][x] -= solution->data.F32[x];
        }
        psFree(solution);
    }

    psMetadataAddImage(ro->analysis, PS_LIST_TAIL, PM_PATTERN_ROW_CORRECTION, PS_META_REPLACE,
                       "Pattern row correction", corr);
    psFree(corr);

    psFree(indices);
    psFree(clip);
    psFree(clipMask);
    psFree(poly);
    psFree(data);

    return true;
}

bool pmPatternRowApply(pmReadout *ro, psImageMaskType maskBad)
{
    PM_ASSERT_READOUT_NON_NULL(ro, false);
    PM_ASSERT_READOUT_IMAGE(ro, false);

    bool mdok;                          // Status of MD lookup
    psImage *corr = psMetadataLookupPtr(&mdok, ro->analysis, PM_PATTERN_ROW_CORRECTION); // Correction
    if (!mdok) {
        // No correction to apply
        return true;
    }

    psImage *image = ro->image; // Image of interest
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    if (corr->numRows != numRows) {
        psError(PS_ERR_BAD_PARAMETER_SIZE, true,
                "Number of rows of image (%d) does not match number of rows of pattern correction (%d)\n",
                numRows, corr->numRows);
        return false;
    }

    int order = corr->numCols - 1;                                        // Polynomial order
    psPolynomial1D *poly = psPolynomial1DAlloc(PS_POLYNOMIAL_ORD, order); // Polynomial to apply
    psVector *indices = psVectorAlloc(numCols, PS_TYPE_F32); // Indices for polynomial
    float norm = 2.0 / (float)numCols;  // Normalisation for indices
    for (int x = 0; x < numCols; x++) {
        indices->data.F32[x] = x * norm - 1.0;
    }

    for (int y = 0; y < numRows; y++) {
        memcpy(poly->coeff, corr->data.F64[y], (order + 1) * PSELEMTYPE_SIZEOF(PS_TYPE_F64));
        psVector *solution = psPolynomial1DEvalVector(poly, indices); // Solution vector
        if (!solution) {
            psWarning("Unable to evaluate polynomial for row %d", y);
            psErrorClear();
            patternMaskRow(ro, y, maskBad);
            continue;
        }

        for (int x = 0; x < numCols; x++) {
            image->data.F32[y][x] -= solution->data.F32[x];
        }
        psFree(solution);
    }

    psFree(poly);
    psFree(indices);

    return true;
}


bool pmPatternCell(pmChip *chip, const psVector *tweak, psStatsOptions bgStat, psStatsOptions cellStat,
                   psImageMaskType maskVal, psImageMaskType maskBad)
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_VECTOR_NON_NULL(tweak, false);
    PS_ASSERT_VECTOR_SIZE(tweak, chip->cells->n, false);
    PS_ASSERT_VECTOR_TYPE(tweak, PS_TYPE_U8, false);

    int numCells = tweak->n;            // Number of cells

    psVector *mean = psVectorAlloc(numCells, PS_TYPE_F32); // Mean for each cell
    psVector *meanMask = psVectorAlloc(numCells, PS_TYPE_VECTOR_MASK); // Mask for means
    psVectorInit(mean, NAN);
    psVectorInit(meanMask, 0);

    // Mask bits
    enum {
        PM_PATTERN_IGNORE = 0x01,       // Ignore this cell
        PM_PATTERN_TWEAK  = 0x02,       // Tweak this cell
        PM_PATTERN_ERROR  = 0x04,       // Error in calculating background
        PM_PATTERN_ALL    = 0xFF,       // All causes
    };

    // Count number of cells to tweak
    int numTweak = 0;                   // Number of cells to tweak
    int numIgnore = 0;                  // Number of cells to ignore
    for (int i = 0; i < numCells; i++) {
        pmCell *cell = chip->cells->data[i]; // Cell of interest
        if (!cell || !cell->data_exists || !cell->process ||
            cell->readouts->n == 0 || cell->readouts->n > 1 || !cell->readouts->data[0]) {
            numIgnore++;
            meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] = PM_PATTERN_IGNORE;
            continue;
        }
        if (tweak->data.U8[i]) {
            meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] = PM_PATTERN_TWEAK;
            numTweak++;
        }
    }
    if (numTweak == 0) {
        // Nothing to do
        psFree(mean);
        psFree(meanMask);
        return true;
    }
    if (numTweak >= numCells - numIgnore) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Cannot pattern-correct all cells within a chip.");
        psFree(mean);
        psFree(meanMask);
        return false;
    }

    // Measure mean of each cell
    // This is not really the perfect thing to do, which would be to take a common mean for the set of cells
    // which aren't being tweaked (because some cells will be heavily masked, so shouldn't be weighted the
    // same as pure cells), but it's simple and fast.
    psStats *bgStats = psStatsAlloc(bgStat); // Statistics on background
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS); // Random number generator
    for (int i = 0; i < numCells; i++) {
        if (meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PM_PATTERN_IGNORE) {
            continue;
        }
        pmCell *cell = chip->cells->data[i]; // Cell of interest
        pmReadout *ro = cell->readouts->data[0]; // Readout of interest

        psStatsInit(bgStats);
#if 1
        if (!psImageBackground(bgStats, NULL, ro->image, ro->mask, maskVal, rng)) {
            psWarning("Unable to measure background for cell %d\n", i);
            psErrorClear();
            meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PM_PATTERN_ERROR;
            continue;
        }
#else
        if (!psImageStats(bgStats, ro->image, ro->mask, maskVal)) {
            psWarning("Unable to measure background for cell %d\n", i);
            psErrorClear();
            meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PM_PATTERN_ERROR;
            continue;
        }
#endif
        mean->data.F32[i] = psStatsGetValue(bgStats, bgStat);
        if (!isfinite(mean->data.F32[i])) {
            psWarning("Non-finite background for cell %d\n", i);
            psErrorClear();
            meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] |= PM_PATTERN_ERROR;
            continue;
        }
    }
    psFree(bgStats);
    psFree(rng);

    psStats *cellStats = psStatsAlloc(cellStat); // Statistics on cells
    if (!psVectorStats(cellStats, mean, NULL, meanMask, PM_PATTERN_ALL)) {
        // an error in psVectorStats implies a programming error
        psError(PS_ERR_UNKNOWN, false, "Unable to calculate mean cell background.");
        psFree(mean);
        psFree(meanMask);
        psFree(cellStats);
        return false;
    }

    float background = psStatsGetValue(cellStats, cellStat); // Background value for chip
    psFree(cellStats);
    if (!isfinite(background)) {
        // this can happen if all data in the image is bad -- in this case, do not correct, but
        // do not treat this as an error (other functions are responsible for check data validity)
        psLogMsg("psModules.detrend", PS_LOG_DETAIL, "Non-finite mean cell background -- skipping correction (data probabaly bad).");
        psFree(mean);
        psFree(meanMask);
        return true;
    }

    psLogMsg("psModules.detrend", PS_LOG_DETAIL, "Mean chip background is %f", background);

    for (int i = 0; i < numCells; i++) {
        if (meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PM_PATTERN_IGNORE) {
            continue;
        }
        if (!(meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PM_PATTERN_TWEAK)) {
            continue;
        }
        pmCell *cell = chip->cells->data[i]; // Cell of interest
        pmReadout *ro = cell->readouts->data[0]; // Readout of interest
        if (meanMask->data.PS_TYPE_VECTOR_MASK_DATA[i] & PM_PATTERN_ERROR) {
            psImageInit(ro->image, NAN);
            psBinaryOp(ro->mask, ro->mask, "|", psScalarAlloc(maskBad, PS_TYPE_IMAGE_MASK));
            psMetadataAddF32(ro->analysis, PS_LIST_TAIL, PM_PATTERN_CELL_CORRECTION, PS_META_REPLACE,
                             "Pattern cell correction solution", NAN);
            continue;
        }
        float correction = background - mean->data.F32[i]; // Correction to apply
        const char *cellName = psMetadataLookupStr(NULL, cell->concepts, "CELL.NAME"); // Name of cell
        psLogMsg("psModules.detrend", PS_LOG_DETAIL, "Correcting background of cell %s by %f",
                 cellName, correction);
        psBinaryOp(ro->image, ro->image, "+", psScalarAlloc(correction, PS_TYPE_F32));
        psMetadataAddF32(ro->analysis, PS_LIST_TAIL, PM_PATTERN_CELL_CORRECTION, PS_META_REPLACE,
                         "Pattern cell correction solution", correction);
    }

    psFree(mean);
    psFree(meanMask);

    return true;
}

bool pmPatternCellApply(pmReadout *ro, psImageMaskType maskBad)
{
    PM_ASSERT_READOUT_NON_NULL(ro, false);
    PM_ASSERT_READOUT_IMAGE(ro, false);

    bool mdok;                          // Status of MD lookup
    float corr = psMetadataLookupF32(&mdok, ro->analysis, PM_PATTERN_CELL_CORRECTION); // Correction to apply
    if (!mdok) {
        // No correction to apply
        return true;
    }

    psImage *image = ro->image, *mask = ro->mask; // Image and mask of interest
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    if (!isfinite(corr)) {
        for (int y = 0; y < numRows; y++) {
            for (int x = 0; x < numCols; x++) {
                image->data.F32[y][x] = NAN;
            }
        }
        if (mask) {
            for (int y = 0; y < numRows; y++) {
                for (int x = 0; x < numCols; x++) {
                    mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] |= maskBad;
                }
            }
        }
    } else {
        for (int y = 0; y < numRows; y++) {
            for (int x = 0; x < numCols; x++) {
                image->data.F32[y][x] += corr;
            }
        }
    }

    return true;
}


