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

#include "pmFPA.h"
#include "pmHDU.h"
#include "pmHDUUtils.h"
#include "pmConcepts.h"
#include "pmFPAHeader.h"

#include "pmFPARead.h"

#define MAX(x,y) ((x) > (y) ? (x) : (y))
#define MIN(x,y) ((x) < (y) ? (x) : (y))

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// File-static functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Carve a readout from the image pixels
static bool readoutCarve(pmReadout *readout, // Readout to be carved up
                         psImage **target, // Target pointer for the carved image
                         psImage *image, // Image that will be carved
                         const psRegion *trimsec, // Trim section
                         const psList *biassecs // Bias sections
                        )
{
    assert(readout);
    assert(image);
    assert(trimsec);
    assert(biassecs);

    // The image corresponding to the trim region
    if (psRegionIsNaN(*trimsec)) {
        psString regionString = psRegionToString(*trimsec);
        psError(PS_ERR_UNKNOWN, true, "Invalid trim section: %s\n", regionString);
        psFree(regionString);
        psFree(readout);
        return false;
    }
    psRegion region = psRegionSet(MAX(trimsec->x0 - readout->col0, 0), // x0
                                  MIN(trimsec->x1 - readout->col0, image->numCols), // x1
                                  MAX(trimsec->y0 - readout->row0, 0), // y0
                                  MIN(trimsec->y1 - readout->row0, image->numRows) // y1
                                 );
    if (readout->image) {
        psFree(readout->image);         // Make way!
    }
    *target = psMemIncrRefCounter(psImageSubset(image, region));

    // Get the list of overscans
    if (readout->bias->n != 0) {
        // Make way!
        psFree(readout->bias);
        readout->bias = psListAlloc(NULL);
    }
    psListIterator *iter = psListIteratorAlloc((psList*)biassecs, PS_LIST_HEAD, false); // Iterator
    psRegion *biassec = NULL;       // A BIASSEC region from the list
    while ((biassec = psListGetAndIncrement(iter))) {
        if (psRegionIsNaN(*biassec)) {
            psString regionString = psRegionToString(*biassec);
            psError(PS_ERR_IO, true, "Invalid bias section: %s\n", regionString);
            psFree(regionString);
            psFree(readout);
            return false;
        }
        psRegion region = psRegionSet(MAX(biassec->x0 - readout->col0, 0), // x0
                                      MIN(biassec->x1 - readout->col0, image->numCols), // x1
                                      MAX(biassec->y0 - readout->row0, 0), // y0
                                      MIN(biassec->y1 - readout->row0, image->numRows) // y1
                                     );
        psImage *overscan = psMemIncrRefCounter(psImageSubset(image, region));
        psListAdd(readout->bias, PS_LIST_TAIL, overscan);
        psFree(overscan);
    }
    psFree(iter);

    return true;
}

// Read a component of a readout
psImage *readoutReadComponent(psFits *fits, // FITS file from which to read
                              const psRegion *region, // Region to read
                              int readdir, // Read direction (1=rows, 2=cols)
                              int min,  // Minimum row/col number to read
                              int max,   // Maximum row/col number to read
                              int z,     // Image plane to read
                              float bad // Bad value
                             )
{
    assert(fits);
    assert(region);

    bool resize = false;                // Do we need to resize the image once read?
    psRegion toRead = psRegionSet(region->x0, region->x1, region->y0, region->y1); // Region to read
    psRegion fullSize = toRead;         // Full sized region
    if (readdir == 1) {
        if (toRead.y0 <= min) {
            toRead.y0 = min;
        } else {
            fullSize.y0 = min;
            resize = true;
        }
        if (toRead.y1 >= max) {
            toRead.y1 = max;
        } else {
            fullSize.y1 = max;
            resize = false;
        }
    } else if (readdir == 2) {
        if (toRead.x0 <= min) {
            toRead.x0 = min;
        } else {
            fullSize.x0 = min;
            resize = true;
        }
        if (toRead.x1 >= max) {
            toRead.x1 = max;
        } else {
            fullSize.x1 = max;
            resize = false;
        }
    } else {
        psAbort(__func__, "Read direction can only be 1 (rows) or 2 (cols).\n");
    }

    psTrace(__func__, 5, "Reading section [%.0f:%.0f,%.0f:%.0f]\n",
            toRead.x0, toRead.x1, toRead.y0, toRead.y1);
    psImage *image = psFitsReadImage(fits, toRead, z); // Desired pixels
    psTrace(__func__, 7, "Image is %dx%d\n", image->numCols, image->numRows);

    // XXX: We only support F32 for now
    if (image->type.type != PS_TYPE_F32) {
        psImage *temp = psImageCopy(NULL, image, PS_TYPE_F32);
        psFree(image);
        image = temp;
    }

    if (resize) {
        // For some reason, the region of interest is smaller than the number of pixels we want.
        psTrace(__func__, 5, "Resizing image to %.0fx%.0f\n",
                fullSize.x1 - fullSize.x0, fullSize.y1 - fullSize.y0);
        psImage *temp = psImageAlloc(fullSize.x1 - fullSize.x0, fullSize.y1 - fullSize.y0, image->type.type);
        psImageInit(temp, bad);
        psImageOverlaySection(temp, image, toRead.x0, toRead.y0, "=");
        psFree(image);
        image = temp;
    }

    return image;
}

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

