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

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

#include "ppStack.h"
#include "ppStackLoop.h"

#define PHOT_SOURCE_MASK (PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_SATSTAR | PM_SOURCE_MODE_BLEND | \
                          PM_SOURCE_MODE_BADPSF | PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | \
                          PM_SOURCE_MODE_CR_LIMIT | PM_SOURCE_MODE_EXT_LIMIT) // Mask to apply to sources

bool ppStackInputPhotometer(const pmReadout *ro, const psArray *sources, const pmConfig *config)
{
    PM_ASSERT_READOUT_NON_NULL(ro, false);
    PS_ASSERT_ARRAY_NON_NULL(sources, false);

    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
    psAssert(recipe, "We've thrown an error on this before.");

    bool mdok;                          // Status of MD lookup

    float zpRadius = psMetadataLookupS32(&mdok, recipe, "PHOT.RADIUS"); // Radius for PHOT measurement
    if (!mdok) {
        psError(PS_ERR_UNKNOWN, true, "Unable to find PHOT.RADIUS in recipe");
        return false;
    }
    float zpSigma = psMetadataLookupF32(&mdok, recipe, "PHOT.SIGMA"); // Gaussian sigma for photometry
    if (!mdok) {
        psError(PS_ERR_UNKNOWN, true, "Unable to find PHOT.SIGMA in recipe");
        return false;
    }
    float zpFrac = psMetadataLookupF32(&mdok, recipe, "PHOT.FRAC"); // Fraction of good pixels for photometry
    if (!mdok) {
        psError(PS_ERR_UNKNOWN, true, "Unable to find PHOT.FRAC in recipe");
        return false;
    }

    psString maskValStr = psMetadataLookupStr(&mdok, recipe, "MASK.VAL"); // Name of bits to mask going in
    if (!mdok || !maskValStr) {
        psError(PS_ERR_UNKNOWN, false, "Unable to find MASK.VAL in recipe");
        return false;
    }
    psImageMaskType maskVal = pmConfigMaskGet(maskValStr, config); // Bits to mask

    psImage *image = ro->image, *mask = ro->mask; // Image and mask from readout

    // Measure background
    psRandom *rng = psRandomAlloc(PS_RANDOM_TAUS); // Random number generator
    psStats *stats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV); // Statistics
    if (!psImageBackground(stats, NULL, image, mask, maskVal, rng)) {
        psError(PS_ERR_UNKNOWN, false, "Unable to measure background for image");
        psFree(stats);
        psFree(rng);
        return false;
    }
    float bg = stats->robustMedian; // Background level
    psFree(stats);
    psFree(rng);

    // Photometer sources
    int numCols = image->numCols, numRows = image->numRows; // Size of image
    int numSources = sources->n; // Number of sources
    int numGood = 0;            // Number of good sources
    float maxMag = -INFINITY;   // Maximum magnitude
    for (int j = 0; j < numSources; j++) {
        pmSource *source = sources->data[j]; // Source of interest
        if (source->mode & PHOT_SOURCE_MASK || !isfinite(source->psfMag)) {
            source->psfMag = NAN;
            continue;
        }
        float x = source->peak->xf, y = source->peak->yf; // Coordinates of source
        int xCentral = x + 0.5, yCentral = y + 0.5; // Central pixel
        // Range of integration
        int xMin = PS_MAX(0, xCentral - zpRadius), xMax = PS_MIN(numCols - 1, xCentral + zpRadius);
        int yMin = PS_MAX(0, yCentral - zpRadius), yMax = PS_MIN(numRows - 1, yCentral + zpRadius);

        // Integrate
        double sumImage = 0.0, sumKernel = 0.0; // Sums from integration
        int numBadPix = 0;         // Number of bad pixels
        for (int v = yMin; v <= yMax; v++) {
            float dy2 = PS_SQR(y - v); // Distance from centroid
            for (int u = xMin; u <= xMax; u++) {
                if (mask->data.PS_TYPE_IMAGE_MASK_DATA[v][u] & maskVal) {
                    numBadPix++;
                    continue;
                }
                float dx2 = PS_SQR(x - u); // Distance from centroid
                double kernel = exp(-0.5 * (dx2 + dy2) / PS_SQR(zpSigma)); // Kernel value
                sumImage += (image->data.F32[v][u] - bg) * kernel;
                sumKernel += kernel;
            }
        }

        if (numBadPix > zpFrac * M_PI * PS_SQR(zpRadius) || !isfinite(sumImage)) {
            source->psfMag = NAN;
        } else {
            source->psfMag = - 2.5 * log10(sumImage / sumKernel * M_PI * PS_SQR(zpRadius));
            if (isfinite(source->psfMag)) {
                numGood++;
                maxMag = PS_MAX(maxMag, source->psfMag);
            }
        }
    }
    psLogMsg("ppStack", PS_LOG_INFO, "Photometered %d good sources in image; max mag %f", numGood, maxMag);

    return true;
}


