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

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

#include "ppStats.h"
#include "ppStatsLoop.h"


static void getMetadata(psMetadata *target, // Target for metadata
                        psMetadata *source, // Source for metadata
                        psList *list    // List containing keywords
    )
{
    psListIterator *iterator = psListIteratorAlloc(list, PS_LIST_HEAD, false); // Iterator
    psString name;                      // Name from iteration
    while ((name = psListGetAndIncrement(iterator))) {
        psMetadataItem *item = psMetadataLookup(source, name); // Item of interest, or NULL
        if (item) {
            psMetadataAddItem(target, item, PS_LIST_TAIL, 0);
        }
    }
    psFree(iterator);
    return;
}

static bool doThis(psList *toDoList,    // List of things to do
                   const char *this     // The name of "this"
    )
{
    if (psListLength(toDoList) == 0) {
        // No list --- do everything
        return true;
    }

    psListIterator *iterator = psListIteratorAlloc(toDoList, PS_LIST_HEAD, false); // Iterator
    psString test;                      // Test string, from iteration
    while ((test = psListGetAndIncrement(iterator))) {
        if (strcmp(this, test) == 0) {
            // It's in the list --- do it
            psFree(iterator);
            return true;
        }
    }
    psFree(iterator);
    // Couldn't find it --- don't do it
    return false;
}

static void addToHierarchy(psMetadata *source, // Source to add
                           psMetadata *target, // Target to which to add
                           const char *name, // Name of source
                           const char *comment // Comment for source
    )
{
    if (psListLength(source->list) > 0 && !psMetadataLookup(target, name)) {
        psMetadataAdd(target, PS_LIST_TAIL, name, PS_DATA_METADATA,
                      comment, source);
    }
    return;
}


static psExit cellStats(psMetadata *chipResults, // Metadata holding the chip results
                        pmCell *cell,     // Cell for which to get statistics
                        psFits *fits,     // FITS file handle
                        ppStatsData *data,// The data
                        const pmConfig *config // Configuration
    )
{
    assert(chipResults);
    assert(cell);
    assert(data);
    assert(config);

    const char *cellName = psMetadataLookupStr(NULL, cell->concepts, "CELL.NAME"); // Name of cell

    // Check to see if this is a cell of interest
    if (!doThis(data->cells, cellName)) {
        return PS_EXIT_SUCCESS;
    }
    pmHDU *hdu = pmHDUFromCell(cell); // HDU for cell
    if (!hdu || hdu->blankPHU) {
        if (fits) {
            // No HDU means there's no data in this cell
            return PS_EXIT_SUCCESS;
        } else {
            psError(PS_ERR_UNKNOWN, false, "Can't find HDU for cell\n");
            return PS_EXIT_CONFIG_ERROR;
        }
    }

    // Cell-level results
    bool mdok;                          // Status of MD lookup
    psMetadata *cellResults = psMemIncrRefCounter(psMetadataLookupMetadata(&mdok, chipResults, cellName));
    if (!mdok || !cellResults) {
        cellResults = psMetadataAlloc();
    }

    if (psListLength(data->headers) > 0 && cell->hdu) {
        if (fits && !pmCellReadHeader(cell, fits)) {
            psError (PS_ERR_IO, false, "trouble reading cell header\n");
            psFree(cellResults);
            return PS_EXIT_DATA_ERROR;
        }
        pmHDU *hdu = cell->hdu;     // HDU for headers
        getMetadata(cellResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits && cell->hdu && !pmCellReadHeader(cell, fits)) {
            psError (PS_ERR_IO, false, "trouble reading cell header\n");
            psFree(cellResults);
            return PS_EXIT_DATA_ERROR;
        }
        pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_ALL, false, config->database);
        getMetadata(cellResults, cell->concepts, data->concepts);
    }

    if (!data->doStats && psListLength(data->summary)) {
        // Nothing further to do --- don't want to waste our time reading the data
        goto cellDone;
    }

    if (fits && !pmCellRead(cell, fits, config->database)) {
        psError (PS_ERR_IO, false, "trouble reading cell data\n");
        psFree(cellResults);
        return PS_EXIT_DATA_ERROR;
    }

    psArray *readouts = cell->readouts; // Array of component readouts
    if (readouts->n == 0) {
        psLogMsg(__func__, PS_LOG_WARN, "No readouts present in cell %s --- skipping\n", cellName);
        goto cellDone;
    }
    if (readouts->n > 1) {
        psLogMsg(__func__, PS_LOG_WARN, "Multiple readouts (%ld) present in cell %s --- "
                 "using only the first.\n", readouts->n, cellName);
    }
    pmReadout *readout = readouts->data[0]; // The readout of interest
    if (!readout->image) {
        psLogMsg(__func__, PS_LOG_WARN, "No image associated with readout in cell %s --- ignored.\n", cellName);
        goto cellDone;
    }

    // Measure basic image statistics (means, stdevs, etc).
    if (data->sample <= 0.0) {
        if (!psImageStats(data->stats, readout->image, readout->mask, data->maskVal)) {
            psWarning("Unable to perform statistics on cell %s --- ignored.\n", cellName);
            goto statsDone;
        }
    } else {
        // Apply sampling
        psImage *image = readout->image; // The image of interest
        psImage *mask = readout->mask; // The mask image
        int numSamples = data->sample * image->numCols * image->numRows; // Number of samples
        int sampleSpace = 1.0 / data->sample; // Space between samples
        psVector *sampleValues = psVectorAlloc(numSamples, PS_TYPE_F32); // Vector of samples
        psVector *sampleMask = NULL;  // Corresponding mask
        if (mask) {
            sampleMask = psVectorAlloc(numSamples, PS_TYPE_U8);
        }
        for (int i = 0; i < numSamples; i++) {
            int j = i * sampleSpace;
            int y = j / image->numCols;
            int x = j % image->numCols;
            sampleValues->data.F32[i] = image->data.F32[y][x];
            if (mask) {
                sampleMask->data.U8[i] = mask->data.U8[y][x];
            }
        }
        if (!psVectorStats(data->stats, sampleValues, NULL, sampleMask, data->maskVal)) {
            psLogMsg(__func__, PS_LOG_WARN, "Unable to perform statistics on cell %s --- "
                     "ignored.\n", cellName);
            psFree(sampleValues);
            psFree(sampleMask);
            goto statsDone;
        }
        psFree(sampleValues);
        psFree(sampleMask);
    }

