#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

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

#include "ppMerge.h"
#include "ppMergeScaleZero.h"

// Get the scale and zero for each chip of each input
bool ppMergeScaleZero(psImage **scales, // The scales for each integration/cell
                      psImage **zeros, // The zeroes for each integration/cell
                      psArray **shutters, // The shutter correction data for each cell
                      ppMergeData *data,// The data
                      const ppMergeOptions *options, // The options
                      const pmConfig *config // The configuration
    )
{
    assert(data);
    assert(options);
    assert(config);

    if (!options->scale && !options->zero && !options->darktime && !options->shutter) {
        return true;                    // We did everything we were asked for
    }

    assert(config->camera);             // Need the camera configuration
    assert(config->arguments);          // Need the list of files

    psArray *filenames = psMetadataLookupPtr(NULL, config->arguments, "INPUT"); // The input file names
    assert(filenames);                  // It should be here --- it's put here in ppMergeConfig

    // Sanity checks
    assert(!options->scale || scales);
    assert(!scales || !*scales || ((*scales)->type.type == PS_TYPE_F32 &&
                                   (*scales)->numCols == data->numCells &&
                                   (*scales)->numRows == filenames->n));
    assert(!options->zero || zeros);
    assert(!zeros || !*zeros || ((*zeros)->type.type == PS_TYPE_F32 &&
                                 (*zeros)->numCols == data->numCells &&
                                 (*zeros)->numRows == filenames->n));
    assert(!options->shutter || shutters);
    assert(!shutters || !*shutters || (*shutters)->n == data->numCells);

    // Allocate the outputs
    if (options->scale) {
        if (*scales) {
            psFree(*scales);
        }
        *scales = psImageAlloc(data->numCells, filenames->n, PS_TYPE_F32);
    }
    if (options->zero) {
        if (*zeros) {
            psFree(*zeros);
        }
        *zeros = psImageAlloc(data->numCells, filenames->n, PS_TYPE_F32);
    }
    psRandom *rng = NULL;               // Random number generator
    if (options->shutter) {
        if (*shutters) {
            psFree(*shutters);
        }
        *shutters = psArrayAlloc(data->numCells);
        rng = psRandomAlloc(PS_RANDOM_TAUS, 0);
    }

    bool fromConcepts = false;          // Do we get the scale and zero points from the concepts?
    bool first = true;                  // Are we on the first cell (that sets the standard for the rest)?
    bool done = false;                  // Are we done going through the list?
    bool mdok = true;                   // Status of MD lookup
    for (long i = 0; i < data->in->n && !done; i++) {
        pmFPA *fpa = data->in->data[i]; // The FPA
        if (!fpa) {
            continue;
        }
        long cellNum = -1;              // Number of the cell
        psArray *chips = fpa->chips;    // The array of chips
        for (long j = 0; j < chips->n && !done; j++) {
            pmChip *chip = chips->data[j]; // The chip
            if (!chip) {
                continue;
            }
            psArray *cells = chip->cells; // The array of cells
            for (long k = 0; k < cells->n && !done; k++) {
                pmCell *cell = cells->data[k]; // The cell
                if (!cell) {
                    continue;
                }
                cellNum++;

                if (options->scale) {
                    float scale = psMetadataLookupF32(&mdok, cell->concepts, "PPMERGE.SCALE"); // The scale
                    if (mdok && !isnan(scale)) {
                        if (!first && !fromConcepts) {
                            psWarning("PPMERGE.SCALE and PPMERGE.ZERO have been set "
                                     "for some, but not all cells --- we will re-measure it for all cells.");
                            done = true;
                            continue;
                        }
                        fromConcepts = true;
                        (*scales)->data.F32[i][cellNum] = scale;
                        psTrace("ppMerge", 9, "Scale for input %ld, chip %ld, cell %ld: %f\n",
                                i, j, k, scale);
                    } else if (!first && fromConcepts) {
                        psWarning("PPMERGE.SCALE and PPMERGE.ZERO have been set "
                                 "for some, but not all cells --- we will re-measure it for all cells.");
                        fromConcepts = false;
                        done = true;
                        continue;
                    }
                }

                if (options->zero) {
                    float zero = psMetadataLookupF32(&mdok, cell->concepts, "PPMERGE.ZERO"); // The zero
                    if (mdok && !isnan(zero)) {
                        if (!first && !fromConcepts) {
                            psWarning("PPMERGE.SCALE and PPMERGE.ZERO have been set "
                                     "for some, but not all cells --- we will re-measure it for all cells.");
                            done = true;
                            continue;
                        }
                        fromConcepts = true;
                        (*zeros)->data.F32[i][cellNum] = zero;
                        psTrace("ppMerge", 9, "Zero for input %ld, chip %ld, cell %ld: %f\n", i, j, k, zero);
                    } else if (!first && fromConcepts) {
                        psWarning("PPMERGE.SCALE and PPMERGE.ZERO have been set "
                                 "for some, but not all cells --- we will re-measure it for all cells.");
                        fromConcepts = false;
                        done = true;
                        continue;
                    }
                }

                first = false;
            }
        }
    }

    if (fromConcepts) {
        // We've already done everything we need to
        psFree(rng);
        return true;
    }

    psImage *background = psImageAlloc(data->numCells, filenames->n, PS_TYPE_F32); // Background measurements
    psImageInit(background, NAN);
    psStats *bgStats = psStatsAlloc(options->mean); // Statistic to measure the background
    psVector *gains = NULL;             // The gains for each cell
    psImage *darktime = NULL;           // The dark time for each integration of each cell
    if (options->scale) {
        gains = psVectorAlloc(data->numCells, PS_TYPE_F32);
    }
    if (options->darktime) {
        darktime = psImageAlloc(data->numCells, filenames->n, PS_TYPE_F32);
    }

    pmFPAview *view = pmFPAviewAlloc(0);

    bool status = true;                 // Status of getting the scale and zero --- did everything go right?
    for (int i = 0; i < filenames->n; i++) {
        psString name = filenames->data[i]; // The name of the file
        if (!name || strlen(name) == 0) {
            continue;
        }
        psTrace("ppMerge", 9, "Opening %s to get background...\n", name);
        psFits *inFile = data->files->data[i]; // The FITS file to read
        pmFPA *fpa = data->in->data[i]; // The FPA for this input
        int cellNum = -1;               // Number of the cell
        psArray *chips = fpa->chips;    // Array of chips
        for (int j = 0; j < chips->n; j++) {
            pmChip *chip = chips->data[j]; // The chip of interest
            if (!chip) {
                continue;
            }
            psArray *cells = chip->cells; // Array of cells
            for (int k = 0; k < cells->n; k++) {
                pmCell *cell = cells->data[k]; // The cell of interest
                if (!cell) {
                    continue;
                }
                cellNum++;

                if (!pmCellReadHeader(cell, inFile)) {
                    continue;
                }

                // Normalising by the exposure time
                if (options->darktime) {
                    darktime->data.F32[i][cellNum] = psMetadataLookupF32(NULL, cell->concepts,
                                                                         "CELL.DARKTIME");
                    if (isnan(darktime->data.F32[i][cellNum])) {
                        psWarning("CELL.DARKTIME for file %s chip %d cell %d is not set.\n", name, j, k);
                        status = false;
                    } else {
                        view->chip = j;
                        view->cell = k;
                        darktime->data.F32[i][cellNum] = pmFPADarkNorm(fpa, view,
                                                                       darktime->data.F32[i][cellNum]);
                        if (isnan(darktime->data.F32[i][cellNum])) {
                            psWarning("Cannot get dark normalisation for file %s chip %d cell %d.\n",
                                      name, j, k);
                            status = false;
                        }
                    }
                }

                // Scaling by the background
                if (options->scale || options->zero) {
                    if (!pmCellRead(cell, inFile, config->database)) {
                        // Nothing here
                        pmCellFreeData(cell);
                        continue;
                    }

                    if (cell->readouts->n > 1) {
                        psWarning("File %s chip %d cell %d contains more than one "
                                 "readout --- ignoring all but the first.\n", name, j, k);
                        status = false;
                    }

                    pmReadout *readout = cell->readouts->data[0]; // The readout of interest
                    psImage *image = readout->image; // The pixels of interest
                    if (!image) {
                        pmCellFreeData(cell);
                        continue;
                    }

                    // Get the gain
                    if (options->scale) {
                        bool mdok = true;   // Status of MD lookup
                        gains->data.F32[cellNum] = psMetadataLookupF32(&mdok, cell->concepts, "CELL.GAIN");
                        if (!mdok || isnan(gains->data.F32[cellNum])) {
                            psWarning("CELL.GAIN for file %s chip %d cell %d is not "
                                     "set.\n", name, j, k);
                            gains->data.F32[cellNum] = NAN;
                            status = false;
                        }
                    }

                    // Get the background
                    int sampleSize = (image->numCols * image->numRows) / options->sample; // Size of sample
                    psVector *sample = psVectorAlloc(sampleSize, PS_TYPE_F32); // Sample of the image
                    psVector *sampleMask = NULL; // Mask for sample
                    if (readout->mask) {
                        sampleMask = psVectorAlloc(sampleSize, PS_TYPE_U8);
                    }
                    psImage *mask = readout->mask; // The mask image
                    for (long i = 0; i < sampleSize; i++) {
                        int j = i * options->sample; // Index into image
                        int x = j % image->numCols; // x index
                        int y = j / image->numCols; // y index
                        sample->data.F32[i] = image->data.F32[y][x];
                        if (readout->mask) {
                            sampleMask->data.U8[i] = mask->data.U8[y][x];
                        }
                    }
                    status = psVectorStats(bgStats, sample, sampleMask, NULL, options->combine->maskVal);
                    if (!status) {
                      psTrace("ppMerge", 3, "failed to get stats for for %s, cell %d is %f\n", name, cellNum,
                              background->data.F32[i][cellNum]);
                      psErrorClear();
                    }
                    psFree(sample);
                    psFree(sampleMask);
                    background->data.F32[i][cellNum] = psStatsGetValue(bgStats, options->mean);
                    psTrace("ppMerge", 3, "Background for %s, cell %d is %f\n", name, cellNum,
                            background->data.F32[i][cellNum]);
                }

                // Shutter correction
                if (options->shutter) {
                    if (!pmCellRead(cell, inFile, config->database)) {
                        // Nothing here
                        pmCellFreeData(cell);
                        continue;
                    }

                    if (cell->readouts->n > 1) {
                        psWarning("File %s chip %d cell %d contains more than one "
                                  "readout --- ignoring all but the first.\n", name, j, k);
                        status = false;
                    }
                    pmReadout *readout = cell->readouts->data[0]; // The readout of interest

                    pmShutterCorrectionData *shutter = (*shutters)->data[cellNum]; // Shutter correction data
                    if (!shutter) {
                        shutter = pmShutterCorrectionDataAlloc(readout->image->numCols,
                                                               readout->image->numRows,
                                                               options->shutterSize);
                        (*shutters)->data[cellNum] = shutter;
                    }
                    if (!pmShutterCorrectionAddReadout(shutter, readout, options->mean, options->stdev,
                                                       options->combine->maskVal, rng)) {
                        psWarning("Can't add file %s chip %d cell %d to shutter correction --- ignored.",
                                  name, j, k);
                        status = false;
                    }
                }


                pmCellFreeData(cell);
            }
            pmChipFreeData(chip);
        }
        pmFPAFreeData(fpa);
    }
    psFree(rng);
    psFree(bgStats);
    psFree(view);

    if (options->scale) {
        // Need to normalize over the focal plane
        if (psTraceGetLevel("ppMerge") > 9) {
            for (int i = 0; i < gains->n; i++) {
                psTrace("ppMerge", 10, "Gain for cell %d is %f\n", i, gains->data.F32[i]);
            }
        }
        psVector *fluxes = NULL;        // Solution to fluxes
        if (!pmFlatNormalize(&fluxes, &gains, background)) {
            psWarning("Normalisation failed to converge --- continuing anyway.\n");
            status = false;
        }
        psFree(gains);

        psImage *scalesDeref = *scales; // Dereference the pointer

        for (int i = 0; i < scalesDeref->numRows; i++) {
            psF32 bg = fluxes->data.F32[i];
            for (int j = 0; j < scalesDeref->numCols; j++) {
                scalesDeref->data.F32[i][j] = bg;
            }
        }
    }

    if (options->darktime) {
        if (!options->scale) {
            // Copy over the exposure times
            *scales = psImageCopy(*scales, darktime, PS_TYPE_F32);
        } else {
            *scales = (psImage*)psBinaryOp(*scales, *scales, "*", darktime);
        }
    }

    if (options->zero) {
        if (!*zeros) {
            // This is much faster than copying!
            *zeros = psMemIncrRefCounter(background);
        } else {
            *zeros = psImageCopy(*zeros, background, PS_TYPE_F32);
        }
    }

    // Diagnostic stuff
    if (scales && *scales && psTraceGetLevel("ppMerge") > 9) {
        psImage *scalesDeref = *scales; // Dereference the pointer
        for (int i = 0; i < scalesDeref->numRows; i++) {
            for (int j = 0; j < scalesDeref->numCols; j++) {
                psTrace("ppMerge", 9, "Scale for exposure %d, cell %d is: %f\n", i, j,
                        scalesDeref->data.F32[i][j]);
            }
        }
    }
    if (zeros && *zeros && psTraceGetLevel("ppMerge") > 9) {
        psImage *zerosDeref = *zeros; // Dereference the pointer
        for (int i = 0; i < zerosDeref->numRows; i++) {
            for (int j = 0; j < zerosDeref->numCols; j++) {
                psTrace("ppMerge", 9, "Zero for exposure %d, cell %d is: %f\n", i, j,
                        zerosDeref->data.F32[i][j]);
            }
        }
    }

    psFree(background);

    return status;
}

