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

#include "pslib.h"
#include "pmFPA.h"
#include "pmChipMosaic.h"

#define MEM_LEAKS 0

// Compare a value with a maximum and minimum
#define COMPARE(value,min,max) \
if ((value) < (min)) { \
    (min) = (value); \
} \
if ((value) > (max)) { \
    (max) = (value); \
}

// Mosaic multiple images, with flips, binning and offsets
psImage *p_pmImageMosaic(const psArray *source, // Images to splice in
                         const psVector *xFlip, const psVector *yFlip, // Need to flip x and y?
                         const psVector *xBinSource, const psVector *yBinSource, // Binning in x and y of
                         // source images
                         int xBinTarget, int yBinTarget, // Binning in x and y of target images
                         const psVector *x0, const psVector *y0 // Offsets for source images on target
                        )
{
    assert(source);
    assert(xFlip && xFlip->type.type == PS_TYPE_U8);
    assert(yFlip && yFlip->type.type == PS_TYPE_U8);
    assert(xBinSource && xBinSource->type.type == PS_TYPE_S32);
    assert(yBinSource && yBinSource->type.type == PS_TYPE_S32);
    assert(x0 && x0->type.type == PS_TYPE_S32);
    assert(y0 && y0->type.type == PS_TYPE_S32);

    // Get the maximum extent of the mosaic image
    int xMin = INT_MAX;
    int xMax = - INT_MAX;
    int yMin = INT_MAX;
    int yMax = - INT_MAX;
    psElemType type = 0;
    for (int i = 0; i < source->n; i++) {
        psImage *image = source->data[i]; // The image of interest
        if (!image) {
            continue;
        }

        // Only implemented for F32 and U8 images so far.
        assert(image->type.type == PS_TYPE_F32 || image->type.type == PS_TYPE_U8);
        // All input types must be the same
        if (type == 0) {
            type = image->type.type;
        }
        assert(type == image->type.type);

        // Size of cell in x and y
        int xParity = xFlip->data.U8[i] ? -1 : 1;
        int yParity = yFlip->data.U8[i] ? -1 : 1;
        psTrace(__func__, 5, "Extent of cell %d: %d -> %d , %d -> %d\n", i, x0->data.S32[i],
                x0->data.S32[i] + xParity * xBinSource->data.S32[i] * image->numCols, y0->data.S32[i],
                y0->data.S32[i] + yParity * yBinSource->data.S32[i] * image->numRows);

        COMPARE(x0->data.S32[i], xMin, xMax);
        COMPARE(y0->data.S32[i], yMin, yMax);
        // Subtract the parity to get the inclusive limit (not exclusive)
        COMPARE(x0->data.S32[i] + xParity * xBinSource->data.S32[i] * image->numCols - xParity, xMin, xMax);
        COMPARE(y0->data.S32[i] + yParity * yBinSource->data.S32[i] * image->numRows - yParity, yMin, yMax);
    }

    // Set up the image
    // Since both upper and lower values are inclusive, we need to add one to the size
    float xSize = (float)(xMax - xMin + 1) / (float)xBinTarget;
    if (xSize - (int)xSize > 0) {
        xSize += 1;
    }
    float ySize = (float)(yMax - yMin + 1) / (float)yBinTarget;
    if (ySize - (int)ySize > 0) {
        ySize += 1;
    }

    psTrace(__func__, 3, "Spliced image will be %dx%d\n", (int)xSize, (int)ySize);
    psImage *mosaic = psImageAlloc((int)xSize, (int)ySize, type); // The mosaic image
    psImageInit(mosaic, 0.0);

    // Next pass through the images to do the mosaicking
    for (int i = 0; i < source->n; i++) {
        psImage *image = source->data[i]; // The image of interest
        if (!image) {
            continue;
        }
        if (xBinSource->data.S32[i] == xBinTarget && yBinSource->data.S32[i] == yBinTarget &&
                xFlip->data.U8[i] == 0 && yFlip->data.U8[i] == 0) {
            // Let someone else do the hard work; useful to test psImageOverlaySection if no other reason
            psImageOverlaySection(mosaic, image, x0->data.S32[i], y0->data.S32[i], "+");
        } else {
            // We have to do the hard work ourself
            for (int y = 0; y < image->numRows; y++) {
                int yParity = yFlip->data.U8[i] ? -1 : 1;
                float yTargetBase = (y0->data.S32[i] + yParity * yBinSource->data.S32[i] * y) / yBinTarget;
                for (int x = 0; x < image->numCols; x++) {
                    int xParity = xFlip->data.U8[i] ? -1 : 1;
                    float xTargetBase = (x0->data.S32[i] + xParity * xBinSource->data.S32[i] * x) /
                                        xBinTarget;

                    // In case the original image is binned but the mosaic is not, we need to fill in the
                    // values in the mosaic.
                    #define FILL_IN(TYPE)                                                                    \
                    for (int j = 0; j < yBinSource->data.S32[i]; j++) {                                      \
                        int yTarget = (int)(yTargetBase + yParity * (float)j / (float)yBinTarget);           \
                        for (int i = 0; i < xBinSource->data.S32[i]; i++) {                                  \
                            int xTarget = (int)(xTargetBase + xParity * (float)i / (float)xBinTarget);       \
                            mosaic->data.TYPE[yTarget][xTarget] += image->data.TYPE[y][x];                   \
                        }                                                                                    \
                    }

                    switch (type) {
                    case PS_TYPE_F32:
                        FILL_IN(F32);
                        break;
                    case PS_TYPE_U8:
                        FILL_IN(U8);
                        break;
                    default:
                        psAbort(__func__, "Should never get here.\n");
                    }

                }
            } // Iterating over input image
        }
    }

    return mosaic;
}