// Read the next readout; return true if we read pixels in.
//
// Note that this doesn't put pixels in the HDU.  It is therefore NOT COMPATIBLE with pmCellWrite,
// pmChipWrite, and pmFPAWrite (or any function that uses or creates an HDU for that matter).
// Use pmReadoutWriteNext to write the data that's read by this function.
bool pmReadoutReadNext(pmReadout *readout, // Readout into which to read
                       psFits *fits,    // FITS file from which to read
                       int z,           // Readout number/plane; zero-offset indexing
                       int numScans     // The number of scans to read
                      )
{
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);
    PS_ASSERT_INT_NONNEGATIVE(z, false);
    PS_ASSERT_INT_NONNEGATIVE(numScans, false);

    // Get the HDU and read the header
    pmCell *cell = readout->parent;     // The parent cell
    if (!pmCellReadHeader(cell, fits)) {
        psError(PS_ERR_IO, false, "Unable to read header for cell!\n");
        return false;
    }

    pmHDU *hdu = pmHDUFromCell(cell);   // The HDU
    if (!hdu) {
        return false;
    }

    // Make sure we have the information we need
    if (!pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_HEADER | PM_CONCEPT_SOURCE_CAMERA |
                            PM_CONCEPT_SOURCE_DEFAULTS, true, NULL)) {
        psError(PS_ERR_IO, false, "Failed to read concepts for cell.\n");
        return false;
    }

    // Get the trim and bias sections
    bool mdok = true;                   // Status of MD lookup
    psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim sections
    if (!mdok || !trimsec || psRegionIsNaN(*trimsec)) {
        psError(PS_ERR_IO, true, "CELL.TRIMSEC is not set.\n");
        return false;
    }
    psList *biassecs = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.BIASSEC"); // Bias sections
    if (!mdok || !biassecs) {
        psError(PS_ERR_IO, true, "CELL.BIASSEC is not set.\n");
        return false;
    }
    int readdir = psMetadataLookupS32(&mdok, cell->concepts, "CELL.READDIR"); // Read direction
    if (!mdok || readdir == 0 || (readdir != 1 && readdir != 2)) {
        psError(PS_ERR_IO, true, "CELL.READDIR is not set to -1 or +1.\n");
        return false;
    }
    float bad = psMetadataLookupF32(&mdok, cell->concepts, "CELL.BAD"); // Bad level
    if (!mdok) {
        psLogMsg(__func__, PS_LOG_WARN, "CELL.BAD is not set --- assuming zero.\n");
        bad = 0.0;
    }

    // Check the third dimension
    int naxis = psMetadataLookupS32(&mdok, hdu->header, "NAXIS"); // The number of axes
    if (!mdok) {
        psError(PS_ERR_IO, true, "Unable to find NAXIS in header for extension %s\n", hdu->extname);
        return false;
    }
    if (naxis == 0) {
        // No pixels to read, as for a PHU.
        return false;
    }
    if (naxis < 2 || naxis > 3) {
        psError(PS_ERR_IO, true, "NAXIS in header of extension %s (= %d) is not valid.\n",
                hdu->extname, naxis);
        return false;
    }
    int naxis3 = 1;                     // The number of image planes
    if (naxis == 3) {
        naxis3 = psMetadataLookupS32(&mdok, hdu->header, "NAXIS3");
        if (!mdok) {
            psError(PS_ERR_IO, true, "Unable to find NAXIS3 in header for extension %s\n", hdu->extname);
            return false;
        }
    }
    if (z >= naxis3) {
        // Nothing to see here.  Move along.
        return false;
    }

    // Get the size of the image plane
    int naxis1 = psMetadataLookupS32(&mdok, hdu->header, "NAXIS1"); // The number of columns
    if (!mdok) {
        psError(PS_ERR_IO, true, "Unable to find NAXIS1 in header for extension %s\n", hdu->extname);
        return false;
    }
    int naxis2 = psMetadataLookupS32(&mdok, hdu->header, "NAXIS2"); // The number of columns
    if (!mdok) {
        psError(PS_ERR_IO, true, "Unable to find NAXIS2 in header for extension %s\n", hdu->extname);
        return false;
    }

    int maxSize;                        // Number of cols,rows in image
    if (readdir == 1) {
        maxSize = PS_MIN(naxis2, trimsec->y1);
    } else {
        maxSize = PS_MIN(naxis1, trimsec->x1);
    }

    // Read the FITS image
    int offset = 0;                     // The offset from the start of the image to where we are now
    if (readout->image) {
        if (readdir == 1) {
            // Reading rows
            readout->row0 += readout->image->numRows;
            offset = readout->row0;
        } else {
            // Reading columns
            readout->col0 += readout->image->numCols;
            offset = readout->col0;
        }
    }

    // Blow away existing data.
    // Do this before returning, so that we're not returning data from a previous read
    psFree(readout->image);
    readout->image = NULL;
    while (readout->bias->n > 0) {
        psListRemove(readout->bias, PS_LIST_HEAD);
    }

    if (offset >= maxSize) {
        // We've read everything there is
        psTrace(__func__, 7, "Read everything.\n");
        psFree(readout->image);
        return false;
    }
    int upper = PS_MIN(offset + numScans, maxSize); // The upper limit for the pixel read
    psTrace(__func__, 7, "offset=%d, upper = %d, image is %dx%d, trimsec [%.0f:%.0f,%.0f:%.0f]\n",
            offset, upper, naxis1, naxis2, trimsec->x0, trimsec->x1, trimsec->y0, trimsec->y1);

    // Get the new the trim section
    readout->image = readoutReadComponent(fits, trimsec, readdir, offset, upper, z, bad); // The image

    // Get the new bias sections
    psListIterator *biassecsIter = psListIteratorAlloc(biassecs, PS_LIST_HEAD, false); // Iterator for BIASSEC
    psRegion *biassec = NULL;           // Bias section from iteration
    while ((biassec = psListGetAndIncrement(biassecsIter))) {
        psImage *bias = readoutReadComponent(fits, biassec, readdir, offset, upper, z, bad); // The bias
        psListAdd(readout->bias, PS_LIST_TAIL, bias);
    }
    psFree(biassecsIter);

    return true;
}


