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

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

#include <pslib.h>
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmHDUUtils.h"
#include "pmHDUGenerate.h"
#include "pmFPAMaskWeight.h"

#define PIXELS_BUFFER 100               // Size of buffer for allocating pixel lists

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// File-static (private) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Create the parent mask images that reside in the HDU
static void createParentMasks(pmHDU *hdu // The HDU for which to create
                             )
{
    assert(hdu);
    assert(hdu->images);

    // Generate the parent mask images
    psArray *images = hdu->images;      // Array of images
    psArray *masks = hdu->masks;        // Array of masks
    if (!masks) {
        masks = psArrayAlloc(images->n);
        hdu->masks = masks;
    }

    for (long i = 0; i < images->n; i++) {
        psImage *image = images->data[i]; // The image for this readout
        if (!image || masks->data[i]) {
            continue;
        }
        masks->data[i] = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
        psImageInit(masks->data[i], 0);
    }

    return;
}

// Create the parent weight images that reside in the HDU
static void createParentWeights(pmHDU *hdu // The HDU for which to create
                               )
{
    assert(hdu);
    assert(hdu->images);

    // Generate the parent mask images
    psArray *images = hdu->images;      // Array of images
    psArray *weights = hdu->weights;    // Array of weight images
    if (!weights) {
        weights = psArrayAlloc(images->n);
        hdu->weights = weights;
    }

    for (long i = 0; i < images->n; i++) {
        psImage *image = images->data[i]; // The image for this readout
        if (!image || weights->data[i]) {
            continue;
        }
        weights->data[i] = psImageAlloc(image->numCols, image->numRows, PS_TYPE_F32);
    }

    return;
}

// Identify a readout within the HDU, on the basis of the image pointer.
// This is a little dirty, but hopefully should work....
static long identifyReadout(pmHDU *hdu, // The HDU containing the readouts
                            pmReadout *readout // The readout to be identified
                           )
{
    assert(hdu);
    assert(readout);

    long index = -1;                    // Index of the readout
    for (long i = 0; i < hdu->images->n && index == -1; i++) {
        if (hdu->images->data[i] == readout->image->parent) {
            index = i;
        }
    }

    return index;
}

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

bool pmReadoutSetMask(pmReadout *readout, psMaskType satMask, psMaskType badMask)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_IMAGE_NON_NULL(readout->image, false);

    pmCell *cell = readout->parent;     // The parent cell
    bool mdok = true;                   // Status of MD lookup

    // Get the "concepts" of interest
    float saturation = psMetadataLookupF32(&mdok, cell->concepts, "CELL.SATURATION"); // Saturation level
    if (!mdok || isnan(saturation)) {
        psError(PS_ERR_IO, true, "CELL.SATURATION is not set --- unable to set mask.\n");
        return false;
    }
    float bad = psMetadataLookupF32(&mdok, cell->concepts, "CELL.BAD"); // Bad level
    if (!mdok || isnan(bad)) {
        psError(PS_ERR_IO, true, "CELL.BAD is not set --- unable to set mask.\n");
        return false;
    }
    psTrace("psModules.camera", 5, "Saturation: %f, bad: %f\n", saturation, bad);


    // Set up the mask
    psImage *image = readout->image;    // The image pixels
    if (!readout->mask) {
        // Generate a (throwaway) mask image, if required
        readout->mask = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
    }
    psImage *mask = readout->mask;      // The mask pixels
    psImageInit(mask, 0);

    // Dereference pointers for speed
    psF32 **imageData = image->data.F32;// The image
    psU8  **maskData  = mask->data.U8;  // The mask

    for (int i = 0; i < image->numRows; i++) {
        for (int j = 0; j < image->numCols; j++) {
            if (imageData[i][j] >= saturation) {
                maskData[i][j] |= satMask;
            }
            if (imageData[i][j] <= bad) {
                maskData[i][j] |= badMask;
            }
            if (!isfinite(imageData[i][j])) {
                maskData[i][j] |= badMask;
            }
        }
    }

    return true;
}

