#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);
    }
    psFree(source);                     // Drop reference
    return;
}


static void 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;
    }

    // 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)) {
            goto cellDone;
        }
        pmHDU *hdu = cell->hdu;     // HDU for headers
        getMetadata(cellResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits && !pmCellReadHeader(cell, fits)) {
            goto cellDone;
        }
        pmConceptsReadCell(cell, PM_CONCEPT_SOURCE_ALL, false, config->database);
        getMetadata(cellResults, cell->concepts, data->concepts);
    }

    if (!data->doStats) {
        // Nothing further to do --- don't want to waste our time reading the data
        if (psListLength(cellResults->list) > 0) {
            psMetadataAdd(chipResults, PS_LIST_TAIL, cellName, PS_DATA_METADATA,
                          "Results for cell", cellResults);
        }
        goto cellDone;
    }

    pmHDU *hdu = pmHDUFromCell(cell); // HDU for cell
    if (!hdu || hdu->blankPHU) {
        goto cellDone;
    }

    if (fits && !pmCellRead(cell, fits, config->database)) {
        psLogMsg(__func__, PS_LOG_WARN, "Unable to read cell %s\n", cellName);
        goto cellDone;
    }

    psArray *readouts = cell->readouts; // Array of component readouts
    if (readouts->n == 0) {
        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;
    }

    // Do the statistics
    if (data->sample <= 0.0) {
        if (!psImageStats(data->stats, readout->image, readout->mask, data->maskVal)) {
            psLogMsg(__func__, PS_LOG_WARN, "Unable to perform statistics on cell %s --- "
                     "ignored.\n", cellName);
            goto cellDone;
        }
    } 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 cellDone;
        }
        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);

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

static bool 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);

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

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

    // 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)) {
            goto chipDone;
        }
        pmHDU *hdu = chip->hdu;     // HDU for headers
        getMetadata(chipResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits && !pmChipReadHeader(chip, fits)) {
            goto chipDone;
        }
        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
        cellStats(chipResults, cell, fits, data, config);
        goto chipDone;
    }

    // Iterate over cells
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // Cell of interest
        cellStats(chipResults, cell, fits, data, config);
    }

chipDone:
    addToHierarchy(chipResults, fpaResults, chipName, "Results for chip");
    if (fits) {
        pmChipFreeData(chip);
    }

    return true;
}


psMetadata *ppStatsLoop(psMetadata *fpaResults, // Metadata to hold the FPA results
                        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);

    pmFPAview *view = psMemIncrRefCounter(data->view); // View for analysis
    if (!view) {
        view = pmFPAviewAlloc(0);
    }

    if (!fpaResults) {
        fpaResults = psMetadataAlloc();
    }

    // Iterate through the FPA
    if (psListLength(data->headers) > 0 && fpa->hdu) {
        if (fits) {
            pmFPAReadHeader(fpa, fits);
        }
        pmHDU *hdu = fpa->hdu;          // HDU for headers
        getMetadata(fpaResults, hdu->header, data->headers);
    }
    if (psListLength(data->concepts) > 0) {
        if (fits) {
            pmFPAReadHeader(fpa, fits);
        }
        pmConceptsReadFPA(fpa, PM_CONCEPT_SOURCE_ALL, false, config->database);
        getMetadata(fpaResults, fpa->concepts, data->concepts);
    }

    psArray *chips = fpa->chips;        // Array of chips
    if (view->chip >= 0) {
        pmChip *chip = chips->data[view->chip]; // Chip of interest
        chipStats(fpaResults, chip, fits, view, data, config);
        goto fpaDone;
    }

    // Iterate over cells
    for (int i = 0; i < chips->n; i++) {
        pmChip *chip = chips->data[i];  // Chip of interest
        chipStats(fpaResults, chip, fits, view, data, config);
    }

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

    return fpaResults;
}