#define WRITE_STAT(SYMBOL, NAME, SOURCE) \
    if (data->stats->options & SYMBOL) { \
        psMetadataAddF32(cellResults, PS_LIST_TAIL, NAME, 0, NULL, data->stats->SOURCE); \
    }

    WRITE_STAT(PS_STAT_SAMPLE_MEAN,     "SAMPLE_MEAN",   sampleMean);
    WRITE_STAT(PS_STAT_SAMPLE_MEDIAN,   "SAMPLE_MEDIAN", sampleMedian);
    WRITE_STAT(PS_STAT_SAMPLE_STDEV,    "SAMPLE_STDEV",  sampleStdev);
    WRITE_STAT(PS_STAT_SAMPLE_QUARTILE, "SAMPLE_LQ",     sampleLQ);
    WRITE_STAT(PS_STAT_SAMPLE_QUARTILE, "SAMPLE_UQ",     sampleUQ);
    WRITE_STAT(PS_STAT_ROBUST_MEDIAN,   "ROBUST_MEDIAN", robustMedian);
    WRITE_STAT(PS_STAT_ROBUST_STDEV,    "ROBUST_STDEV",  robustStdev);
    WRITE_STAT(PS_STAT_ROBUST_QUARTILE, "ROBUST_LQ",     robustLQ);
    WRITE_STAT(PS_STAT_ROBUST_QUARTILE, "ROBUST_UQ",     robustUQ);
    WRITE_STAT(PS_STAT_ROBUST_QUARTILE, "ROBUST_N50",    robustN50);
    WRITE_STAT(PS_STAT_FITTED_MEAN,     "FITTED_MEAN",   fittedMean);
    WRITE_STAT(PS_STAT_FITTED_STDEV,    "FITTED_STDEV",  fittedStdev);
    WRITE_STAT(PS_STAT_CLIPPED_MEAN,    "CLIPPED_MEAN",  clippedMean);
    WRITE_STAT(PS_STAT_CLIPPED_STDEV,   "CLIPPED_STDEV", clippedStdev);

    // measure other types of statistics tests