// XXX this function creates the mask pixels, or uses the existing mask
// pixels.  currently, it will set mask bits if (value <= BAD) or (value >= SATURATION)
// should we optionally ignore these tests?
bool pmReadoutGenerateMask(pmReadout *readout, psMaskType satMask, psMaskType badMask)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);

    pmCell *cell = readout->parent;     // The parent cell
    bool mdok = true;                   // Status of MD lookup

    // Create the mask image if required
    if (!readout->mask) {
        psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
        if (!mdok || psRegionIsNaN(*trimsec)) {
            psError(PS_ERR_IO, true, "CELL.TRIMSEC is not set --- unable to set mask.\n");
            return false;
        }

        pmHDU *hdu = pmHDUFromCell(cell);   // The HDU containing the cell's pixels
        PS_ASSERT_PTR_NON_NULL(hdu, false);
        if (!hdu->images && !pmHDUGenerateForCell(cell)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate HDU for cell.\n");
            return false;
        }

        createParentMasks(hdu);

        // Need to identify which readout we're working with....
        long index = identifyReadout(hdu, readout); // Index of the readout
        if (index == -1) {
            psError(PS_ERR_UNKNOWN, true, "Unable to identify readout image in HDU.\n");
            return false;
        }

        psImage *mask = psImageSubset(hdu->masks->data[index], *trimsec); // The mask pixels
        if (!mask) {
            psString trimsecString = psRegionToString(*trimsec);
            psError(PS_ERR_UNKNOWN, false, "Unable to set mask from HDU with trimsec: %s.\n", trimsecString);
            psFree(trimsecString);
            return false;
        }
        psImageInit(mask, 0);
        assert (readout->mask == NULL); // or else this is a memory leak.
        readout->mask = mask;
    }

    return pmReadoutSetMask(readout, satMask, badMask);
}

bool pmReadoutSetWeight(pmReadout *readout, bool poisson)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);

    pmCell *cell = readout->parent;     // The parent cell
    bool mdok = true;                   // Status of MD lookup

    // Get the "concepts" of interest
    float gain = psMetadataLookupF32(&mdok, cell->concepts, "CELL.GAIN"); // Cell gain
    if (!mdok || isnan(gain)) {
        psError(PS_ERR_IO, true, "CELL.GAIN is not set --- unable to set weight.\n");
        return false;
    }
    float readnoise = psMetadataLookupF32(&mdok, cell->concepts, "CELL.READNOISE"); // Cell read noise
    if (!mdok || isnan(readnoise)) {
        psError(PS_ERR_IO, true, "CELL.READNOISE is not set --- unable to set weight.\n");
        return false;
    }

    if (poisson) {
        // Set weight image to the variance in ADU = f/g + rn^2
        psImage *image = readout->image;    // The image pixels
        readout->weight = (psImage*)psBinaryOp(readout->weight, image, "/", psScalarAlloc(gain, PS_TYPE_F32));

        // a negative weight is non-sensical. if the image value drops below 1, the weight must be 1.
        readout->weight = (psImage*)psUnaryOp(readout->weight, readout->weight, "abs");
        readout->weight = (psImage*)psBinaryOp(readout->weight, readout->weight, "max",
                                               psScalarAlloc(1, PS_TYPE_F32));
    } else {
        // Just use the read noise
        if (!readout->weight) {
            readout->weight = psImageAlloc(readout->image->numCols, readout->image->numRows, PS_TYPE_F32);
        }
        psImageInit(readout->weight, 0.0);
    }

    readout->weight = (psImage*)psBinaryOp(readout->weight, readout->weight, "+",
                                           psScalarAlloc(readnoise*readnoise/gain/gain, PS_TYPE_F32));

    return true;
}

// this function creates the weight pixels, or uses the existing weight pixels.  it will set
// the noise pixel values only if the weight image is not supplied
bool pmReadoutGenerateWeight(pmReadout *readout, bool poisson)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);

    pmCell *cell = readout->parent;     // The parent cell
    bool mdok = true;                   // Status of MD lookup

    // Create the weight image if required
    if (readout->weight)
        return true;

    psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
    if (!mdok || psRegionIsNaN(*trimsec)) {
        psError(PS_ERR_IO, true, "CELL.TRIMSEC is not set --- unable to set weight.\n");
        return false;
    }

    pmHDU *hdu = pmHDUFromCell(cell);   // The HDU containing the cell's pixels
    PS_ASSERT_PTR_NON_NULL(hdu, false);
    if (!hdu->images && !pmHDUGenerateForCell(cell)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to generate HDU for cell.\n");
        return false;
    }

    createParentWeights(hdu);

    // Need to identify which readout we're working with....
    long index = identifyReadout(hdu, readout); // Index of the readout
    if (index == -1) {
        psError(PS_ERR_UNKNOWN, true, "Unable to identify readout image in HDU.\n");
        return false;
    }

    psImage *weight = psImageSubset(hdu->weights->data[index], *trimsec); // The weight pixels
    if (!weight) {
        psString trimsecString = psRegionToString(*trimsec);
        psError(PS_ERR_UNKNOWN, false, "Unable to set weight from HDU with trimsec: %s.\n",
                trimsecString);
        psFree(trimsecString);
        return false;
    }
    psImageInit(weight, 0);
    readout->weight = weight;

    return pmReadoutSetWeight(readout, poisson);
}

