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

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

#include "ppImageOptions.h"
#include "ppImageDetrendFringe.h"

bool ppImageDetrendFringeMeasure(pmReadout *readout, pmCell *fringe, const ppImageOptions *options)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(fringe, false);
    PS_ASSERT_PTR_NON_NULL(options, false);

    // Reference fringe measurements
    psArray *references = psMemIncrRefCounter(psMetadataLookupPtr(NULL, fringe->analysis,
                                                                  "FRINGE.MEASUREMENTS"));
    if (!references) {
        references = pmFringesParse(fringe); // Reference fringes
        if (!references) {
            psError(PS_ERR_IO, false, "Unable to find fringe references.\n");
            return false;
        }
        psMetadataAdd(fringe->analysis, PS_LIST_TAIL, "FRINGE.MEASUREMENTS", PS_DATA_UNKNOWN,
                      "Fringe measurements", references);
    }

    pmFringeStats *reference = references->data[0]; // Take the first as representative
    pmFringeRegions *regions = reference->regions; // Regions to measure
    pmFringeStats *measurements = pmFringeStatsMeasure(regions, readout, options->maskValue); // Fringe stats

    // Normalise measurements by the exposure time
    bool mdok;                          // Status of MD lookup
    float expTime = psMetadataLookupF32(&mdok, readout->parent->concepts, "CELL.EXPOSURE"); // Exp. time
    if (!mdok || !isfinite(expTime)) {
        psError(PS_ERR_UNKNOWN, false, "CELL.EXPOSURE is not set for --- can't normalise fringes\n");
        psFree(measurements);
        return false;
    }
    if (expTime == 0) {
        psWarning("Exposure time is zero --- are you sure you want fringe subtraction?\n");
        expTime = 1.0;
    }
    psBinaryOp(measurements->f, measurements->f, "*", psScalarAlloc(1.0 / expTime, PS_TYPE_F32));
    psBinaryOp(measurements->df, measurements->df, "*", psScalarAlloc(1.0 / expTime, PS_TYPE_F32));

    // Science fringe measurements
    pmFringeStats *previous = psMetadataLookupPtr(NULL, readout->parent->analysis, "FRINGE.MEASUREMENTS");
    if (previous) {
        // Multiple readouts: concatenate
        psArray *concatenate = psArrayAlloc(2); // Array to hold fringes

        // Concatenate science measurements
        concatenate->data[0] = previous;
        concatenate->data[1] = measurements;
        pmFringeStats *new = pmFringeStatsConcatenate(concatenate, NULL, NULL); // New measurements
        psFree(measurements);
        measurements = new;

        // Concatenate reference measurements (duplication, so the science and reference line up)
        for (int i = 0; i < references->n; i++) {
            concatenate->data[0] = concatenate->data[1] = references->data[i];
            pmFringeStats *refNew = pmFringeStatsConcatenate(concatenate, NULL, NULL);
            psFree(references->data[i]);
            references->data[i] = refNew;
        }
        concatenate->data[0] = concatenate->data[1] = NULL;
        psFree(concatenate);
    }

    psMetadataAdd(readout->parent->analysis, PS_LIST_TAIL, "FRINGE.MEASUREMENTS",
                  PS_DATA_UNKNOWN | PS_META_REPLACE, "Fringe measurements", measurements);
    psFree(measurements);

    psFree(references);

    return true;
}


// Pull the fringes out of the cell analysis FRINGE.MEASUREMENTS for a chip
static psArray *getFringes(const pmChip *chip // Chip of interest
    )
{
    psArray *cells = chip->cells;       // Component cells
    psArray *fringes = psArrayAlloc(cells->n); // Fringes, to return
    for (int i = 0; i < cells->n; i++) {
        pmCell *cell = cells->data[i];  // Cell of interest
        fringes->data[i] = psMemIncrRefCounter(psMetadataLookupPtr(NULL, cell->analysis,
                                                                   "FRINGE.MEASUREMENTS"));
    }

    return fringes;
}


