#include <stdio.h>
#include "pslib.h"
#include "pmFPA.h"
#include "pmFPAMaskWeight.h"

bool pmReadoutSetMask(pmReadout *readout // Readout for which to set mask
                     )
{
    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 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(NULL, 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;
    }

    psImage *image = readout->image;    // The image pixels
    psImage *mask = readout->mask;      // The mask pixels
    if (!mask) {
        mask = psImageAlloc(image->numCols, image->numRows, PS_TYPE_U8);
        readout->mask = mask;
        psImageInit(mask, 0);
    }

    // Set up the mask
    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;
            }
        }
    }

    return true;
}

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

    psImage *image = readout->image;    // The image pixels
    psImage *weight = readout->weight;  // The weight pixels
    if (!weight) {
        weight = psImageAlloc(image->numCols, image->numRows, PS_TYPE_F32);
        readout->weight = weight;
        psImageInit(weight, 0.0);
    }

    // Set weight image to the variance in ADU = f/g + 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;
}



bool pmCellSetMaskWeight(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 |= pmReadoutSetMask(readout);
        success |= pmReadoutSetWeight(readout);
    }

    return success;
}
