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

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

#include "pmHDUGenerate.h"

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

// Add cells in a chip to a list
static bool addCellsFromChip(psList *list, // List of cells
                             const pmChip *chip // The chip from which to add cells
                            )
{
    assert(list);
    assert(chip);

    psArray *cells = chip->cells;       // Array of cells
    bool result = true;                 // Result of adding cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // A cell
        if (!cell->hdu) {               // Don't add cells that have their own HDU
            result |= psListAdd(list, PS_LIST_TAIL, cell);
        }
    }

    return result;
}

// Add cells in an FPA to a list
static bool addCellsFromFPA(psList *list, // List of cells
                            const pmFPA *fpa // The FPA from which to add cells
                           )
{
    assert(list);
    assert(fpa);

    psArray *chips = fpa->chips;        // Array of chips
    bool result = true;                 // Result of adding cells
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // A chip
        if (! chip->hdu) {              // Don't add chips that have their own HDU
            result |= addCellsFromChip(list, chip);
        }
    }

    return result;
}

// Get the maximum extent of the HDU from the trimsec and biassecs
static bool sizeHDU(int *xSize, int *ySize, // Size of HDU
                    psList *cells       // List of cells
                   )
{
    psListIterator *cellsIter = psListIteratorAlloc(cells, PS_LIST_HEAD, false); // Iterator for cells
    pmCell *cell = NULL;                // The cell from iteration
    bool mdok = true;                   // Status of MD lookup
    *xSize = 0;
    *ySize = 0;
    while ((cell = psListGetAndIncrement(cellsIter))) {
        psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
        if (mdok && trimsec && !psRegionIsNaN(*trimsec)) {
            *xSize = PS_MAX(trimsec->x1, *xSize);
            *ySize = PS_MAX(trimsec->y1, *ySize);
        } else {
            psFree(cellsIter);
            return false;
        }
        psList *biassecs = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.BIASSEC"); // Bias sections
        if (mdok && biassecs) {
            psListIterator *biassecsIter = psListIteratorAlloc(biassecs, PS_LIST_HEAD, false); // Iterator
            psRegion *biassec = NULL;   // The bias section
            while ((biassec = psListGetAndIncrement(biassecsIter))) {
                if (!psRegionIsNaN(*biassec)) {
                    *xSize = PS_MAX(biassec->x1, *xSize);
                    *ySize = PS_MAX(biassec->y1, *ySize);
                } else {
                    psFree(biassecsIter);
                    psFree(cellsIter);
                    return false;
                }
            }
            psFree(biassecsIter);
        }
    }
    psFree(cellsIter);

    return (*xSize != 0 && *ySize != 0);
}


static psRegion *sectionForImage(int *position, // Position on the output image, updated
                                 const psImage *image, // Image containing the sizes and offsets
                                 int readdir // Read direction, 1=rows, 2=cols
                                )
{
    psRegion *region = NULL;            // The region to return
    switch (readdir) {
    case 1:                           // Read direction is rows
        region = psRegionAlloc(*position, *position + image->numCols, image->row0,
                               image->row0 + image->numRows);
        *position += image->numCols;
        break;
    case 2:                           // Read direction is columns
        region = psRegionAlloc(image->col0, image->col0 + image->numCols, *position,
                               *position + image->numRows);
        *position += image->numRows;
        break;
    default:
        psAbort(__func__, "Shouldn't ever get here!\n");
    }

    return region;
}