bool pmReadoutGenerateMaskWeight(pmReadout *readout, psMaskType satMask, psMaskType badMask, bool poisson)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);

    bool success = true;                // Was everything successful?

    success &= pmReadoutGenerateMask(readout, satMask, badMask);
    success &= pmReadoutGenerateWeight(readout, poisson);

    return success;
}

bool pmCellGenerateMaskWeight(pmCell *cell, psMaskType satMask, psMaskType badMask, bool poisson)
{
    PS_ASSERT_PTR_NON_NULL(cell, false);

    bool success = true;                // Was everything successful?
    psArray *readouts = cell->readouts; // Array of readouts
    for (int i = 0; i < readouts->n; i++) {
        pmReadout *readout = readouts->data[i]; // The readout
        success &= pmReadoutGenerateMaskWeight(readout, poisson, satMask, badMask);
    }

    return success;
}


bool pmReadoutWeightRenorm(const pmReadout *readout, psMaskType maskVal, psStatsOptions meanStat,
                           psStatsOptions stdevStat, int width, psRandom *rng)
{
    PM_ASSERT_READOUT_NON_NULL(readout, false);
    PM_ASSERT_READOUT_IMAGE(readout, false);
    PM_ASSERT_READOUT_WEIGHT(readout, false);
    PS_ASSERT_INT_POSITIVE(width, false);

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

    psImage *image = readout->image, *mask = readout->mask, *weight = readout->weight; // Readout images
    int numCols = image->numCols, numRows = image->numRows; // Size of images
    int xNum = numCols / width + 1, yNum = numRows / width + 1; // Number of renormalisation regions
    float xSize = numCols / (float)xNum, ySize = numRows / (float)yNum; // Size of renormalisation regions

    psStats *meanStats = psStatsAlloc(meanStat), *stdevStats = psStatsAlloc(stdevStat); // Statistics
    psVector *buffer = NULL;

    for (int j = 0; j < yNum; j++) {
        // Bounds in y
        int yMin = j * ySize;
        int yMax = (j + 1) * ySize;
        for (int i = 0; i < xNum; i++) {
            // Bounds in x
            int xMin = i * xSize;
            int xMax = (i + 1) * xSize;

            psRegion region = psRegionSet(xMin, xMax, yMin, yMax); // Region of interest
            psImage *subImage = psImageSubset(image, region); // Sub-image of the image pixels
            psImage *subWeight = psImageSubset(weight, region); // Sub image of the weight pixels
            psImage *subMask = mask ? psImageSubset(mask, region) : NULL; // Sub-image of the mask pixels

            if (!psImageBackground(stdevStats, &buffer, subImage, subMask, maskVal, rng) ||
                !psImageBackground(meanStats, &buffer, subWeight, subMask, maskVal, rng)) {
                // Nothing we can do about it, but don't want to keel over and die, so do our best to flag it.
                psString regionStr = psRegionToString(region); // String with region
                psWarning("Unable to measure statistics over %s", regionStr);
                psFree(regionStr);
                psErrorClear();
                psImageInit(subWeight, NAN);
                if (subMask) {
                    psImageInit(subMask, maskVal);
                }
            } else {
                double meanVar = psStatsGetValue(meanStats, meanStat); // Mean of variance map
                double stdev = psStatsGetValue(stdevStats, stdevStat); // Standard deviation of image
                psTrace("psModules.camera", 3,
                        "Region [%d:%d,%d:%d] has variance %lf, but mean of variance map is %lf\n",
                        xMin, xMax, yMin, yMax, PS_SQR(stdev), meanVar);
                psBinaryOp(subWeight, subWeight, "*", psScalarAlloc(PS_SQR(stdev) / meanVar, PS_TYPE_F32));
            }

            psFree(subImage);
            psFree(subWeight);
            psFree(subMask);
        }
    }
    psFree(meanStats);
    psFree(stdevStats);
    psFree(rng);
    psFree(buffer);

    return true;
}


bool pmReadoutMaskNonfinite(pmReadout *readout, psMaskType maskVal)
{
    PM_ASSERT_READOUT_NON_NULL(readout, false);
    PM_ASSERT_READOUT_IMAGE(readout, false);

    psImage *image = readout->image;    // Readout's image
    psImage *weight = readout->weight;  // Readout's weight
    int numCols = image->numCols, numRows = image->numRows; // Size of image

    if (!readout->mask) {
        readout->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
    }
    psImage *mask = readout->mask;      // Readout's mask

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (!isfinite(image->data.F32[y][x]) || (weight && !isfinite(weight->data.F32[y][x]))) {
                mask->data.PS_TYPE_MASK_DATA[y][x] |= maskVal;
            }
        }
    }

    return true;
}



