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

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

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// 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);
        masks->n = 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);
    }

    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);
        weights->n = 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// Readout for which to set mask
                     )
{
    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(__func__, 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

    // 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] |= PM_MASK_SAT;
            }
            if (imageData[i][j] <= bad) {
                maskData[i][j] |= PM_MASK_BAD;
            }
        }
    }

    return true;
}

bool pmReadoutGenerateMask(pmReadout *readout // Readout for which to generate mask
                          )
{
    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);
        readout->mask = mask;
    }

    return pmReadoutSetMask(readout);
}

bool pmReadoutSetWeight(pmReadout *readout // Readout for which to set weight
                       )
{
    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;
    }

    // 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));
    readout->weight = (psImage*)psBinaryOp(readout->weight, readout->weight, "+",
                                           psScalarAlloc(readnoise*readnoise/gain/gain, PS_TYPE_F32));

    return true;
}

bool pmReadoutGenerateWeight(pmReadout *readout // Readout for which to generate weight
                            )
{
    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) {
        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);
}

bool pmCellGenerateMaskWeight(pmCell *cell // Cell for which to set weights
                             )
{
    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 |= pmReadoutGenerateMask(readout);
        success |= pmReadoutGenerateWeight(readout);
    }

    return success;
}

bool pmReadoutGenerateMaskWeight(pmReadout *readout // Readout for which to set mask and weights
                                )
{
    PS_ASSERT_PTR_NON_NULL(readout, false);

    bool success = true;                // Was everything successful?

    success |= pmReadoutGenerateMask(readout);
    success |= pmReadoutGenerateWeight(readout);

    return success;
}