// Read in the cell, and allocate the readouts
bool pmCellRead(pmCell *cell,           // Cell to read into
                psFits *fits,           // FITS file from which to read
                psDB *db                // Database handle, for "concepts" ingest
               )
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    pmHDU *hdu = pmHDUFromCell(cell);   // The HDU
    if (!hdu) {
        return false;                    // Nothing to see here; move along
    }
    if (!hdu->images && !pmHDURead(hdu, fits)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to read HDU for cell.\n");
        return false;
    }

    if (!pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
        psError(PS_ERR_UNKNOWN, false, "Failed to read concepts for cell");
        return false;
    }

    // Having read the cell, we now have to cut it up
    psRegion *trimsec = psMetadataLookupPtr(NULL, cell->concepts, "CELL.TRIMSEC");
    psList *biassecs = psMetadataLookupPtr(NULL, cell->concepts, "CELL.BIASSEC");

    // Iterate over each of the image planes
    for (int i = 0; i < hdu->images->n; i++) {
        psImage *image = hdu->images->data[i]; // The i-th plane

        // XXX: Type conversion here to support the modules, which don't have multiple type support yet
        if (image->type.type != PS_TYPE_F32) {
            psImage *temp = psImageCopy(NULL, image, PS_TYPE_F32); // Temporary image
            psFree(hdu->images->data[i]);
            hdu->images->data[i] = temp;
            image = temp;
        }

        pmReadout *readout;             // Readout into which to read
        if (cell->readouts->n > i && cell->readouts->data[i]) {
            readout = psMemIncrRefCounter(cell->readouts->data[i]);
        } else {
            readout = pmReadoutAlloc(cell);
        }
        if (!readoutCarve(readout, &readout->image, image, trimsec, biassecs)) {
            psError(PS_ERR_UNEXPECTED_NULL, false,
                    "Unable to carve readout into image and bias sections for %d the plane.", i);
            return NULL;
        }
        psFree(readout);                // Drop reference
    }

    pmCellSetDataStatus(cell, true);
    return true;
}