statsDone:
    // count saturated pixels
    if (psListLength(data->summary) > 0) {
        bool get_nSatPixels = false;
        bool get_fSatPixels = false;

        psListIterator *iterator = psListIteratorAlloc(data->summary, PS_LIST_HEAD, false);
        psString choice;
        while ((choice = psListGetAndIncrement(iterator))) {
            if (!strcasecmp (choice, "SAT_PIXEL_NUM"))  get_nSatPixels = true;
            if (!strcasecmp (choice, "SAT_PIXEL_FRAC")) get_fSatPixels = true;
        }

        if (!get_nSatPixels && !get_fSatPixels) {
            goto cellDone;
        }

        // Get the "concepts" of interest
        float saturation = psMetadataLookupF32(&mdok, cell->concepts, "CELL.SATURATION"); // Saturation level
        if (!mdok || isnan(saturation)) {
            psLogMsg(__func__, PS_LOG_WARN, "CELL.SATURATION is not set --- unable to measure N_SAT_PIXELS.\n");
            if (get_nSatPixels) psMetadataAddS32(cellResults, PS_LIST_TAIL, "SAT_PIXEL_NUM", 0, NULL, 0);
            if (get_fSatPixels) psMetadataAddF32(cellResults, PS_LIST_TAIL, "SAT_PIXEL_FRAC", 0, NULL, NAN);
            goto cellDone;
        }

        int nSatPixels = 0;
        for (int j = 0; j < readout->image->numRows; j++) {
            for (int i = 0; i < readout->image->numCols; i++) {
                if (readout->image->data.F32[j][i] >= saturation) {
                    nSatPixels ++;
                }
            }
        }
        if (get_nSatPixels) psMetadataAddS32(cellResults, PS_LIST_TAIL, "SAT_PIXEL_NUM", 0, NULL, nSatPixels);
        if (get_fSatPixels) psMetadataAddF32(cellResults, PS_LIST_TAIL, "SAT_PIXEL_FRAC", 0, NULL, nSatPixels / (double)(readout->image->numRows * readout->image->numCols));
    }

cellDone:
    // Add the cell results to the chip
    addToHierarchy(cellResults, chipResults, cellName, "Results for cell");
    psFree (cellResults);
    if (fits) {
        pmCellFreeData(cell);
    }
    return PS_EXIT_SUCCESS;
}

static psExit chipStats(psMetadata *fpaResults, // Metadata holding the fpa results
                        pmChip *chip,     // Chip for which to get statistics
                        psFits *fits,     // FITS file handle
                        pmFPAview *view,  // View for analysis
                        ppStatsData *data,// The data
                        const pmConfig *config // Configuration
    )
{
    assert(fpaResults);
    assert(chip);
    assert(view);
    assert(data);
    assert(config);

    psExit result = PS_EXIT_SUCCESS;

    const char *chipName = psMetadataLookupStr(NULL, chip->concepts, "CHIP.NAME"); // Name of chip

    // Check to see if this is a chip of interest
    if (!doThis(data->chips, chipName)) {
        return PS_EXIT_SUCCESS;
    }

    psArray *cells = chip->cells;       // Array of cells
    if (view->cell >= cells->n) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Desired cell view (%d) doesn't match "
                "number of cells (%ld)\n", view->cell, cells->n);
        return PS_EXIT_CONFIG_ERROR;
    }

    // Chip-level results
    bool mdok;                          // Status of MD lookup
    psMetadata *chipResults = psMemIncrRefCounter(psMetadataLookupMetadata(&mdok, fpaResults, chipName));
    if (!mdok || !chipResults) {
        chipResults = psMetadataAlloc();
    }

    if (psListLength(data->headers) > 0 && chip->hdu) {
        if (fits && !pmChipReadHeader(chip, fits)) {
            psError (PS_ERR_IO, false, "trouble reading chip header\n");
            psFree(chipResults);
            return PS_EXIT_DATA_ERROR;
        }
        pmHDU *hdu = chip->hdu;     // HDU for headers
        getMetadata(chipResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits && chip->hdu && !pmChipReadHeader(chip, fits)) {
            psError (PS_ERR_IO, false, "trouble reading chip header\n");
            psFree(chipResults);
            return PS_EXIT_DATA_ERROR;
        }
        pmConceptsReadChip(chip, PM_CONCEPT_SOURCE_ALL, false, false, config->database);
        getMetadata(chipResults, chip->concepts, data->concepts);
    }

    if (view->cell >= 0) {
        pmCell *cell = cells->data[view->cell]; // Cell of interest
        result = cellStats(chipResults, cell, fits, data, config);
        if (result != PS_EXIT_SUCCESS) {
            psError(PS_ERR_UNKNOWN, false, "trouble with cell stats for %d\n", view->cell);
        }
        addToHierarchy(chipResults, fpaResults, chipName, "Results for chip");
        psFree (chipResults);
        if (fits) {
            pmChipFreeData(chip);
        }
        return result;
    }

    // Iterate over cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // Cell of interest
        result = cellStats(chipResults, cell, fits, data, config);
        if (result != PS_EXIT_SUCCESS) {
            psError(PS_ERR_UNKNOWN, false, "trouble with cell stats for %d\n", i);
            if (fits) {
                pmChipFreeData(chip);
            }
            psFree(chipResults);
            return result;
        }
    }

    addToHierarchy(chipResults, fpaResults, chipName, "Results for chip");
    psFree (chipResults);
    if (fits) {
        pmChipFreeData(chip);
    }
    return PS_EXIT_SUCCESS;
}