// Set the concepts in the new cell, based on the values in the old one
static bool cellConcepts(pmCell *target,// Target cell
                         psArray *sources, // Source cells
                         int xBin, int yBin // Binning
                        )
{
    bool success = true;                // Result of setting everything
    float gain       = 0.0;             // Gain
    float readnoise  = 0.0;             // Read noise
    float saturation = 0.0;             // Saturation level
    float bad        = 0.0;             // Bad level
    float exposure   = 0.0;             // Exposure time
    float darktime   = 0.0;             // Dark time
    double time      = 0.0;             // Time of observation
    psTimeType timeSys = 0;             // Time system

    int nCells = 0;                     // Number of cells;
    for (int i = 0; i < sources->n; i++) {
        pmCell *cell = sources->data[i];// The cell of interest
        if (!cell) {
            continue;
        }

        nCells++;
        gain       += psMetadataLookupF32(NULL, cell->concepts, "CELL.GAIN");
        readnoise  += psMetadataLookupF32(NULL, cell->concepts, "CELL.READNOISE");
        saturation += psMetadataLookupF32(NULL, cell->concepts, "CELL.SATURATION");
        bad        += psMetadataLookupF32(NULL, cell->concepts, "CELL.BAD");
        exposure   += psMetadataLookupF32(NULL, cell->concepts, "CELL.EXPOSURE");
        darktime   += psMetadataLookupF32(NULL, cell->concepts, "CELL.DARKTIME");
        time       += psTimeToMJD(psMetadataLookupPtr(NULL, cell->concepts, "CELL.TIME"));
        if (i == 0) {
            timeSys = psMetadataLookupS32(NULL, cell->concepts, "CELL.TIMESYS");
        } else if (timeSys != psMetadataLookupS32(NULL, cell->concepts, "CELL.TIMESYS")) {
            psLogMsg(__func__, PS_LOG_ERROR, "Differing time systems in use: %d vs %d\n", timeSys,
                     psMetadataLookupS32(NULL, cell->concepts, "CELL.TIMESYS"));
            success = false;
        }
    }
    gain       /= (float)nCells;
    readnoise  /= (float)nCells;
    saturation /= (float)nCells;
    bad        /= (float)nCells;
    exposure   /= (float)nCells;
    darktime   /= (float)nCells;
    psTime *timePtr = psTimeFromMJD(time/(double)nCells);
    timePtr = psTimeConvert(timePtr, timeSys);

    // XXX *REALLY* need a generic "concept update" function that handles the type and comments transparently.
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.GAIN", PS_META_REPLACE, "Gain (e/ADU)", gain);
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.READNOISE", PS_META_REPLACE, "Read noise (e)", readnoise);
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.SATURATION", PS_META_REPLACE, "Saturation level (ADU)", saturation);
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.BAD", PS_META_REPLACE, "Bad level (ADU)", bad);
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.EXPOSURE", PS_META_REPLACE, "Exposure time (sec)", exposure);
    psMetadataAddF32(target->concepts, PS_LIST_TAIL, "CELL.DARKTIME", PS_META_REPLACE, "Time since last CCD flush (sec)", darktime);
    psMetadataAddPtr(target->concepts, PS_LIST_TAIL, "CELL.TIME", PS_DATA_TIME | PS_META_REPLACE, "Time of observation", timePtr);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.TIMESYS", PS_META_REPLACE, "Time system", timeSys);
    psFree(timePtr);

    // Now fill in the ones I know by other means
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.X0", PS_META_REPLACE, "Position of (0,0) on the chip", 0);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.Y0", PS_META_REPLACE, "Position of (0,0) on the chip", 0);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.XPARITY", PS_META_REPLACE, "Orientation in x compared to the rest of the FPA", 1);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.YPARITY", PS_META_REPLACE, "Orientation in x compared to the rest of the FPA", 1);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.XBIN", PS_META_REPLACE, "Binning in x", xBin);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.YBIN", PS_META_REPLACE, "Binning in x", yBin);
    psMetadataAddS32(target->concepts, PS_LIST_TAIL, "CELL.READDIR", PS_META_REPLACE, "Read direction (faked)", 1);
    psRegion *trimsec = psMetadataLookupPtr(NULL, target->concepts, "CELL.TRIMSEC");
    trimsec->x0 = trimsec->x1 = trimsec->y0 = trimsec->y1 = 0.0;

    return success;
}