// Solve the fringe system: we have science fringe measurements for each cell, and an array of reference
// fringe measurements for each cell.  Need to concatenate these together first, and then solve.
bool ppImageDetrendFringeSolve(pmChip *scienceChip, const pmChip *refChip, const ppImageOptions *options)
{
    PS_ASSERT_PTR_NON_NULL(scienceChip, NULL);
    PS_ASSERT_PTR_NON_NULL(refChip, NULL);
    PS_ASSERT_PTR_NON_NULL(options, NULL);

    psArray *science = getFringes(scienceChip); // Fringe measurements on science chip
    pmFringeStats *scienceCat = pmFringeStatsConcatenate(science, NULL, NULL); // Science fringes
    psFree(science);

    // Need to transform the array of cells each with an array of fringes --> array of fringes for the chip as
    // a whole
    psArray *references = getFringes(refChip); // Fringe measurements on reference chip
    int numRefs = ((psArray*)references->data[0])->n; // Number of reference fringes
    psArray *referencesCat = psArrayAlloc(numRefs); // Reference fringes
    for (int i = 0; i < numRefs; i++) { // Iterate over fringes
        psArray *refs = psArrayAlloc(references->n); // Array of fringes for each cell
        for (int j = 0; j < references->n; j++) { // Iterate over cells
            psArray *ref = references->data[j]; // Array of references for this cell
            refs->data[j] = psMemIncrRefCounter(ref->data[i]);
        }
        referencesCat->data[i] = pmFringeStatsConcatenate(refs, NULL, NULL);
        psFree(refs);
    }
    psFree(references);

    // Now we can solve
    psTrace("ppImage", 3, "Solving fringe system...\n");
    pmFringeScale *solution = pmFringeScaleMeasure(scienceCat, referencesCat, options->fringeRej,
                                                   options->fringeIter, options->fringeKeep);

    psMetadataAdd(scienceChip->analysis, PS_LIST_TAIL, "FRINGE.SOLUTION", PS_DATA_UNKNOWN,
                  "Fringe solution", solution);
    psFree(solution);
    psFree(scienceCat);
    psFree(referencesCat);

    return true;
}


bool ppImageDetrendFringeGenerate(pmCell *science, pmCell *fringes)
{
    PS_ASSERT_PTR_NON_NULL(science, false);
    PS_ASSERT_PTR_NON_NULL(fringes, false);

    pmFringeScale *solution = psMetadataLookupPtr(NULL, science->parent->analysis, "FRINGE.SOLUTION");
    if (!solution) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find fringe solution.\n");
        return false;
    }
    assert(fringes->readouts->n == solution->nFringeFrames);

    bool mdok;                          // Status of MD lookup
    float expTime = psMetadataLookupF32(&mdok, science->concepts, "CELL.EXPOSURE"); // Exp. time
    if (!mdok || !isfinite(expTime)) {
        psError(PS_ERR_UNKNOWN, true, "CELL.EXPOSURE is not set --- can't renormalise fringes\n");
        return false;
    }
    if (expTime == 0) {
        psWarning("Exposure time is zero --- are you sure you want fringe subtraction?\n");
        expTime = 1.0;
    }

    pmHDU *hdu = pmHDUFromCell(science);// HDU  of interest

    // Construct the fringe image from the scale
    psTrace("ppImage", 3, "Generating fringe correction...\n");
    psImage *sumFringe = NULL; // Sum of the fringes
    for (int i = 0; i < solution->nFringeFrames; i++) {
        pmReadout *fringeRO = fringes->readouts->data[i]; // Fringe readout
        psImage *fringe = fringeRO->image; // Fringe image
        psTrace("ppImage", 5, "Scale for fringe component %d is %f\n",
                i, solution->coeff->data.F32[i + 1]);
        psBinaryOp(fringe, fringe, "*", psScalarAlloc(solution->coeff->data.F32[i + 1], PS_TYPE_F32));
        if (!sumFringe) {
            sumFringe = psImageCopy(NULL, fringe, PS_TYPE_F32);
        } else {
            psBinaryOp(sumFringe, sumFringe, "+", fringe);
        }

        psVector *md5 = psImageMD5(fringe); // md5 hash
        psString md5string = psMD5toString(md5); // String
        psFree(md5);
        psStringPrepend(&md5string, "Fringe image %d (scale %.3f) MD5: ",
                        i, solution->coeff->data.F32[i + 1]);
        psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                         md5string, "");
        psFree(md5string);
    }
    if (expTime != 1.0) {
        psBinaryOp(sumFringe, sumFringe, "*", psScalarAlloc(expTime, PS_TYPE_F32));
    }

    // Apply the correction to all readouts
    psArray *readouts = science->readouts; // Array of science readouts
    for (int i = 0; i < readouts->n; i++) {
        pmReadout *readout = readouts->data[i]; // Readout of interest
        if (!readout || !readout->data_exists) {
            continue;
        }

        // XXX: Make generic, so subregions may be subtracted as well
        if (!psBinaryOp(readout->image, readout->image, "-", sumFringe)) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to subtract fringe.\n");
            return false;
        }
    }
    psFree(sumFringe);

    psTime *time = psTimeGetNow(PS_TIME_TAI); // The time now, used for reporting
    psString timeString = psTimeToISO(time); // String with time
    psFree(time);
    psStringPrepend(&timeString, "Fringe correction completed at ");
    psMetadataAddStr(hdu->header, PS_LIST_TAIL, "HISTORY", PS_META_DUPLICATE_OK,
                     timeString, "");
    psFree(timeString);

    return sumFringe;
}

