#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"

// XXX this function is now sufficiently different for the major types, it would make sense to just
// split it into three: BASIC, SHUTTER, DARK

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);
    }

    int nThreads = psMetadataLookupS32 (&mdok, arguments, "NTHREADS");
    if (!mdok) nThreads = 0;

    // General combination parameters
    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
    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

    // set the mask and mark bit values based on the named masks
    psMaskType maskVal;
    psMaskType markVal;
    if (!pmConfigMaskSetBits (&maskVal, &markVal, config)) {
        psError (PS_ERR_UNKNOWN, true, "Unable to define the mask bit values");
        return false;
    }

    pmCombineParams *combination = pmCombineParamsAlloc(combineStat); // Combination parameters
    combination->maskVal = maskVal;
    combination->blank = pmConfigMaskGet("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;
            }

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

            pmReadout *outRO = pmReadoutAlloc(outCell);

            // open the input files (we need to do the work ourselves)
            for (int i = 0; i < numFiles; i++) {
                if (!ppMergeFileOpenInput(config, view, i)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to open file %d", i);
                    goto ERROR;
                }
            }

            ppMergeFileGroup *fileGroup = NULL;
            psArray *fileGroups = psArrayAlloc(nThreads + 1);

            // Generate readouts for each input file in each file group
            for (int i = 0; i < fileGroups->n; i++) {
                psArray *readouts = psArrayAlloc(numFiles); // Input readouts
                for (int j = 0; j < numFiles; j++) {
                    pmFPAfile *input = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", j);
                    pmCell *inCell = pmFPAviewThisCell(view, input->fpa); // Input cell
                    readouts->data[j] = pmReadoutAlloc(inCell);
                }

                fileGroup = ppMergeFileGroupAlloc();
                fileGroup->readouts = readouts;
                fileGroup->read = false;
                fileGroup->busy = false;
                fileGroup->lastScan = 0;
                fileGroup->firstScan = 0;
                fileGroups->data[i] = fileGroup;
            }

            // call the init functions
            switch (type) {
              case PPMERGE_TYPE_BIAS:
              case PPMERGE_TYPE_FLAT:
              case PPMERGE_TYPE_FRINGE:
                psAssert (fileGroups->n > 0, "no valid file groups defined");
                fileGroup = fileGroups->data[0];
                if (!pmReadoutCombinePrepare(outRO, fileGroup->readouts, combination)) {
                    goto ERROR;
                }
                break;
              case PPMERGE_TYPE_DARK:
                psAssert (fileGroups->n > 0, "no valid file groups defined");
                fileGroup = fileGroups->data[0];
                if (!pmDarkCombinePrepare(outCell, fileGroup->readouts, darkOrdinates, darkNorm)) {
                    goto ERROR;
                }
                break;
              case PPMERGE_TYPE_SHUTTER:
                psAssert (fileGroups->n > 0, "no valid file groups defined");
                fileGroup = fileGroups->data[0];
                if (!pmShutterCorrectionGeneratePrepare(outRO, pattern, fileGroup->readouts, maskVal)) {
                    goto ERROR;
                }
                break;
              default:
                psAbort("Should never get here.");
            }

            // Read input data by chunks
            psTimerStart("ppMergeLoop");
            for (int numChunk = 0; true; numChunk++) {

                bool status = false;
                fileGroup = ppMergeReadChunk(&status, fileGroups, config, numChunk);
                if (!status) {
                    // Something went wrong
                    goto ERROR;
                }
                if (!fileGroup) {
                    // Nothing more to read
                    break;
                }

                // Start a job
                switch (type) {
                  case PPMERGE_TYPE_BIAS:
                  case PPMERGE_TYPE_FLAT:
                  case PPMERGE_TYPE_FRINGE: {
                      psThreadJob *job = psThreadJobAlloc("PPMERGE_READOUT_COMBINE"); // Job to start

                      // Construct the arguments for this job
                      psArrayAdd(job->args, 1, outRO);
                      psArrayAdd(job->args, 1, fileGroup);
                      psArrayAdd(job->args, 1, zeros);
                      psArrayAdd(job->args, 1, scales);
                      psArrayAdd(job->args, 1, combination);

                      // call: pmReadoutCombine(outRO, fileGroup->readouts, zeros, scales, combination);
                      if (!psThreadJobAddPending(job)) {
                          goto ERROR;
                      }
                      break;
                  }
                  case PPMERGE_TYPE_DARK: {
                      psThreadJob *job = psThreadJobAlloc ("PPMERGE_DARK_COMBINE"); // Job to start

                      // construct the arguments for this job
                      psArrayAdd(job->args, 1, outCell);
                      psArrayAdd(job->args, 1, fileGroup);
                      psArrayAdd(job->args, 1, psScalarAlloc(iter, PS_TYPE_S32));
                      psArrayAdd(job->args, 1, psScalarAlloc(rej, PS_TYPE_F32));
                      psArrayAdd(job->args, 1, psScalarAlloc(maskVal, PS_TYPE_U8));

                      // call: pmDarkCombine(outCell, fileGroup->readouts, iter, rej, maskVal);
                      if (!psThreadJobAddPending(job)) {
                          goto ERROR;
                      }
                      break;
                  }
                  case PPMERGE_TYPE_SHUTTER: {
                      psThreadJob *job = psThreadJobAlloc ("PPMERGE_SHUTTER_CORRECTION");

                      // construct the arguments for this job
                      psArrayAdd(job->args, 1, outRO);
                      psArrayAdd(job->args, 1, pattern);
                      psArrayAdd(job->args, 1, fileGroup);
                      psArrayAdd(job->args, 1, psScalarAlloc(shutterRef, PS_TYPE_F32));
                      psArrayAdd(job->args, 1, shutters->data[cellNum]);
                      psArrayAdd(job->args, 1, psScalarAlloc(iter, PS_TYPE_S32));
                      psArrayAdd(job->args, 1, psScalarAlloc(rej, PS_TYPE_F32));
                      psArrayAdd(job->args, 1, psScalarAlloc(maskVal, PS_TYPE_U8));

                      // call: pmShutterCorrectionGenerate(outRO, pattern, fileGroup->readouts, shutterRef,
                      //                                   shutters->data[cellNum], iter, rej, maskVal);
                      if (!psThreadJobAddPending (job)) {
                          goto ERROR;
                      }
                      break;
                  }
                  default:
                    psAbort("Should never get here.");
                }
            }

            // Wait for the threads to finish and manage results
            if (!psThreadPoolWait(false)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to combine images.");
                return false;
            }

            // we don't care about the results, just dump the done queue jobs
            psThreadJob *job = NULL;    // Job to dump
            while ((job = psThreadJobGetDone())) {
                psFree (job);
            }

            psFree(fileGroups);

            // XXX eventually need to keep both the shutter and the pattern, as we do with dark
            psFree(pattern);

            // Get list of cells for concepts averaging
            psList *inCells = psListAlloc(NULL); // List of cells
            for (int i = 0; i < numFiles; i++) {
                pmFPAfile *input = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", i);
                pmCell *inCell = pmFPAviewThisCell(view, input->fpa); // Input cell
                psListAdd(inCells, PS_LIST_TAIL, inCell);
            }
            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);

            // 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 (stats && !ppStatsFPA(stats, outFPA, view, maskVal, 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;
}

