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

#include <stdio.h>
#include <string.h>

#include <pslib.h>
#include <psmodules.h>
#include <ppStats.h>

#include "ppMerge.h"

bool ppMergeLoop(pmConfig *config)
{
    assert(config);

    psMetadata *arguments = config->arguments; // Arguments
    ppMergeType type = psMetadataLookupS32(NULL, config->arguments, "TYPE"); // Type of frame
    int numFiles = psMetadataLookupS32(NULL, arguments, "INPUTS.NUM"); // Number of input files
    bool mdok;                          // Status of MD lookup
    bool haveMasks = psMetadataLookupBool(&mdok, arguments, "INPUTS.MASKS"); // Do we have masks?
    bool haveWeights = psMetadataLookupBool(&mdok, arguments, "INPUTS.WEIGHTS"); // Do we have weights?

    psArray *inputs = ppMergeFileDataLevel(config, "PPMERGE.INPUT", PM_FPA_LEVEL_READOUT); // Input images
    psArray *masks = NULL, *weights = NULL; // Input masks and weights
    if (haveMasks) {
        masks = ppMergeFileDataLevel(config, "PPMERGE.INPUT.MASK", PM_FPA_LEVEL_READOUT);
    }
    if (haveWeights) {
        weights = ppMergeFileDataLevel(config, "PPMERGE.INPUT.WEIGHT", PM_FPA_LEVEL_READOUT);
    }

    // General combination parameters
    int rows = psMetadataLookupS32(NULL, arguments, "ROWS"); // Number of rows to read per chunk
    int iter = psMetadataLookupS32(NULL, arguments, "ITER"); // Number of rejection iterations
    float rej = psMetadataLookupF32(NULL, arguments, "REJ"); // Rejection level
    float fraclow = psMetadataLookupF32(NULL, arguments, "FRACLOW"); // Reject fraction of low pixels
    float frachigh = psMetadataLookupF32(NULL, arguments, "FRACHIGH"); // Reject fraction of hi pixels
    int nKeep = psMetadataLookupS32(NULL, arguments, "NKEEP"); // Minimum number of values to keep
    psMaskType maskVal = psMetadataLookupU8(NULL, arguments, "MASKVAL"); // Value to mask
    psStatsOptions combineStat = psMetadataLookupS32(NULL, arguments, "COMBINE"); // Combination statistic
    bool useWeights = psMetadataLookupBool(NULL, arguments, "WEIGHTS"); // Use weights?

    // Fringe parameters
    int fringeNum = psMetadataLookupS32(NULL, arguments, "FRINGE.NUM"); // Number of fringe points
    int fringeSize = psMetadataLookupS32(NULL, arguments, "FRINGE.SIZE"); // Size of fringe regions
    int fringeSmoothX = psMetadataLookupS32(NULL, arguments, "FRINGE.XSMOOTH"); // Smoothing regions in x
    int fringeSmoothY = psMetadataLookupS32(NULL, arguments, "FRINGE.YSMOOTH"); // Smoothing regions in y

    pmCombineParams *combination = pmCombineParamsAlloc(combineStat); // Combination parameters
    combination->maskVal = maskVal;
    combination->blank = pmConfigMask("BLANK", config);
    combination->nKeep = nKeep;
    combination->fracHigh = frachigh;
    combination->fracLow = fraclow;
    combination->iter = iter;
    combination->rej = rej;
    combination->weights = useWeights;

    psMetadata *stats = NULL;           // Statistics for output
    if (psMetadataLookup(config->arguments, "STATS.NAME")) {
        stats = psMetadataAlloc();
        psMetadataAddMetadata(config->arguments, PS_LIST_TAIL, "STATS.DATA", 0, "Statistics output", stats);
    }

    pmFPAview *view = pmFPAviewAlloc(0); // View to component of interest

    // Retrieve data placed on analysis
    psVector *scales = NULL, *zeros = NULL; // Scale and zeroes for combination
    psArray *shutters = NULL;           // Shutter correction data
    switch (type) {
      case PPMERGE_TYPE_FRINGE:
        zeros = psMetadataLookupPtr(NULL, arguments, "ZEROS");
        if (!zeros) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find ZEROS");
            goto ERROR;
        }
        // Flow through
      case PPMERGE_TYPE_FLAT:
        scales = psMetadataLookupPtr(NULL, arguments, "SCALES");
        if (!scales) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find SCALES");
            goto ERROR;
        }
        break;
      case PPMERGE_TYPE_SHUTTER:
        shutters = psMetadataLookupPtr(NULL, arguments, "SHUTTER");
        if (!shutters) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find SHUTTER");
            goto ERROR;
        }
        break;
      case PPMERGE_TYPE_MASK:
      case PPMERGE_TYPE_BIAS:
      case PPMERGE_TYPE_DARK:
        break;
      default:
        psAbort("Should never get here.");
    }

    // Dark parameters
    psArray *darkOrdinates = psMetadataLookupPtr(NULL, arguments, "DARK.ORDINATES"); // Dark info
    psString darkNorm = psMetadataLookupStr(&mdok, arguments, "DARK.NORM"); // Dark normalisation


    if (!ppMergeFileActivate(config, PPMERGE_FILES_ALL, true)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to activate files.");
        goto ERROR;
    }
    psString outName = ppMergeOutputFile(config); // Name of output file
    pmFPAfile *output = psMetadataLookupPtr(NULL, config->files, outName); // Output file
    psFree(outName);
    assert(output && output->fpa);
    pmFPA *outFPA = output->fpa;        // Output FPA
    int cellNum = 0;                    // Index of cell
    if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
        goto ERROR;
    }
    pmChip *outChip;                    // Chip of interest
    while ((outChip = pmFPAviewNextChip(view, outFPA, 1))) {
        if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
            goto ERROR;
        }
        pmCell *outCell;                // Cell of interest
        while ((outCell = pmFPAviewNextCell(view, outFPA, 1))) {
            if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
                goto ERROR;
            }

            pmHDU *hdu = pmHDUGetLowest(outFPA, outChip, outCell); // HDU for cell
            if (!hdu || hdu->blankPHU) {
                // No data here
                continue;
            }

            pmReadout *outRO = pmReadoutAlloc(outCell);

            psArray *readouts = psArrayAlloc(numFiles); // Input readouts
            for (int i = 0; i < numFiles; i++) {
                // We need to do some of the opening ourselves
                if (!ppMergeFileOpenInput(config, view, i)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to open file %d", i);
                    goto ERROR;
                }

                pmFPAfile *input = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", i);
                pmCell *inCell = pmFPAviewThisCell(view, input->fpa); // Input cell
                readouts->data[i] = pmReadoutAlloc(inCell);
            }

            float shutterRef = NAN;     // Reference shutter correction
            if (type == PPMERGE_TYPE_SHUTTER) {
                shutterRef = pmShutterCorrectionReference(shutters->data[cellNum]);
            }

            // Read convolutions by chunks
            bool more = true;               // More to read?
            for (int numChunk = 0; more; numChunk++) {
                psTrace("ppStack", 2, "Initial stack of chunk %d....\n", numChunk);
                for (int i = 0; i < numFiles; i++) {
                    pmReadout *inRO = readouts->data[i]; // Input readout

                    // Read a chunk from a file
                    #define READ_CHUNK(NAME,TYPE) { \
                        pmFPAfile *file = pmFPAfileSelectSingle(config->files, NAME, i); \
                        if (!pmReadoutReadChunk##TYPE(inRO, file->fits, 0, rows, 0, config)) { \
                            psError(PS_ERR_IO, false, "Unable to read chunk %d for file %s %d", \
                                    numChunk, NAME, i); \
                            psFree(readouts); \
                            psFree(outRO); \
                            goto ERROR; \
                        } \
                    }

                    READ_CHUNK("PPMERGE.INPUT", /* Blank */);
                    if (haveMasks) {
                        READ_CHUNK("PPMERGE.INPUT.MASK", Mask);
                    }
                    if (haveWeights) {
                        READ_CHUNK("PPMERGE.INPUT.WEIGHT", Weight);
                    }
                }

                switch (type) {
                  case PPMERGE_TYPE_SHUTTER:
                    if (!pmShutterCorrectionGenerate(outRO, NULL, readouts, shutterRef,
                                                     shutters->data[cellNum], iter, rej, maskVal)) {
                        psFree(readouts);
                        psFree(outRO);
                        goto ERROR;
                    }
                    break;
                  case PPMERGE_TYPE_DARK:
                    if (!pmDarkCombine(outCell, readouts, darkOrdinates, darkNorm, iter, rej, maskVal)) {
                        psFree(readouts);
                        psFree(outRO);
                        goto ERROR;
                    }
                    break;
                  case PPMERGE_TYPE_BIAS:
                  case PPMERGE_TYPE_FLAT:
                  case PPMERGE_TYPE_FRINGE:
                    if (!pmReadoutCombine(outRO, readouts, zeros, scales, combination)) {
                        psFree(readouts);
                        psFree(outRO);
                        goto ERROR;
                    }
                    break;
                  default:
                    psAbort("Should never get here.");
                }


                for (int i = 0; i < numFiles && more; i++) {
                    pmReadout *inRO = readouts->data[i];

                    // Check to see if there's more chunks to read
                    #define MORE_CHUNK(NAME,TYPE) { \
                        pmFPAfile *file = pmFPAfileSelectSingle(config->files, NAME, i); \
                        more &= pmReadoutMore##TYPE(inRO, file->fits, 0, rows, config); \
                    }

                    MORE_CHUNK("PPMERGE.INPUT", /* Blank */);
                    if (haveMasks) {
                        MORE_CHUNK("PPMERGE.INPUT.MASK", Mask);
                    }
                    if (haveWeights) {
                        MORE_CHUNK("PPMERGE.INPUT.WEIGHT", Weight);
                    }
                }
            }

            // Get list of cells for concepts averaging
            psList *inCells = psListAlloc(NULL); // List of cells
            for (int i = 0; i < numFiles; i++) {
                pmReadout *readout = readouts->data[i]; // Readout of interest
                psListAdd(inCells, PS_LIST_TAIL, readout->parent);
            }
            if (!pmConceptsAverageCells(outCell, inCells, NULL, NULL, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to average cell concepts.");
                psFree(inCells);
                psFree(outRO);
                goto ERROR;
            }
            psFree(inCells);

            psFree(readouts);

            // Plug supplementary images into their own FPAs
            {
                pmCell *countsCell = pmFPAfileThisCell(config->files, view, "PPMERGE.OUTPUT.COUNT");
                pmReadout *countsRO = pmReadoutAlloc(countsCell); // Readout with count of inputs per pixel
                psImage *counts = psMetadataLookupPtr(NULL, outRO->analysis, PM_READOUT_STACK_ANALYSIS_COUNT);
                countsRO->image = psImageCopy(countsRO->image, counts, PS_TYPE_F32);
                psMetadataRemoveKey(outRO->analysis, PM_READOUT_STACK_ANALYSIS_COUNT);
                countsRO->data_exists = countsCell->data_exists = countsCell->parent->data_exists = true;
                psFree(countsRO);

                pmCell *sigmaCell = pmFPAfileThisCell(config->files, view, "PPMERGE.OUTPUT.SIGMA");
                pmReadout *sigmaRO = pmReadoutAlloc(sigmaCell); // Readout with stdev per pixel
                psImage *sigma = psMetadataLookupPtr(NULL, outRO->analysis, PM_READOUT_STACK_ANALYSIS_SIGMA);
                sigmaRO->image = psImageCopy(sigmaRO->image, sigma, PS_TYPE_F32);
                psMetadataRemoveKey(outRO->analysis, PM_READOUT_STACK_ANALYSIS_SIGMA);
                sigmaRO->data_exists = sigmaCell->data_exists = sigmaCell->parent->data_exists = true;
                psFree(sigmaRO);
            }

            // Measure the fringes for this cell
            //
            // XXX Need to deal with multiple components: we will do this by building up the components
            // Read the existing fringe measurements
            // Use existing regions to measure fringe statistics
            // Add the new fringe measurements to the existing fringe measurements.
            // Write the appended fringe measurements.
            // Read in the "output" file to get the existing components.
            // Put the new readout into the cell after the existing readouts.
            if (type == PPMERGE_TYPE_FRINGE && outRO) {
                pmFringeRegions *regions = pmFringeRegionsAlloc(fringeNum, fringeSize, fringeSize,
                                                                fringeSmoothX, fringeSmoothY);
                pmFringeStats *fringe = pmFringeStatsMeasure(regions, outRO, maskVal);
                psFree(regions);
                if (!fringe) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to measure fringe statistics.\n");
                    psFree(outRO);
                    goto ERROR;
                }

                psArray *fringes = psArrayAlloc(1); // Array of fringes
                fringes->data[0] = fringe;

                pmFringesFormat(outCell, NULL, fringes);
                psFree(fringes);        // Drop reference
            }

            if (!ppStatsFPA(stats, outFPA, view, maskVal | pmConfigMask("BLANK", config), config)) {
                psError(PS_ERR_UNKNOWN, true, "Unable to generate stats for image.");
                goto ERROR;
            }

            psFree(outRO);
            cellNum++;

            if (!pmFPAfileIOChecks(config, view, PM_FPA_AFTER)) {
                goto ERROR;
            }
        }

        if (outChip->data_exists) {
            psList *inChips = psListAlloc(NULL);
            for (int i=0; i < numFiles; i++) {
                pmChip *chip = pmFPAviewThisChip(view, ((pmFPAfile *)inputs->data[i])->fpa);
                psListAdd(inChips, PS_LIST_TAIL, chip);
            }
            if (!pmConceptsAverageChips(outChip, inChips, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to average Chip concepts.");
                psFree(inChips);
                goto ERROR;
            }
            psFree(inChips);
        }
        if (!pmFPAfileIOChecks(config, view, PM_FPA_AFTER)) {
            goto ERROR;
        }
    }

    // Get list of FPAs for concepts averaging
    psList *inFPAs = psListAlloc(NULL); // List of FPAs
    for (int i = 0; i < numFiles; i++) {
        pmFPAfile *input = inputs->data[i]; // Input file
        psListAdd(inFPAs, PS_LIST_TAIL, input->fpa);
    }
    if (!pmConceptsAverageFPAs(output->fpa, inFPAs)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to average FPA concepts.");
        psFree(inFPAs);
        goto ERROR;
    }
    psFree(inFPAs);


    if (!pmFPAfileIOChecks(config, view, PM_FPA_AFTER)) {
        goto ERROR;
    }

    psFree(view);
    psFree(combination);
    psFree(inputs);
    psFree(masks);
    psFree(weights);
    psFree(stats);
    return true;

ERROR:
    psFree(view);
    psFree(combination);
    psFree(inputs);
    psFree(masks);
    psFree(weights);
    psFree(stats);
    return false;
}