bool pmReadoutMaskApply(pmReadout *readout, psMaskType maskVal)
{
    PM_ASSERT_READOUT_NON_NULL(readout, false);
    PM_ASSERT_READOUT_IMAGE(readout, false);
    PM_ASSERT_READOUT_MASK(readout, false);

    int numCols = readout->image->numCols, numRows = readout->image->numRows; // Size of image
    psMaskType **maskData = readout->mask->data.PS_TYPE_MASK_DATA; // Dereference mask
    psF32 **imageData = readout->image->data.F32;// Dereference image
    psF32 **weightData = readout->weight ? readout->weight->data.F32 : NULL; // Dereference weight map

    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (maskData[y][x] & maskVal) {
                imageData[y][x] = NAN;
                if (weightData) {
                    weightData[y][x] = NAN;
                }
            }
        }
    }

    return true;
}


bool pmReadoutInterpolateBadPixels(pmReadout *readout, psMaskType maskVal, psImageInterpolateMode mode,
                                   float poorFrac, psMaskType maskPoor, psMaskType maskBad)
{
    PM_ASSERT_READOUT_NON_NULL(readout, false);
    PM_ASSERT_READOUT_IMAGE(readout, false);
    PM_ASSERT_READOUT_MASK(readout, false);
    if (!maskVal) {
        return true;
    }

    psImage *image = readout->image;    // Image
    psImage *mask = readout->mask;      // Mask
    psImage *weight = readout->weight;  // Weight map

    psImageInterpolateOptions *interp = psImageInterpolateOptionsAlloc(mode, image, weight, mask, maskVal,
                                                                       NAN, NAN, maskBad, maskPoor, poorFrac);
    interp->shifting = false;           // Turn off "exact shifts" so we get proper interpolation

    int numCols = mask->numCols, numRows = mask->numRows; // Size of image

    psPixels *pixels = psPixelsAllocEmpty(PIXELS_BUFFER); // Pixels that have been interpolated
    psVector *imagePix = psVectorAllocEmpty(PIXELS_BUFFER, PS_TYPE_F32); // Corresponding values for image
    psVector *weightPix = psVectorAllocEmpty(PIXELS_BUFFER, PS_TYPE_F32); // Corresponding values for weight
    psVector *maskPix = psVectorAllocEmpty(PIXELS_BUFFER, PS_TYPE_MASK); // Corresponding values for mask

    long numBad = 0;                    // Number of bad pixels interpolated
    for (int y = 0; y < numRows; y++) {
        for (int x = 0; x < numCols; x++) {
            if (mask->data.PS_TYPE_MASK_DATA[y][x] & maskVal) {
                double imageValue, weightValue; // Image and weight value from interpolation
                psMaskType maskValue = 0; // Mask value from interpolation
                psImageInterpolateStatus status = psImageInterpolate(&imageValue, &weightValue, &maskValue,
                                                                     x, y, interp);
                if (status == PS_INTERPOLATE_STATUS_ERROR || status == PS_INTERPOLATE_STATUS_OFF) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to interpolate readout at %d,%d", x, y);
                    psFree(interp);
                    psFree(pixels);
                    psFree(imagePix);
                    psFree(weightPix);
                    psFree(maskPix);
                    return false;
                }
                if (status == PS_INTERPOLATE_STATUS_BAD) {
                    // It's still bad: couldn't interpolate enough
                    continue;
                }

                pixels = psPixelsAdd(pixels, PIXELS_BUFFER, x, y);
                imagePix = psVectorExtend(imagePix, PIXELS_BUFFER, 1);
                weightPix = psVectorExtend(weightPix, PIXELS_BUFFER, 1);
                maskPix = psVectorExtend(maskPix, PIXELS_BUFFER, 1);
                imagePix->data.F32[numBad] = imageValue;
                weightPix->data.F32[numBad] = weightValue;
                maskPix->data.PS_TYPE_MASK_DATA[numBad] = maskValue;
                numBad++;
            }
        }
    }

    psFree(interp);

    for (long i = 0; i < numBad; i++) {
        int x = pixels->data[i].x, y = pixels->data[i].y; // Coordinates of pixel
        image->data.F32[y][x] = imagePix->data.F32[i];
        weight->data.F32[y][x] = weightPix->data.F32[i];
        mask->data.PS_TYPE_MASK_DATA[y][x] = maskPix->data.PS_TYPE_MASK_DATA[i];
    }

    psFree(pixels);
    psFree(imagePix);
    psFree(weightPix);
    psFree(maskPix);

    psLogMsg("psModules.camera", PS_LOG_INFO, "Interpolated over %ld pixels", numBad);

    return true;
}