static bool doBiasSections(int *position, // Position on the output image, updated
                           int *readdir,// Read direction for cells
                           pmCell *cell // Cell
                          )
{
    psMetadataItem *biassecItem = psMetadataLookup(cell->concepts, "CELL.BIASSEC"); // Bias sections
    if (!biassecItem) {
        psLogMsg(__func__, PS_LOG_WARN, "CELL.BIASSEC has not been initialised in cell --- "
                 "ignored.\n");
        return false;
    }
    psFree(biassecItem->data.V);        // Blow away the old list
    psList *biassecs = psListAlloc(NULL);
    biassecItem->data.V = biassecs;

    bool mdok = true;                   // Status of MD lookup
    int cellreaddir = psMetadataLookupS32(&mdok, cell->concepts, "CELL.READDIR"); // Read direction
    if (!mdok || (cellreaddir != 1 && cellreaddir != 2)) {
        // Probably unnecessary, but just in case....
        psLogMsg(__func__, PS_LOG_WARN, "CELL.READDIR is not set in cell --- ignored.\n");
        return false;
    }
    if (*readdir == 0) {
        *readdir = cellreaddir;
    } else if (*readdir != cellreaddir) {
        psLogMsg(__func__, PS_LOG_WARN, "CELL.READDIR does not match read direction for HDU --- ignored.\n");
        return false;
    }

    pmReadout *readout = cell->readouts->data[0]; // The first readout, as representative
    psList *biases = readout->bias; // The bias images from the readout

    psListIterator *biasIter = psListIteratorAlloc(biases, PS_LIST_HEAD, true); // Iterator for biases
    psImage *bias = NULL;       // Bias image from iteration
    while ((bias = psListGetAndIncrement(biasIter))) {
        // Construct a region
        psRegion *biassec = sectionForImage(position, bias, *readdir);
        psListAdd(biassecs, PS_LIST_TAIL, biassec);
        psFree(biassec);        // Drop reference
    }
    psFree(biasIter);

    return true;
}

// Generate CELL.TRIMSEC and CELL.BIASSEC for the cells
static bool generateTrimBias(psList *cells // List of cells below the HDU
                            )
{
    pmCell *cell = NULL;                // Cell from iteration
    int numCells = cells->n;            // Number of cells
    int cellNum = 0;                    // The cell number
    int position = 0;                   // Position on the image
    bool mdok = true;                   // Status of MD lookup
    int readdir = 0;                    // Read direction (1=rows, 2=cols)

    // First run through to do the LHS biases
    psListIterator *cellsIter = psListIteratorAlloc(cells, PS_LIST_HEAD, false); // Iterator for cells
    bool done = false;                  // Done with iteration (due to being halfway through)?
    while ((cell = psListGetAndIncrement(cellsIter)) && !done) {
        if (cellNum <= numCells/2 - 1) {
            doBiasSections(&position, &readdir, cell);
            cellNum++;
        } else {
            done = true;
        }
    }

    // Second run through to do the trim sections
    psListIteratorSet(cellsIter, PS_LIST_HEAD);
    while ((cell = psListGetAndIncrement(cellsIter))) {
        psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
        if (!mdok || !trimsec) {
            psLogMsg(__func__, PS_LOG_WARN, "CELL.TRIMSEC has not been initialised in cell --- "
                     "ignored.\n");
            continue;
        }

        int cellreaddir = psMetadataLookupS32(&mdok, cell->concepts, "CELL.READDIR"); // Read direction
        if (!mdok || (cellreaddir != 1 && cellreaddir != 2)) {
            // Probably unnecessary, but just in case....
            psLogMsg(__func__, PS_LOG_WARN, "CELL.READDIR is not set in cell --- ignored.\n");
            continue;
        }
        if (readdir == 0 && mdok && cellreaddir != 0) {
            readdir = cellreaddir;
        } else if (readdir != cellreaddir) {
            psLogMsg(__func__, PS_LOG_WARN, "CELL.READDIR for cells within the HDU do not match!\n");
            cellreaddir = readdir;
        }

        pmReadout *readout = cell->readouts->data[0]; // The first readout, as representative
        psImage *image = readout->image;// The proper image
        if (!image) {
            continue;
        }
        psFree(trimsec);                // Drop old reference
        trimsec = sectionForImage(&position, image, cellreaddir); // New reference
    }

    // A final run through to do the RHS biases
    psListIteratorSet(cellsIter, cellNum);
    while ((cell = psListGetAndIncrement(cellsIter))) {
        doBiasSections(&position, &readdir, cell);
    }

    // Clean up
    psFree(cellsIter);

    return (position > 0);
}

