/** @file ppMergeMask.c
 *
 *  @brief
 *
 *  @ingroup ppMerge
 *
 *  @author IfA
 *  @version $Revision: 1.18 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-02-06 02:44:31 $
 *  Copyright 2009 Institute for Astronomy, University of Hawaii
 */

#include "ppMerge.h"

static bool mergeMask(pmConfig *config, ///< Configuration
                      const pmFPAview *view, ///< View to chip
                      bool writeOut,     ///< Write output?
                      psRandom *rng,    ///< Random number generator
                      psMetadata *stats ///< Statistics output
    )
{
    assert(config);
    assert(view);
    assert(view->chip != -1 && view->cell == -1 && view->readout == -1);

    bool mdok;                          ///< Status of MD lookup
    int numFiles = psMetadataLookupS32(NULL, config->arguments, "INPUTS.NUM"); ///< Number of input files
    psStatsOptions meanStat = psMetadataLookupS32(NULL, config->arguments, "MEAN"); ///< Statistic for mean
    psStatsOptions stdevStat = psMetadataLookupS32(NULL, config->arguments, "STDEV"); ///< Statistic for stdev
    int sample = psMetadataLookupS32(NULL, config->arguments, "SAMPLE"); ///< Size of sample for statistics
    bool chipStats = psMetadataLookupBool(&mdok, config->arguments, "MASK.CHIPSTATS"); ///< Statistics on chip?
    float maskSuspect = psMetadataLookupF32(NULL, config->arguments, "MASK.SUSPECT"); ///< Threshold for suspect pixels
    float maskBad = psMetadataLookupF32(NULL, config->arguments, "MASK.BAD"); ///< Threshold for bad pixels
    pmMaskIdentifyMode maskMode = psMetadataLookupS32(NULL, config->arguments, "MASK.MODE"); ///< Mode for identifying bad pixels
    int maskGrow = psMetadataLookupS32(NULL, config->arguments, "MASK.GROW"); ///< Radius to grow mask

    bool smoothSuspect = psMetadataLookupBool(&mdok, config->arguments, "MASK.SMOOTH.SUSPECT"); ///< Radius to grow mask
    float smoothScale = psMetadataLookupF32(&mdok, config->arguments, "MASK.SMOOTH.SCALE"); ///< Radius to grow mask

    psImageMaskType markVal, maskValRaw;
    if (!pmConfigMaskSetBits(&maskValRaw, &markVal, config)) {
        psError(PS_ERR_UNKNOWN, true, "Unable to define the mask bit values");
        return false;
    }

    char *maskOutName = psMetadataLookupStr (&mdok, config->arguments, "MASK.SET.VALUE");
    psImageMaskType maskValOut = pmConfigMaskGet (maskOutName, config);
    if (!maskValOut) {
        psError (PS_ERR_UNKNOWN, true, "Undefined output mask bit value");
        return false;
    }

    psStats *statistics = psStatsAlloc(meanStat | stdevStat); // Statistics for background

    psString outName = ppMergeOutputFile(config); ///< Name of output file
    pmChip *outChip = pmFPAfileThisChip(config->files, view, outName); ///< Output chip
    psFree(outName);

    int numCells = 1;
    if (chipStats) {
        // count the number of active cells for this chip:
        numCells = 0;
        for (int i = 0; i < outChip->cells->n; i++) {
            pmCell *cell = outChip->cells->data[i];
            if (!cell->process) continue;
            numCells ++;
        }
    }

    // For each input file, get the statistics, which can be calculated at the chip or cell levels
    psVector *values = psVectorAlloc(sample, PS_TYPE_F32); ///< Pixel values for statistics
    pmFPAview *inView = pmFPAviewAlloc(0); ///< View for input
    for (int i = 0; i < numFiles; i++) {
        pmFPAfileActivate(config->files, false, NULL);
        psArray *files = ppMergeFileActivateSingle(config, PPMERGE_FILES_INPUT, true, i); ///< Input files
        pmFPAfile *input = files->data[0]; ///< Input file
        psFree(files);
        pmFPA *inFPA = input->fpa;  ///< Input FPA
        *inView = *view;

        int valueIndex = 0;             ///< Index for vector of pixel values

        pmCell *inCell;                 ///< Input cell
        while ((inCell = pmFPAviewNextCell(inView, inFPA, 1))) {

            // the output FPA structure carries the information about which cells to process
            pmCell *outCell = pmFPAfileThisCell(config->files, inView, "PPMERGE.OUTPUT.MASK"); // Output cell
            if (!outCell->process) continue;

            pmHDU *hdu = pmHDUFromCell(inCell); ///< HDU for cell
            if (!hdu || hdu->blankPHU) {
                // No data here
                continue;
            }
            psTrace("ppMerge", 1, "Getting suspect pixels for file %d chip %d cell %d",
                    i, inView->chip, inView->cell);

            if (!pmFPAfileIOChecks(config, inView, PM_FPA_BEFORE)) {
                psFree(inView);
                goto MERGE_MASK_ERROR;
            }

            if (!ppMergeFileOpenInput(config, view, i)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to open file %d", i);
                psFree(inView);
                goto MERGE_MASK_ERROR;
            }

            if (inCell->readouts->n > 1) {
                psError(PS_ERR_BAD_PARAMETER_VALUE, true,
                        "File %d chip %d cell %d contains more than one readout (%ld)",
                        i, inView->chip, inView->cell, inCell->readouts->n);
                psFree(inView);
                goto MERGE_MASK_ERROR;
            }

            pmReadout *readout;
            if (inCell->readouts && inCell->readouts->n == 1) {
                readout = psMemIncrRefCounter(inCell->readouts->data[0]); ///< Input readout
            } else {
                readout = pmReadoutAlloc(inCell);
            }

            if (!ppMergeFileReadInput(config, readout, i, 0)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to read readout %d", i);
                psFree(inView);
                psFree(readout);
                goto MERGE_MASK_ERROR;
            }

            pmReadout *outRO = NULL;    ///< Output readout
            if (outCell->readouts && outCell->readouts->n == 1) {
                outRO = psMemIncrRefCounter(outCell->readouts->data[0]);
            } else {
                outRO = pmReadoutAlloc(outCell);
            }
            psImage *outMask = outRO->mask;    ///< Output mask image (for iterative generation of mask)

            psImage *image = readout->image, *mask = readout->mask; ///< Image and mask
            int numCols = readout->image->numCols, numRows = readout->image->numRows; ///< Image size
            int numPix = numCols * numRows; ///< Number of pixels
            int num = PS_MIN(numPix, sample / numCells); ///< Number of values to add
            if (!chipStats) {
                valueIndex = 0;
            }
            for (int i = 0; i < num; i++) {
                int pixel = numPix * psRandomUniform(rng);
                int x = pixel % numCols;
                int y = pixel / numCols;
                if (mask && (mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] & maskValRaw)) continue;
                if (outMask && (outMask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] & maskValOut)) continue;
                if (!isfinite(image->data.F32[y][x])) continue;

                values->data.F32[valueIndex++] = image->data.F32[y][x];
            }

            if (!chipStats) {
                values->n = valueIndex;
                if (!psVectorStats(statistics, values, NULL, NULL, 0)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to do statistics on readout.");
                    psFree(inView);
                    psFree(readout);
                    goto MERGE_MASK_ERROR;
                }

                float mean = psStatsGetValue(statistics, meanStat);
                float stdev = psStatsGetValue(statistics, stdevStat);

                // this function increments the count for each suspect pixel in each input plane
                // maskValRaw is used to test for valid input pixels
                if (!pmMaskFlagSuspectPixels(outRO, readout, mean, stdev, maskSuspect, maskValRaw)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to find suspect values in file %d", i);
                    psFree(inView);
                    psFree(readout);
                    goto MERGE_MASK_ERROR;
                }
                pmCellFreeData(inCell);

                if (!pmFPAfileIOChecks(config, inView, PM_FPA_AFTER)) {
                    psFree(inView);
                    psFree(readout);
                    goto MERGE_MASK_ERROR;
                }
            }
            psFree(readout);
            psFree(outRO);
        }

        // Additional run through cells if we want chip-level statistics
        if (chipStats && valueIndex > 0) {
            values->n = valueIndex;
            if (!psVectorStats(statistics, values, NULL, NULL, 0)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to do statistics on chip.");
                goto MERGE_MASK_ERROR;
            }
            inView->cell = -1;
            while ((inCell = pmFPAviewNextCell(inView, inFPA, 1))) {
                // the output FPA structure carries the information about which cells to process
                pmCell *outCell = pmFPAfileThisCell(config->files, inView, "PPMERGE.OUTPUT.MASK"); // Output cell
                if (!outCell->process) continue;

                pmHDU *hdu = pmHDUFromCell(inCell); ///< HDU for cell
                if (!hdu || hdu->blankPHU) continue;

                pmReadout *readout = inCell->readouts->data[0]; ///< Readout of interest

                inView->readout = 0;
                pmReadout *outRO = pmFPAfileThisReadout(config->files, inView, "PPMERGE.OUTPUT.MASK");

                float mean = psStatsGetValue(statistics, meanStat);
                float stdev = psStatsGetValue(statistics, stdevStat);

                if (!pmMaskFlagSuspectPixels(outRO, readout, mean, stdev, maskSuspect, maskValRaw)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to find suspect values in file %d", i);
                    goto MERGE_MASK_ERROR;
                }
                pmCellFreeData(inCell);

                inView->readout = -1;
                if (!pmFPAfileIOChecks(config, inView, PM_FPA_AFTER)) {
                    psFree(inView);
                    goto MERGE_MASK_ERROR;
                }
            }
        }
    }
    psFree(inView);
    psFree(statistics); statistics = NULL;
    psFree(values); values = NULL;


    // Another run through the chip to threshold on the suspects
    pmFPAfileActivate(config->files, false, NULL);
    if (writeOut) {
        ppMergeFileActivate(config, PPMERGE_FILES_OUTPUT, true);
    }
    pmFPA *outFPA = outChip->parent;    ///< Output FPA
    pmCell *outCell;                    ///< Output cell
    pmFPAview *outView = pmFPAviewAlloc(0); ///< View into output FPA
    *outView = *view;
    while ((outCell = pmFPAviewNextCell(outView, outFPA, 1))) {

        // skip inactive cells
        if (!outCell->process) continue;

        pmHDU *hdu = pmHDUFromCell(outCell); ///< HDU for cell
        if (!hdu || hdu->blankPHU) continue;

        psTrace("ppMerge", 1, "Getting bad pixels for chip %d cell %d", outView->chip, outView->cell);

        assert(outCell->readouts && outCell->readouts->n == 1);
        pmReadout *outRO = outCell->readouts->data[0]; ///< Output readout

        if (smoothSuspect) {
            // XXX test output of suspect pixel image
            psImage *suspects = psMetadataLookupPtr(NULL, outRO->analysis, PM_MASK_ANALYSIS_SUSPECT); // Suspect img
            assert (suspects);
            psImageSmooth (suspects, smoothScale, 3); // extend smoothing region to 3-sigma
        }

        // set the bad pixels to the value 'maskVal'
        if (!pmMaskIdentifyBadPixels(outRO, maskValOut, maskBad, maskMode)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to mask bad pixels");
            goto MERGE_MASK_ERROR;
        }

        // Supplementary outputs
        {
            // The counts image is fairly useless, but it preserves the model
            pmCell *countsCell = pmFPAfileThisCell(config->files, outView, "PPMERGE.OUTPUT.COUNT");
            pmReadout *countsRO = pmReadoutAlloc(countsCell); ///< Readout with count of inputs per pixel
            countsRO->image = psImageAlloc(outRO->mask->numCols, outRO->mask->numRows, PS_TYPE_F32);
            psImageInit(countsRO->image, numFiles);
            countsRO->data_exists = countsCell->data_exists = countsCell->parent->data_exists = true;
            psFree(countsRO);

            pmCell *sigmaCell = pmFPAfileThisCell(config->files, outView, "PPMERGE.OUTPUT.SIGMA");
            pmReadout *sigmaRO = pmReadoutAlloc(sigmaCell); ///< Readout with suspect image
            psImage *suspect = psMetadataLookupPtr(NULL, outRO->analysis, PM_MASK_ANALYSIS_SUSPECT);
            sigmaRO->image = psImageCopy(sigmaRO->image, suspect, PS_TYPE_F32);
            psMetadataRemoveKey(outRO->analysis, PM_MASK_ANALYSIS_SUSPECT);
            sigmaRO->data_exists = sigmaCell->data_exists = sigmaCell->parent->data_exists = true;
            psFree(sigmaRO);
        }

        if (maskGrow > 0) {
            psImage *grown = psImageGrowMask(NULL, outRO->mask, maskValOut, maskGrow, maskValOut); ///< Grown mask
            psFree(outRO->mask);
            outRO->mask = grown;
        }

        if (writeOut) {
            if (!pmFPAfileIOChecks(config, outView, PM_FPA_BEFORE)) {
                psFree(outView);
                goto MERGE_MASK_ERROR;
            }

            outRO->data_exists = outCell->data_exists = outChip->data_exists = true;

            // Average concepts
            psList *cells = psListAlloc(NULL); ///< List of cells, for concept averaging
            for (int i = 0; i < numFiles; i++) {
                pmFPAfile *inFile = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", i); ///< Input file
                pmCell *inCell = pmFPAviewThisCell(outView, inFile->fpa); ///< Input cell
                psListAdd(cells, PS_LIST_TAIL, inCell);
            }
            if (!pmConceptsAverageCells(outCell, cells, NULL, NULL, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to average cell concepts.");
                psFree(cells);
                psFree(outRO);
                psFree(outView);
                return false;
            }
            psFree(cells);

            // Statistics on the merged cell using a fake image
            outRO->image = psImageAlloc(outRO->mask->numCols, outRO->mask->numRows, PS_TYPE_F32);
            psImageInit(outRO->image, 1.0);
            if (!ppStatsFPA(stats, outRO->parent->parent->parent, outView, maskValOut, config)) {
                psError(PS_ERR_UNEXPECTED_NULL, true, "Unable to generate stats for image.");
                psFree(outRO);
                psFree(outView);
                return false;
            }
            psFree(outRO->image);
            outRO->image = NULL;

            // Write
            if (!pmFPAfileIOChecks(config, outView, PM_FPA_AFTER)) {
                psFree(outView);
                goto MERGE_MASK_ERROR;
            }
        }
    }
    psFree(outView);

    ppMergeFileActivate(config, PPMERGE_FILES_ALL, true);

    return true;


MERGE_MASK_ERROR:
    psFree(statistics);
    psFree(values);
    return false;
}

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

    bool mdok;                          // Status of MD lookup
    int numFiles = psMetadataLookupS32(NULL, config->arguments, "INPUTS.NUM"); // Number of inputs
    bool haveMasks = psMetadataLookupBool(&mdok, config->arguments, "INPUTS.MASKS"); // Do we have masks?
    bool haveVariances = psMetadataLookupBool(&mdok, config->arguments, "INPUTS.VARIANCES"); // Got variances?
    int iter = psMetadataLookupS32(NULL, config->arguments, "ITER"); // Number of rejection iterations

    PS_ASSERT_INT_POSITIVE(iter, false);

    pmFPAview *view = pmFPAviewAlloc(0); ///< View to component of interest
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS, 0); ///< Random number generator

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

    psArray *inputs = ppMergeFileDataLevel(config, "PPMERGE.INPUT", PM_FPA_LEVEL_READOUT); ///< Input images
    psFree(inputs);
    if (haveMasks) {
        psArray *masks = ppMergeFileDataLevel(config, "PPMERGE.INPUT.MASK", PM_FPA_LEVEL_READOUT);
        psFree(masks);
    }
    if (haveVariances) {
        psArray *variances = ppMergeFileDataLevel(config, "PPMERGE.INPUT.VARIANCE", PM_FPA_LEVEL_READOUT);
        psFree(variances);
    }

    if (!ppMergeFileActivate(config, PPMERGE_FILES_INPUT, true)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to activate files.");
        goto PPMERGE_MASK_ERROR;
    }
    if (!ppMergeFileActivate(config, PPMERGE_FILES_OUTPUT, true)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to activate files.");
        goto PPMERGE_MASK_ERROR;
    }

    // XXX this function should use pmConfigMaskReadHeader () to get the named values defined
    // for the input masks.

    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

    if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
        goto PPMERGE_MASK_ERROR;
    }
    pmChip *outChip;                    ///< Chip of interest
    while ((outChip = pmFPAviewNextChip(view, outFPA, 1))) {

        if (!outChip->process) continue;

        if (!pmFPAfileIOChecks(config, view, PM_FPA_BEFORE)) {
            goto PPMERGE_MASK_ERROR;
        }

        for (int i = 0; i < iter; i++) {
            if (!mergeMask(config, view, (i == iter - 1), rng, stats)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to merge chip %d", view->chip);
                goto PPMERGE_MASK_ERROR;
            }
        }

        if (outChip->data_exists) {
            psList *inChips = psListAlloc(NULL);
            for (int i=0; i < numFiles; i++) {
                pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", i); // Input file
                pmChip *chip = pmFPAviewThisChip(view, file->fpa);
                psListAdd(inChips, PS_LIST_TAIL, chip);
            }

            // XXX I need to call pmConfigMaskWriteHeader for the PHU somewhere, after it is created!

            if (!pmConceptsAverageChips(outChip, inChips, true)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to average Chip concepts.");
                psFree(inChips);
                goto PPMERGE_MASK_ERROR;
            }
            psFree(inChips);
        }
        if (!pmFPAfileIOChecks(config, view, PM_FPA_AFTER)) {
            goto PPMERGE_MASK_ERROR;
        }
    }

    psList *fpaList = psListAlloc(NULL);///< List of FPAs for concept averaging
    for (int i = 0; i < numFiles; i++) {
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPMERGE.INPUT", i); ///< Input file
        psListAdd(fpaList, PS_LIST_TAIL, file->fpa);
    }
    if (!pmConceptsAverageFPAs(outFPA, fpaList)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to average FPA concepts.");
        psFree(fpaList);
        goto PPMERGE_MASK_ERROR;
    }
    psFree(fpaList);

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

    psFree(stats);
    psFree(view);
    psFree(rng);

    return true;

PPMERGE_MASK_ERROR:
    psFree(stats);
    psFree(view);
    psFree(rng);
    return false;
}