// Preparation iteration: Load the sources, and get a target PSF model
bool ppStackPrepare(ppStackOptions *options, pmConfig *config)
{
    psAssert(options, "Require options");
    psAssert(config, "Require configuration");

    psMetadata *recipe = psMetadataLookupMetadata(NULL, config->recipes, PPSTACK_RECIPE); // ppStack recipe
    psAssert(recipe, "We've thrown an error on this before.");

    int num = options->num = psMetadataLookupS32(NULL, config->arguments, "INPUTS.NUM");
    options->sourceLists = psArrayAlloc(num); // Individual lists of sources for matching

    options->inputMask = psVectorAlloc(num, PS_TYPE_VECTOR_MASK); // Mask for inputs
    psVectorInit(options->inputMask, 0);

    if (psMetadataLookupBool(NULL, config->arguments, "HAVE.PSF")) {
        pmFPAfileActivate(config->files, false, NULL);
        ppStackFileActivation(config, PPSTACK_FILES_PREPARE, true);
        pmFPAview *view = ppStackFilesIterateDown(config);
        if (!view) {
            return false;
        }

        // Generate list of PSFs
        psMetadataIterator *fileIter = psMetadataIteratorAlloc(config->files, PS_LIST_HEAD,
                                                               "^PPSTACK.INPUT$"); // Iterator
        psMetadataItem *fileItem; // Item from iteration
        psArray *psfs = psArrayAlloc(num); // PSFs for PSF envelope
        int numCols = 0, numRows = 0;   // Size of image
        int index = 0;                  // Index for file
        while ((fileItem = psMetadataGetAndIncrement(fileIter))) {
            assert(fileItem->type == PS_DATA_UNKNOWN);
            pmFPAfile *inputFile = fileItem->data.V; // An input file
            pmChip *chip = pmFPAviewThisChip(view, inputFile->fpa); // The chip: holds the PSF
            pmPSF *psf = psMetadataLookupPtr(NULL, chip->analysis, "PSPHOT.PSF"); // PSF
            if (!psf) {
                psError(PS_ERR_UNKNOWN, false, "Unable to find PSF.");
                psFree(view);
                psFree(fileIter);
                psFree(psfs);
                return false;
            }
            psfs->data[index] = psMemIncrRefCounter(psf);
            psMetadataRemoveKey(chip->analysis, "PSPHOT.PSF");

            pmCell *cell = pmFPAviewThisCell(view, inputFile->fpa); // Cell of interest
            pmHDU *hdu = pmHDUFromCell(cell);
            assert(hdu && hdu->header);
            int naxis1 = psMetadataLookupS32(NULL, hdu->header, "NAXIS1"); // Number of columns
            int naxis2 = psMetadataLookupS32(NULL, hdu->header, "NAXIS2"); // Number of rows
            if (naxis1 <= 0 || naxis2 <= 0) {
                psError(PS_ERR_UNKNOWN, false, "Unable to determine size of image from PSF.");
                psFree(view);
                psFree(fileIter);
                psFree(psfs);
                return false;
            }
            if (numCols == 0 && numRows == 0) {
                numCols = naxis1;
                numRows = naxis2;
            }

            pmReadout *ro = pmFPAviewThisReadout(view, inputFile->fpa); // Readout with sources
            psArray *sources = psMetadataLookupPtr(NULL, ro->analysis, "PSPHOT.SOURCES"); // Sources
            if (!sources) {
                psError(PS_ERR_UNEXPECTED_NULL, false, "Unable to find sources in readout.");
                return NULL;
            }
            options->sourceLists->data[index] = psMemIncrRefCounter(sources);

            // Re-do photometry if we don't trust the source lists
            if (psMetadataLookupBool(NULL, recipe, "PHOT")) {
                psTrace("ppStack", 2, "Photometering input %d of %d....\n", index, num);
                pmFPAfileActivate(config->files, false, NULL);
                ppStackFileActivationSingle(config, PPSTACK_FILES_CONVOLVE, true, index);
                pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.INPUT", index); // File
                pmFPAview *photView = ppStackFilesIterateDown(config);
                if (!photView) {
                    psFree(view);
                    return false;
                }

                pmReadout *ro = pmFPAviewThisReadout(view, file->fpa); // Readout of interest

                if (!ppStackInputPhotometer(ro, sources, config)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to do photometry on input sources");
                    psFree(view);
                    psFree(photView);
                    return false;
                }

                psFree(photView);
                if (!ppStackFilesIterateUp(config)) {
                    psFree(view);
                    return false;
                }
                pmFPAfileActivate(config->files, false, NULL);
                ppStackFileActivation(config, PPSTACK_FILES_PREPARE, true);
            }

            index++;
        }
        psFree(fileIter);

        // Zero point calibration
        options->sumExposure = ppStackSourcesTransparency(options->sourceLists, options->inputMask,
                                                          view, config);
        if (!isfinite(options->sumExposure) || options->sumExposure <= 0) {
            psError(PS_ERR_UNKNOWN, false, "Unable to calculate transparency differences");
            psFree(view);
            return false;
        }

        // Generate target PSF
        options->psf = ppStackPSF(config, numCols, numRows, psfs, options->inputMask);
        psFree(psfs);
        if (!options->psf) {
            psError(PS_ERR_UNKNOWN, false, "Unable to determine output PSF.");
            psFree(view);
            return false;
        }
        psMetadataAddPtr(config->arguments, PS_LIST_TAIL, "PSF.TARGET", PS_DATA_UNKNOWN,
                         "Target PSF for stack", options->psf);

        pmChip *outChip = pmFPAfileThisChip(config->files, view, "PPSTACK.OUTPUT"); // Output chip
        psMetadataAddPtr(outChip->analysis, PS_LIST_TAIL, "PSPHOT.PSF", PS_DATA_UNKNOWN,
                         "Target PSF", options->psf);
        outChip->data_exists = true;

        psFree(view);

        if (!ppStackFilesIterateUp(config)) {
            return false;
        }
    }

    return true;
}