// Mosaic a chip together into a single image
int pmChipMosaic(pmChip *chip,// Chip to mosaic
                 int xBinChip, int yBinChip // Binning of mosaic image in x and y
                )
{

    psArray *cells = chip->cells;       // The array of cells
    psArray *images = psArrayAlloc(cells->n); // Array of images that will be mosaicked
    psArray *weights = psArrayAlloc(cells->n); // Array of weight images to be mosaicked
    psArray *masks = psArrayAlloc(cells->n); // Array of mask images to be mosaicked
    psVector *x0 = psVectorAlloc(cells->n, PS_TYPE_S32); // Origin x coordinates
    psVector *y0 = psVectorAlloc(cells->n, PS_TYPE_S32); // Origin y coordinates
    psVector *xBin = psVectorAlloc(cells->n, PS_TYPE_S32); // Binning in x
    psVector *yBin = psVectorAlloc(cells->n, PS_TYPE_S32); // Binning in y
    psVector *xFlip = psVectorAlloc(cells->n, PS_TYPE_U8); // Flip in x?
    psVector *yFlip = psVectorAlloc(cells->n, PS_TYPE_U8); // Flip in y?

    // Set up the required inputs
    psTrace(__func__, 1, "Mosaicking %d cells...\n", cells->n);
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // The cell of interest
        if (!cell) {
            continue;
        }
        x0->data.S32[i] = psMetadataLookupS32(NULL, cell->concepts, "CELL.X0");
        y0->data.S32[i] = psMetadataLookupS32(NULL, cell->concepts, "CELL.Y0");
        psTrace(__func__, 5, "Cell %d: x0=%d y0=%d\n", i, x0->data.S32[i], y0->data.S32[i]);
        xBin->data.S32[i] = psMetadataLookupS32(NULL, cell->concepts, "CELL.XBIN");
        yBin->data.S32[i] = psMetadataLookupS32(NULL, cell->concepts, "CELL.XBIN");
        int xParity = psMetadataLookupS32(NULL, cell->concepts, "CELL.XPARITY");
        int yParity = psMetadataLookupS32(NULL, cell->concepts, "CELL.YPARITY");
        if (xParity == 1) {
            xFlip->data.U8[i] = 0;
        } else if (xParity == -1) {
            xFlip->data.U8[i] = 1;
        } else {
            psLogMsg(__func__, PS_LOG_WARN, "The x parity of cell %d is not +/- 1 (it's %d) --- "
                     "assuming +1.\n", i, xParity);
            xFlip->data.U8[i] = 0;
        }
        if (yParity == 1) {
            yFlip->data.U8[i] = 0;
        } else if (yParity == -1) {
            yFlip->data.U8[i] = 1;
        } else {
            psLogMsg(__func__, PS_LOG_WARN, "The y parity of cell %d is not +/- 1 (it's %d) --- "
                     "assuming +1.\n", i, yParity);
            yFlip->data.U8[i] = 0;
        }

        psArray *readouts = cell->readouts; // The array of readouts
        if (readouts->n > 1) {
            psLogMsg(__func__, PS_LOG_WARN, "Cell %d contains more than one readout --- only the first will "
                     "be mosaicked.\n", i);
        }
        pmReadout *readout = readouts->data[0]; // The only readout we'll bother with

        // The images to put into the mosaic
        images->data[i]  = psMemIncrRefCounter(readout->image);
        weights->data[i] = psMemIncrRefCounter(readout->weight);
        masks->data[i]   = psMemIncrRefCounter(readout->mask);
    }
    // Mosaic the images together and we're done
    psImage *image = p_pmImageMosaic(images, xFlip, yFlip, xBin, yBin, xBinChip, yBinChip, x0, y0);
    psImage *weight = p_pmImageMosaic(weights, xFlip, yFlip, xBin, yBin, xBinChip, yBinChip, x0, y0);
    psImage *mask = p_pmImageMosaic(masks, xFlip, yFlip, xBin, yBin, xBinChip, yBinChip, x0, y0);

    // Clean up
    psFree(x0);
    psFree(y0);
    psFree(xBin);
    psFree(yBin);
    psFree(xFlip);
    psFree(yFlip);
    psFree(images);
    psFree(weights);
    psFree(masks);
    int nCells = cells->n;

    // Fix up the HDU
    if (chip->parent->hdu) {
        psLogMsg(__func__, PS_LOG_WARN, "The original format has the entire FPA in a single extension.  "
                 "The FPA hierarchy may be invalid following the pmChipMosaic.\n");
    } else {
        if (! chip->hdu) {
            psString chipName = psMetadataLookupStr(NULL, chip->concepts, "CHIP.NAME");
            chip->hdu = p_pmHDUAlloc(chipName);
        }
        p_pmHDU *hdu = chip->hdu;
        psArrayElementsFree(hdu->images);
        psArrayElementsFree(hdu->weights);
        psArrayElementsFree(hdu->masks);
        hdu->images  = psArrayRealloc(hdu->images,1);
        hdu->weights = psArrayRealloc(hdu->weights, 1);
        hdu->masks   = psArrayRealloc(hdu->masks, 1);
        hdu->images->data[0]  = image;
        hdu->weights->data[0] = weight;
        hdu->masks->data[0]   = mask;
        psMetadataAddS32(hdu->header, PS_LIST_TAIL, "NAXIS1", PS_META_REPLACE, "Number of columns", image->numCols);
        psMetadataAddS32(hdu->header, PS_LIST_TAIL, "NAXIS2", PS_META_REPLACE, "Number of rows", image->numRows);
    }

    // Chop off all the component cells, and construct a new one
    pmCell *newCell = pmCellAlloc(NULL, NULL, __func__); // New cell
    cellConcepts(newCell, cells, xBinChip, yBinChip);
    pmChipFreeCells(chip);
    // Have to put in the new cell manually, since we didn't want to put it in before blowing the cells away.
    newCell->parent = chip;
    psArrayAdd(chip->cells, 1, newCell);
    newCell->exists = true;
    newCell->process = true;

    // Now make a new readout to go in the new cell
    pmReadout *newReadout = pmReadoutAlloc(newCell); // New readout
    // Want the readouts to contain a subimage, but that subimage is the whole image.
    // This preserves the relationship there was before, where freeing the parent frees the child.
    psRegion entire = {0.0, 0.0, 0.0, 0.0};
    newReadout->image = psMemIncrRefCounter(psImageSubset(image, entire));
    newReadout->weight = psMemIncrRefCounter(psImageSubset(weight, entire));
    newReadout->mask = psMemIncrRefCounter(psImageSubset(mask, entire));
    // Drop references
    psFree(newReadout);
    psFree(newCell);

    // Well, we've stuffed around with the camera configuration, so it's no longer valid...
    #if 0

    psFree(chip->parent->camera);
    chip->parent->camera = NULL;
    #endif

    return nCells;
}


int pmFPAMosaicCells(pmFPA *fpa,        // FPA
                     int xBinChip, int yBinChip // Binning of mosaic image in x and y
                    )
{
    assert(fpa);

    int numChips = 0;
    psArray *chips = fpa->chips;        // Component chips
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // The chip of interest
        if (! chip || ! chip->exists || ! chip->process) {
            continue;
        }

        numChips++;
        pmChipMosaic(chip, xBinChip, yBinChip); // Mosaic of cells within the chip

    }

    return numChips;

}