// Generate the HDU, given a list of cells below that HDU.  This is the main engine function, that does all
// the work.
static bool generateHDU(pmHDU *hdu,     // HDU to generate
                        psList *cells   // List of cells
                       )
{
    // Check the number of readouts is consistent within the HDU
    int numReadouts = -1;               // Number of readouts
    psElemType type = 0;                // Type of readout images
    {
        psListIterator *iter = psListIteratorAlloc(cells, PS_LIST_HEAD, false); // Iterator for cells
        pmCell *cell = NULL;                // The cell from iteration
        while ((cell = psListGetAndIncrement(iter)))
        {
            psArray *readouts = cell->readouts;
            if (numReadouts == -1) {
                numReadouts = readouts->n;
            } else if (readouts->n != numReadouts) {
                psError(PS_ERR_IO, true, "Number of readouts doesn't match: %d vs %d\n", readouts->n,
                        numReadouts);
                return false;
            }
            for (int i = 0; i < numReadouts; i++) {
                pmReadout *readout = readouts->data[i]; // The readout
                if (!readout || !readout->image) {
                    continue;
                }
                psElemType imageType = readout->image->type.type; // Type for this image
                if (type == 0) {
                    type = imageType;
                } else if (type != imageType) {
                    psLogMsg(__func__, PS_LOG_WARN, "Images within the HDU are of different types "
                             "(%x vs %x) --- promoting\n", type, imageType);
                    type = PS_MAX(type, imageType);
                }
            }
        }
        psFree(iter);
    }
    if (numReadouts == 0 || type == 0) {
        // Nothing from which to create an HDU
        psFree(cells);
        return false;
    }

    // Get the size of the HDU, either from existing trimsec and biassec, or generate these and try again
    int xSize = 0, ySize = 0;           // Size of HDU
    if (!sizeHDU(&xSize, &ySize, cells) && !(generateTrimBias(cells) && sizeHDU(&xSize, &ySize, cells))) {
        psError(PS_ERR_IO, true, "Unable to determine size of HDU!\n");
        return false;
    }

    // Generate the HDU
    hdu->images = psArrayAlloc(numReadouts);
    hdu->images->n = numReadouts;
    for (int i = 0; i < numReadouts; i++) {
        psImage *image = psImageAlloc(xSize, ySize, type);
        psImageInit(image, 0.0);
        hdu->images->data[i] = image;
    }

    // Insert the pixels into the HDU
    {
        psListIterator *iter = psListIteratorAlloc(cells, PS_LIST_HEAD, false); // Iterator for cells
        pmCell *cell = NULL;           // The cell from iteration
        bool mdok = true;               // Result of MD lookup
        while ((cell = psListGetAndIncrement(iter)))
        {
            psRegion *trimsec = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.TRIMSEC"); // Trim section
            if (!mdok || !trimsec) {
                psAbort(__func__, "Shouldn't ever get here --- CELL.TRIMSEC should have been set above.\n");
            }
            psList *biassecs = psMetadataLookupPtr(&mdok, cell->concepts, "CELL.BIASSEC"); // Bias secionts
            if (!mdok || !biassecs) {
                psAbort(__func__, "Shouldn't ever get here --- CELL.BIASSEC should have been set above.\n");
            }
            psListIterator *biassecsIter = psListIteratorAlloc(biassecs, PS_LIST_HEAD, false); // Iterator

            psArray *readouts = cell->readouts; // Array of readouts
            psArray *hduImages = hdu->images; // Array of images in the HDU
            for (int i = 0; i < readouts->n; i++) {
                pmReadout *readout = readouts->data[i]; // The readout of interest
                if (!readout) {
                    continue;
                }
                psImage *image = readout->image; // The image pixels
                psImage *hduImage = hduImages->data[i]; // The HDU image of interest

                if (image->numCols != trimsec->x1 - trimsec->x0 ||
                        image->numRows != trimsec->y1 - trimsec->y0) {
                    psString trimsecString = psRegionToString(*trimsec);
                    psLogMsg(__func__, PS_LOG_WARN, "Image size (%dx%d) does not match CELL.TRIMSEC (%s).\n",
                             image->numCols, image->numRows, trimsecString);
                    psFree(trimsecString);
                }
                psImageOverlaySection(hduImage, image, trimsec->x0, trimsec->y0, "=");

                // Reference the HDU version, so that subsequent changes will touch the HDU
                psFree(image);
                readout->image = psImageSubset(hduImage, *trimsec);

                if (biassecs->n != readout->bias->n) {
                    psLogMsg(__func__, PS_LOG_WARN, "Number of bias sections (%d) and number of biases (%d)"
                             " do not match.\n", biassecs->n, readout->bias->n);
                }
                psListIterator *biasIter = psListIteratorAlloc(readout->bias, PS_LIST_HEAD, false); // Iteratr
                psImage *bias = NULL;   // Bias image, from iteration
                psListIteratorSet(biassecsIter, PS_LIST_HEAD);
                psRegion *biassec = NULL; // Bias region, from iteration
                psList *newBias = psListAlloc(NULL); // New list of bias images
                while ((bias = psListGetAndIncrement(biasIter)) &&
                        (biassec = psListGetAndIncrement(biassecsIter))) {
                    if (bias->numCols != biassec->x1 - biassec->x0 ||
                            bias->numRows != biassec->y1 - biassec->y0) {
                        psString biassecString = psRegionToString(*biassec);
                        psLogMsg(__func__, PS_LOG_WARN, "Bias size (%dx%d) does not match CELL.BIASSEC (%s)."
                                 "\n", bias->numCols, bias->numRows, biassecString);
                        psFree(biassecString);
                    }

                    psImageOverlaySection(hduImage, bias, biassec->x0, biassec->y0, "=");

                    // Reference the HDU version, so that subsequent changes will touch the HDU
                    bias = psImageSubset(hduImage, *biassec);
                    psListAdd(newBias, PS_LIST_TAIL, bias);
                }
                psFree(biasIter);

                // Add on the new list of bias images
                psFree(readout->bias);
                readout->bias = newBias;
            }
            psFree(biassecsIter);
        } // Iterating over cells within the HDU
        psFree(iter);
    }
    psFree(cells);

    return true;
}

