#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(PPSTACK_ERR_CONFIG, true, "Unable to find PHOT.RADIUS in recipe");
        return false;
    }
    float zpSigma = psMetadataLookupF32(&mdok, recipe, "PHOT.SIGMA"); // Gaussian sigma for photometry
    if (!mdok) {
        psError(PPSTACK_ERR_CONFIG, 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(PPSTACK_ERR_CONFIG, 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(PPSTACK_ERR_CONFIG, 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(PPSTACK_ERR_DATA, 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);
    options->exposures = psVectorAlloc(options->num, PS_TYPE_F32);
    psVectorInit(options->exposures, NAN);

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

    psArray *psfs = psArrayAlloc(num); // PSFs for PSF envelope
    int numCols = 0, numRows = 0;   // Size of image
    options->sumExposure = 0.0;
    for (int i = 0; i < num; i++) {
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PPSTACK.INPUT", i); // File of interest
        pmCell *cell = pmFPAviewThisCell(view, file->fpa); // Cell of interest

        options->exposures->data.F32[i] = psMetadataLookupF32(NULL, cell->concepts, "CELL.EXPOSURE");
        options->sumExposure += options->exposures->data.F32[i];

        // Get list of PSFs, to determine target PSF
        if (options->convolve) {
            pmChip *chip = pmFPAviewThisChip(view, file->fpa); // The chip: holds the PSF
            pmPSF *psf = psMetadataLookupPtr(NULL, chip->analysis, "PSPHOT.PSF"); // PSF
            if (!psf) {
                psError(PPSTACK_ERR_PROG, false, "Unable to find PSF.");
                psFree(view);
                psFree(psfs);
                return false;
            }
            psfs->data[i] = psMemIncrRefCounter(psf);
            psMetadataRemoveKey(chip->analysis, "PSPHOT.PSF");

            pmCell *cell = pmFPAviewThisCell(view, file->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(PPSTACK_ERR_PROG, false, "Unable to determine size of image from PSF.");
                psFree(view);
                psFree(psfs);
                return false;
            }
            if (numCols == 0 && numRows == 0) {
                numCols = naxis1;
                numRows = naxis2;
            }
        }


        bool redoPhot = psMetadataLookupBool(NULL, recipe, "PHOT");

        pmDetections *detections = NULL;
        if (options->convolve || options->matchZPs || options->photometry || redoPhot) {
            pmReadout *ro = pmFPAviewThisReadout(view, file->fpa); // Readout with sources
            detections = psMetadataLookupPtr(NULL, ro->analysis, "PSPHOT.DETECTIONS"); // Sources
            if (!detections || !detections->allSources) {
                psWarning("No detections found for image %d --- rejecting.", i);
                options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] = PPSTACK_MASK_CAL;
                continue;
            }
            psAssert (detections->allSources, "missing sources?");

            options->sourceLists->data[i] = psMemIncrRefCounter(detections->allSources);
        }

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

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

            if (!ppStackInputPhotometer(ro, detections->allSources, config)) {
                psError(psErrorCodeLast(), 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);
        }
    }

    psString log = psStringCopy("Input seeing FWHMs:\n"); // Log message
    bool havePSFs = false;                                // Do we have any PSFs?
    options->inputSeeing = psVectorAlloc(num, PS_TYPE_F32);
    psVectorInit(options->inputSeeing, NAN);
    for (int i = 0; i < num; i++) {
        pmPSF *psf = psfs->data[i];     // PSF for image
        if (!psf) {
            continue;
        }
        havePSFs = true;

        int xNum = PS_MAX(psf->trendNx, 1), yNum = PS_MAX(psf->trendNy, 1); // Number of realisations
        double sumFWHM = 0.0;           // FWHM for image
        int numFWHM = 0;                // Number of FWHM measurements
        for (int y = 0; y < yNum; y++) {
            float yPos = ((float)y + 0.5) / (float)yNum * numRows;
            for (int x = 0; x < xNum; x++) {
                float xPos = ((float)x + 0.5) / (float)xNum * numCols;
                float fwhm = pmPSFtoFWHM(psf, xPos, yPos); // FWHM for image
                if (isfinite(fwhm)) {
                    sumFWHM += fwhm;
                    numFWHM++;
                }
            }
        }
        if (numFWHM == 0) {
            options->inputSeeing->data.F32[i] = NAN;
            options->inputMask->data.PS_TYPE_VECTOR_MASK_DATA[i] = PPSTACK_MASK_PSF;
            psLogMsg("ppStack", PS_LOG_INFO, "Unable to measure PSF FWHM for image %d --- rejected.", i);
        } else {
            options->inputSeeing->data.F32[i] = sumFWHM / (float)numFWHM;
        }

        psStringAppend(&log, "Input %d: %f\n", i, options->inputSeeing->data.F32[i]);
    }
    if (havePSFs) {
        psLogMsg("ppStack", PS_LOG_INFO, "%s", log);
    }
    psFree(log);

    // Generate target PSF
    if (options->convolve) {
        options->psf = ppStackPSF(config, numCols, numRows, psfs, options->inputMask);
        psFree(psfs);
        if (!options->psf) {
            psError(psErrorCodeLast(), 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);
        options->targetSeeing = pmPSFtoFWHM(options->psf, 0.5 * numCols, 0.5 * numRows); // FWHM for target
        psLogMsg("ppStack", PS_LOG_INFO, "Target seeing FWHM: %f\n", options->targetSeeing);

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

    // Zero point calibration
    if (!ppStackSourcesTransparency(options, view, config)) {
        psError(PPSTACK_ERR_DATA, false, "Unable to calculate transparency differences");
        psFree(view);
        return false;
    }
    psFree(view);

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

    return true;
}
