#include <stdio.h>
#include "pslib.h"
#include "pmFPA.h"
#include "pmMaskBadPixels.h"
#include "pmHDUUtils.h"

// Get the bias images for a readout, using the CELL.BIASSEC
psList *pmReadoutGetBias(pmReadout *readout // Readout for which to get the bias sections
                        )
{
    pmCell *cell = readout->parent;     // The parent cell
    psList *sections = (psList*)psMetadataLookupPtr(NULL, cell->concepts, "CELL.BIASSEC"); // CELL.BIASSEC

    psImage *image = readout->image;    // The image that contains both the biassec and the trimsec

    psList *images = psListAlloc(NULL); // List of images from bias sections
    psListIterator *sectionsIter = psListIteratorAlloc(sections, PS_LIST_HEAD, true); // Iterator
    psRegion *region = NULL;            // Bias region from list
    while ((region = psListGetAndIncrement(sectionsIter))) {
        psImage *bias = psMemIncrRefCounter(psImageSubset(image, *region)); // Image from bias section
        psListAdd(images, PS_LIST_TAIL, bias);
        psFree(bias);
    }
    psFree(sectionsIter);

    return images;
}

bool pmReadoutSetWeights(pmReadout *readout)
{

    pmCell *cell = readout->parent;

    float gain = psMetadataLookupF32(NULL, cell->concepts, "CELL.GAIN"); // Cell gain
    float readnoise = psMetadataLookupF32(NULL, cell->concepts, "CELL.READNOISE"); // Cell read noise
    psRegion *trimsec = psMetadataLookupPtr(NULL, cell->concepts, "CELL.TRIMSEC"); // Trim section
    float saturation = psMetadataLookupF32(NULL, cell->concepts, "CELL.SATURATION"); // Saturation level
    float bad = psMetadataLookupF32(NULL, cell->concepts, "CELL.BAD"); // Bad level

    pmHDU *hdu = pmHDUFromCell (cell);

    psArray *pixels = hdu->images;      // Array of images
    psArray *weights = hdu->weights;    // Array of weight images
    psArray *masks = hdu->masks;        // Array of mask images
    // Generate the weights and masks if required
    if (! weights) {
        weights = psArrayAlloc(pixels->n);
        for (int i = 0; i < pixels->n; i++) {
            psImage *image = pixels->data[i];
            weights->data[i] = psImageAlloc(image->numCols, image->numRows, PS_TYPE_F32);
            psImageInit(weights->data[i], 0.0);
        }
        hdu->weights = weights;
    }
    if (! masks) {
        masks = psArrayAlloc(pixels->n);
        for (int i = 0; i < pixels->n; i++) {
            psImage *image = pixels->data[i];
            masks->data[i] = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
            psImageInit(masks->data[i], 0);
        }
        hdu->masks = masks;
    }

    // Set the pixels
    psArray *readouts = cell->readouts; // Array of readouts
    for (int i = 0; i < readouts->n; i++) {
        pmReadout *readout = readouts->data[i]; // The readout of interest

        if (! readout->weight) {
            readout->weight = psMemIncrRefCounter(psImageSubset(weights->data[i], *trimsec));
        }
        if (! readout->mask) {
            readout->mask = psMemIncrRefCounter(psImageSubset(masks->data[i], *trimsec));
        }

        // Set up the mask
        psImage *image = readout->image;// Pixels
        psImage *mask = readout->mask;  // Mask image
        for (int i = 0; i < image->numRows; i++) {
            for (int j = 0; j < image->numCols; j++) {
                if (image->data.F32[i][j] >= saturation) {
                    mask->data.U8[i][j] = PM_MASK_SAT;
                }
                if (image->data.F32[i][j] <= bad) {
                    mask->data.U8[i][j] = PM_MASK_BAD;
                }
            }
        }

        psImage *weight = readout->weight;  // Mask image
        float rnoise = PS_SQR(readnoise/gain);
        for (int i = 0; i < image->numRows; i++) {
            for (int j = 0; j < image->numCols; j++) {
                if (!mask->data.U8[i][j]) {
                    weight->data.F32[i][j] = image->data.F32[i][j]/ gain + rnoise;
                }
            }
        }

        // Set weight image to the variance = g*f + rn^2
        // psBinaryOp(readout->weight, image, "/", psScalarAlloc(gain, PS_TYPE_F32));
        // psBinaryOp(readout->weight, readout->weight, "+",
        // psScalarAlloc(readnoise*readnoise/gain/gain, PS_TYPE_F32));
    }

    return true;
}

