#include <stdio.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;       // Size of image
    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;
}

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
    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

    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;
        }
        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(indices);
    psFree(clip);
    psFree(clipMask);
    psFree(poly);
    psFree(data);

    return true;
}