// Return the level that an extension applies to
static pmFPALevel extensionLevel(pmHDU *hdu // HDU to check
                                )
{
    if (!hdu->format) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "HDU does not have a camera format.\n");
        return PM_FPA_LEVEL_NONE;
    }
    bool mdok = true;                   // Status of MD lookup
    psMetadata *file = psMetadataLookupMD(&mdok, hdu->format, "FILE"); // File information for camera format
    if (!mdok || !file) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Can't file FILE information for camera format "
                "configuration.\n");
        return PM_FPA_LEVEL_NONE;
    }
    psString extensions = psMetadataLookupStr(&mdok, file, "EXTENSIONS"); // Where the HDUs are
    if (!mdok || !extensions || strlen(extensions) == 0) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Can't find EXTENSIONS in the FILE information of the camera "
                "format configuration.\n");
        return PM_FPA_LEVEL_NONE;
    }
    if (strcasecmp(extensions, "CELL") == 0) {
        return PM_FPA_LEVEL_CELL;
    }
    if (strcasecmp(extensions, "CHIP") == 0) {
        return PM_FPA_LEVEL_CHIP;
    }
    if (strcasecmp(extensions, "FPA") == 0) {
        return PM_FPA_LEVEL_FPA;
    }
    if (strcasecmp(extensions, "NONE") == 0) {
        return PM_FPA_LEVEL_NONE;
    }
    // No idea what it is
    psError(PS_ERR_IO, true, "EXTENSIONS (%s) in FILE information is not FPA, CHIP, CELL or NONE.\n",
            extensions);
    return PM_FPA_LEVEL_NONE;
}

// Generate HDU for cells belonging to a chip --- just an iterator
static bool generateForCells(pmChip *chip // Chip for which to generate HDUs
                            )
{
    psArray *cells = chip->cells;       // Array of cells
    bool status = true;                 // Status of HDU generation
    for (long i = 0; i < cells->n; i++) {
        status |= pmHDUGenerateForCell(cells->data[i]);
    }
    return status;
}

// Generate HDU for chips belonging to an FPA --- just an iterator
static bool generateForChips(pmFPA *fpa // FPA for which to generate HDUs
                            )
{
    psArray *chips = fpa->chips;        // Array of chips
    bool status = true;                 // Status of HDU generation
    for (long i = 0; i < chips->n; i++) {
        status |= pmHDUGenerateForChip(chips->data[i]);
    }
    return status;
}


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