psMetadata *ppStatsLoop(psExit *result,
                        ppStatsData *data, // The data
                        const pmConfig *config // Configuration
    )
{
    PS_ASSERT_PTR_NON_NULL(data, NULL);
    pmFPA *fpa = data->fpa;             // FPA to analyse
    psFits *fits = data->fits;          // FITS file handle
    PS_ASSERT_PTR_NON_NULL(fpa, NULL);

    *result = PS_EXIT_SUCCESS;

    // allocate or bump the ref counter (so we can just free below)
    pmFPAview *view = (data->view) ? psMemIncrRefCounter(data->view) : pmFPAviewAlloc(0);

    // allocate a new one if needed
    psMetadata *newResults = psMetadataAlloc();

    // Iterate through the FPA
    if (psListLength(data->headers) > 0 && fpa->hdu) {
        if (fits && !pmFPAReadHeader(fpa, fits)) {
            psError(PS_ERR_IO, false, "Unable to read header for FPA.");
            psFree(view);
            psFree(newResults);
            *result = PS_EXIT_DATA_ERROR;
            return NULL;
        }
        pmHDU *hdu = fpa->hdu;          // HDU for headers
        getMetadata(newResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits && fpa->hdu && !pmFPAReadHeader(fpa, fits)) {
            psError(PS_ERR_IO, false, "Unable to read header for FPA.");
            psFree(view);
            psFree(newResults);
            *result = PS_EXIT_DATA_ERROR;
            return NULL;
        }
        pmConceptsReadFPA(fpa, PM_CONCEPT_SOURCE_ALL, false, config->database);
        getMetadata(newResults, fpa->concepts, data->concepts);
    }

    psArray *chips = fpa->chips;        // Array of chips
    if (view->chip >= 0) {
        pmChip *chip = chips->data[view->chip]; // Chip of interest
        *result = chipStats(newResults, chip, fits, view, data, config);
        if (*result != PS_EXIT_SUCCESS) {
            psError(PS_ERR_UNKNOWN, false, "trouble with stats for cell %d\n", view->cell);
            psFree (view);
            psFree (newResults);
            return NULL;
        }
        if (fits) {
            pmFPAFreeData(fpa);
        }
        psFree(view);
        return newResults;
    }

    // Iterate over cells
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // Chip of interest
        *result = chipStats(newResults, chip, fits, view, data, config);
        if (*result != PS_EXIT_SUCCESS) {
            psError(PS_ERR_UNKNOWN, false, "trouble with stats for chip %d\n", i);
            psFree (view);
            psFree (newResults);
            return NULL;
        }
    }

    if (fits) {
        pmFPAFreeData(fpa);
    }
    psFree(view);

    return newResults;
}