// Read in the component cells
bool pmChipRead(pmChip *chip,           // Chip to read into
                psFits *fits,           // FITS file from which to read
                psDB *db                // Database handle, for "concepts" ingest
               )
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *cells = chip->cells;       // Array of cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // The cell of interest
        success |= pmCellRead(cell, fits, db);
    }
    if (success) {
        if (!pmConceptsReadChip(chip, PM_CONCEPT_SOURCE_HEADER, true, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
        // XXX probably could just use chip->data_exists
        pmChipSetDataStatus(chip, true);
    }

    return success;
}

// Read in the component chips
bool pmFPARead(pmFPA *fpa,              // FPA to read into
               psFits *fits,            // FITS file from which to read
               psDB *db                 // Database handle, for "concepts" ingest
              )
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *chips = fpa->chips;        // Array of chips
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // The cell of interest
        success |= pmChipRead(chip, fits, db);
    }
    if (success) {
        if (!pmConceptsReadFPA(fpa, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
    } else {
        psError(PS_ERR_UNKNOWN, false, "Unable to read any chips in FPA");
    }

    return success;
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Reading the mask
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Read the mask into the cell, and allocate the readouts
bool pmCellReadMask(pmCell *cell,           // Cell to read into
                    psFits *fits,           // FITS file from which to read
                    psDB *db                // Database handle, for "concepts" ingest
                   )
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    pmHDU *hdu = pmHDUFromCell(cell);   // The HDU
    if (!hdu) {
        return false;                    // Nothing to see here; move along
    }
    if (!hdu->images && !pmHDURead(hdu, fits)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to read HDU for cell.\n");
        return false;
    }

    if (!pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
        psError(PS_ERR_UNKNOWN, false, "Failed to read concepts for cell");
        return false;
    }

    // Having read the cell, we now have to cut it up
    psRegion *trimsec = psMetadataLookupPtr(NULL, cell->concepts, "CELL.TRIMSEC");

    // Iterate over each of the image planes
    for (int i = 0; i < hdu->images->n; i++) {
        psImage *image = hdu->images->data[i]; // The i-th plane

        if (image->type.type != PS_TYPE_U8) {
            psImage *temp = psImageCopy(NULL, image, PS_TYPE_U8); // Temporary image
            psFree(hdu->images->data[i]);
            hdu->images->data[i] = temp;
            image = temp;
        }

        pmReadout *readout;             // Readout into which to read
        if (cell->readouts->n > i && cell->readouts->data[i]) {
            readout = psMemIncrRefCounter(cell->readouts->data[i]);
        } else {
            readout = pmReadoutAlloc(cell);
        }
        if (!readoutCarve(readout, &readout->mask, image, trimsec, NULL)) {
            psError(PS_ERR_UNEXPECTED_NULL, false,
                    "Unable to carve readout into image and bias sections for %d the plane.", i);
            return NULL;
        }
        psFree(readout);                // Drop reference
    }

    pmCellSetDataStatus(cell, true);
    return true;
}

// Read the mask into the component cells
bool pmChipReadMask(pmChip *chip,           // Chip to read into
                    psFits *fits,           // FITS file from which to read
                    psDB *db                // Database handle, for "concepts" ingest
                   )
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *cells = chip->cells;       // Array of cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // The cell of interest
        success |= pmCellReadMask(cell, fits, db);
    }
    if (success) {
        if (!pmConceptsReadChip(chip, PM_CONCEPT_SOURCE_HEADER, true, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
        // XXX probably could just use chip->data_exists
        pmChipSetDataStatus(chip, true);
    }

    return success;
}

// Read the mask into the component chips
bool pmFPAReadMask(pmFPA *fpa,              // FPA to read into
                   psFits *fits,            // FITS file from which to read
                   psDB *db                 // Database handle, for "concepts" ingest
                  )
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *chips = fpa->chips;        // Array of chips
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // The cell of interest
        success |= pmChipReadMask(chip, fits, db);
    }
    if (success) {
        if (!pmConceptsReadFPA(fpa, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
    } else {
        psError(PS_ERR_UNKNOWN, false, "Unable to read any chips in FPA");
    }

    return success;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Reading the weight map
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Read the weight map into the cell, and allocate the readouts
bool pmCellReadWeight(pmCell *cell,           // Cell to read into
                      psFits *fits,           // FITS file from which to read
                      psDB *db                // Database handle, for "concepts" ingest
                     )
{
    PS_ASSERT_PTR_NON_NULL(cell, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    pmHDU *hdu = pmHDUFromCell(cell);   // The HDU
    if (!hdu) {
        return false;                    // Nothing to see here; move along
    }
    if (!hdu->images && !pmHDURead(hdu, fits)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to read HDU for cell.\n");
        return false;
    }

    if (!pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
        psError(PS_ERR_UNKNOWN, false, "Failed to read concepts for cell");
        return false;
    }

    // Having read the cell, we now have to cut it up
    psRegion *trimsec = psMetadataLookupPtr(NULL, cell->concepts, "CELL.TRIMSEC");

    // Iterate over each of the image planes
    for (int i = 0; i < hdu->images->n; i++) {
        psImage *image = hdu->images->data[i]; // The i-th plane

        if (image->type.type != PS_TYPE_F32) {
            psImage *temp = psImageCopy(NULL, image, PS_TYPE_F32); // Temporary image
            psFree(hdu->images->data[i]);
            hdu->images->data[i] = temp;
            image = temp;
        }

        pmReadout *readout;             // Readout into which to read
        if (cell->readouts->n > i && cell->readouts->data[i]) {
            readout = psMemIncrRefCounter(cell->readouts->data[i]);
        } else {
            readout = pmReadoutAlloc(cell);
        }
        if (!readoutCarve(readout, &readout->weight, image, trimsec, NULL)) {
            psError(PS_ERR_UNEXPECTED_NULL, false,
                    "Unable to carve readout into image and bias sections for %d the plane.", i);
            return NULL;
        }
        psFree(readout);                // Drop reference
    }

    pmCellSetDataStatus(cell, true);
    return true;
}

// Read the weight map into the component cells
bool pmChipReadWeight(pmChip *chip,           // Chip to read into
                      psFits *fits,           // FITS file from which to read
                      psDB *db                // Database handle, for "concepts" ingest
                     )
{
    PS_ASSERT_PTR_NON_NULL(chip, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *cells = chip->cells;       // Array of cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // The cell of interest
        success |= pmCellReadWeight(cell, fits, db);
    }
    if (success) {
        if (!pmConceptsReadChip(chip, PM_CONCEPT_SOURCE_HEADER, true, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
        // XXX probably could just use chip->data_exists
        pmChipSetDataStatus(chip, true);
    }

    return success;
}

// Read the weight map into the component chips
bool pmFPAReadWeight(pmFPA *fpa,              // FPA to read into
                     psFits *fits,            // FITS file from which to read
                     psDB *db                 // Database handle, for "concepts" ingest
                    )
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);
    PS_ASSERT_PTR_NON_NULL(fits, false);

    bool success = false;               // Were we able to read at least one HDU?
    psArray *chips = fpa->chips;        // Array of chips
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // The cell of interest
        success |= pmChipReadWeight(chip, fits, db);
    }
    if (success) {
        if (!pmConceptsReadFPA(fpa, PM_CONCEPT_SOURCE_HEADER, true, NULL)) {
            psError(PS_ERR_IO, false, "Failed to read concepts for FPA.\n");
            return false;
        }
    } else {
        psError(PS_ERR_UNKNOWN, false, "Unable to read any chips in FPA");
    }

    return success;
}