// Generate an HDU from a cell with pixels
bool pmHDUGenerateForCell(pmCell *cell  // The cell
                         )
{
    PS_ASSERT_PTR_NON_NULL(cell, false);

    // Get the HDU and a list of cells below it
    pmHDU *hdu = pmHDUFromCell(cell); // The HDU in the cell
    if (!hdu) {
        psError(PS_ERR_UNEXPECTED_NULL, true, "Can't find an HDU for cell.\n");
        return false;
    }
    if (hdu->images) {
        // It's already here!
        return true;
    }

    pmFPALevel extLevel = extensionLevel(hdu);
    switch (extLevel) {
    case PM_FPA_LEVEL_NONE:
    case PM_FPA_LEVEL_CELL: {
            psList *cells = psListAlloc(NULL); // List of cells below the HDU

            if (cell->hdu) {
                psListAdd(cells, PS_LIST_TAIL, cell);
            } else {
                pmChip *chip = cell->parent;    // The parent chip
                if (chip->hdu) {
                    addCellsFromChip(cells, chip);
                } else {
                    pmFPA *fpa = chip->parent;  // The parent FPA
                    if (fpa->hdu) {
                        addCellsFromFPA(cells, fpa);
                    }
                }
            }
            if (cells->n == 0) {
                // Nothing to do
                return true;
            }

            return generateHDU(hdu, cells);
        }
    case PM_FPA_LEVEL_CHIP:
    case PM_FPA_LEVEL_FPA:
        return pmHDUGenerateForChip(cell->parent);
    default:
        psAbort(__func__, "Shouldn't ever get here.\n");
    }
    return false;
}

// Generate an HDU from a chip with pixels
bool pmHDUGenerateForChip(pmChip *chip  // The chip
                         )
{
    PS_ASSERT_PTR_NON_NULL(chip, false);

    // Get the HDU and a list of cells below it
    pmHDU *hdu = pmHDUFromChip(chip);   // The HDU in the chip
    if (!hdu) {
        // Nothing here; need to look further down
        return generateForCells(chip);
    }
    if (hdu->images) {
        // It's already here!
        return true;
    }

    pmFPALevel extLevel = extensionLevel(hdu);
    switch (extLevel) {
    case PM_FPA_LEVEL_NONE:
    case PM_FPA_LEVEL_CELL:
        // Work on lower levels
        return generateForCells(chip);
    case PM_FPA_LEVEL_CHIP: {
            // Work on this level
            psList *cells = psListAlloc(NULL);  // List of cells below the HDU
            if (chip->hdu) {
                addCellsFromChip(cells, chip);
            } else {
                pmFPA *fpa = chip->parent;  // The parent FPA
                if (fpa->hdu) {
                    addCellsFromFPA(cells, fpa);
                }
            }
            if (cells->n == 0) {
                // Nothing to do
                return true;
            }

            return generateHDU(hdu, cells);
        }
    case PM_FPA_LEVEL_FPA:
        // Work on higher levels
        return pmHDUGenerateForFPA(chip->parent);
    default:
        psAbort(__func__, "Shouldn't ever get here.\n");
    }
    return false;
}


// Generate an HDU from an FPA with pixels
bool pmHDUGenerateForFPA(pmFPA *fpa     // The fpa
                        )
{
    PS_ASSERT_PTR_NON_NULL(fpa, false);

    // Get the HDU and a list of cells below it
    pmHDU *hdu = pmHDUFromFPA(fpa);     // The HDU in the FPA
    if (!hdu) {
        // Nothing here; need to look further down
        return generateForChips(fpa);
    }
    if (hdu->images) {
        // It's already here!
        return true;
    }

    pmFPALevel extLevel = extensionLevel(hdu);
    switch (extLevel) {
    case PM_FPA_LEVEL_NONE:
    case PM_FPA_LEVEL_CELL:
    case PM_FPA_LEVEL_CHIP:
        // Work on lower levels
        return generateForChips(fpa);
    case PM_FPA_LEVEL_FPA: {
            // Work on this level
            psList *cells = psListAlloc(NULL); // List of cells below the HDU
            if (fpa->hdu) {
                addCellsFromFPA(cells, fpa);
            }
            if (cells->n == 0) {
                // Nothing to do
                return true;
            }

            return generateHDU(hdu, cells);
        }
    default:
        psAbort(__func__, "Shouldn't ever get here.\n");
    }
    return false;
}
