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


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

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmReadoutFake.h"

#include "pmMoments.h"
#include "pmResiduals.h"
#include "pmGrowthCurve.h"
#include "pmTrend2D.h"
#include "pmPSF.h"
#include "pmModel.h"
#include "pmModelClass.h"
#include "pmModelUtils.h"
#include "pmPeaks.h"
#include "pmSource.h"
#include "pmSourceUtils.h"
#include "pmSourceFitModel.h"
#include "pmPSFtry.h"


#include "pmPSFEnvelope.h"



#define TESTING                         // Enable test output
#define PEAK_FLUX 1.0e3                 // Peak flux for each source
#define SKY_VALUE 0.0e0                 // Sky value for fake image
#define WEIGHT_VAL 1.0e0                // Weighting for image
#define PSF_STATS PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV // Statistics options for measuring PSF



// XXX To do:
//
// * PSF variation when only a portion of the image is present (e.g., the edge of an FPA overlapping a
// skycell) may mean a disastrously weird PSF in the missing regions.  To counter this, get a region of
// validity for each PSF (perhaps from an associated mask, or have the user work it out), and taper the PSF
// when outside this region (perhaps multiply the peak flux by a Gaussian whose arguments are the distance
// from the valid region, and a width that the user supplies).


// We deliberately do not include the calculation of and storing of residuals (data - model) for the PSF
// model, because (1) there is no code in psModules to do this, and we're not going to implement it here; and
// (2) this is intended to generate "nice" or "ideal" PSFs to feed into pmSubtraction (PSF matching code), so
// any residuals will hopefully be dealt with by that.


pmPSF *pmPSFEnvelope(int numCols, int numRows, // Size of original image
                     const psArray *inputs, // Input PSF models
                     int instances, // Number of instances per dimension
                     int radius,        // Radius of each PSF
                     const char *modelName,// Name of PSF model to use
                     int xOrder, int yOrder // Order for PSF variation fit
                     )
{
    PS_ASSERT_INT_POSITIVE(numCols, NULL);
    PS_ASSERT_INT_POSITIVE(numRows, NULL);
    PS_ASSERT_ARRAY_NON_NULL(inputs, NULL);
    PS_ASSERT_INT_POSITIVE(instances, NULL);
    PS_ASSERT_INT_POSITIVE(radius, NULL);
    PS_ASSERT_STRING_NON_EMPTY(modelName, NULL);
    PS_ASSERT_INT_NONNEGATIVE(xOrder, NULL);
    PS_ASSERT_INT_NONNEGATIVE(yOrder, NULL);

    float xOrigSpacing = (float)(numCols - 2 * radius) / (float)(instances - 1); // Spacing between instances
    float yOrigSpacing = (float)(numRows - 2 * radius) / (float)(instances - 1); // Spacing between instances
    int fakeSpacing = 2 * radius + 1;   // Spacing between instances (x and y) in the fake image
    int fakeSize = instances * fakeSpacing; // Size of fake image

    // Generate list of fake sources (instances of the PSF)
    int numFakes = PS_SQR(instances);   // Number of fake sources
    psArray *fakes = psArrayAlloc(numFakes); // Fake sources
    psVector *xOffset = psVectorAlloc(numFakes, PS_TYPE_S32); // X offset from fake position to image
    psVector *yOffset = psVectorAlloc(numFakes, PS_TYPE_S32); // Y offset from fake position to image
    for (int j = 0, index = 0; j < instances; j++) {
        float yOrig = j * yOrigSpacing + radius; // Source position in original image
        float yFake = j * fakeSpacing + radius; // Position in fake image
        int dy = yFake - yOrig;         // Difference between fake and original position

        for (int i = 0; i < instances; i++, index++) {
            float xOrig = i * xOrigSpacing + radius; // Source position in original image
            float xFake = i * fakeSpacing + radius; // Position in fake image
            int dx = xFake - xOrig;     // Difference between fake and original position

            pmSource *fake = pmSourceAlloc(); // Fake source
            fake->peak = pmPeakAlloc(xFake - dx, yFake - dy, PEAK_FLUX, PM_PEAK_LONE);
            fake->type = PM_SOURCE_TYPE_STAR;

#ifdef TESTING
            // Required to generate model image
            fake->psfMag = -2.5 * log10(PEAK_FLUX);
#endif

            psTrace("psModules.imcombine", 5, "Source %d: %.2f,%.2f\n",
                    index, xOrig, yOrig);

            fakes->data[index] = fake;
            xOffset->data.S32[index] = dx;
            yOffset->data.S32[index] = dy;
        }
    }

    // Generate fake images with each PSF, and take the envelope
    psImage *envelope = psImageAlloc(fakeSize, fakeSize, PS_TYPE_F32); // Image with envelope of PSFs
    psImageInit(envelope, SKY_VALUE);
    pmReadout *fakeRO = pmReadoutAlloc(NULL); // Fake readout
    for (int i = 0; i < inputs->n; i++) {
        pmPSF *psf = inputs->data[i];   // PSF of interest
        pmResiduals *resid = psf->residuals;// PSF residuals
        psf->residuals = NULL;
        if (!pmReadoutFakeFromSources(fakeRO, fakeSize, fakeSize, fakes, xOffset, yOffset, psf,
                                      NAN, radius, true)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate fake readout.");
            psFree(envelope);
            psFree(yOffset);
            psFree(xOffset);
            psFree(fakes);
            psf->residuals = resid;
            return NULL;
        }
        psf->residuals = resid;

        // Need to renormalise sources so they all have the same peak.  You would think they do have the same
        // peak already, but it seems that the residual map messes things up by adding extra flux
        for (int j = 0; j < numFakes; j++) {
            pmSource *source = fakes->data[j]; // Fake source
            float x = source->peak->xf + xOffset->data.S32[j]; // x coordinate of source
            float y = source->peak->yf + yOffset->data.S32[j]; // y coordinate of source

            double flux = fakeRO->image->data.F32[(int)y][(int)x];
            float norm = PEAK_FLUX / flux; // Normalisation for source
            psRegion region = psRegionSet(x - radius, x + radius, y - radius, y + radius); // PSF region
            psImage *subImage = psImageSubset(fakeRO->image, region); // Subimage of fake PSF
            psImage *subEnv = psImageSubset(envelope, region); // Subimage of envelope
            psBinaryOp(subImage, subImage, "*", psScalarAlloc(norm, PS_TYPE_F32));
            psBinaryOp(subEnv, subEnv, "MAX", subImage);
            psFree(subImage);
            psFree(subEnv);
        }

#ifdef TESTING
        {
            // Write out the PSF field
            psString name = NULL;
            psStringAppend(&name, "psf_field_%03d.fits", i);
            psFits *fits = psFitsOpen(name, "w");
            psFitsWriteImage(fits, NULL, fakeRO->image, 0, NULL);
            psFitsClose(fits);
            psFree(name);
        }
#endif

    }
    psFree(fakeRO);

#ifdef TESTING
    {
        // Write out the envelope
        psFits *fits = psFitsOpen("psf_field_envelope.fits", "w");
        psFitsWriteImage(fits, NULL, envelope, 0, NULL);
        psFitsClose(fits);
    }
#endif

    // Put the fake sources onto a full-size image
    pmReadout *readout = pmReadoutAlloc(NULL); // Readout to contain envelope pixels
    readout->image = psImageAlloc(numCols, numRows, PS_TYPE_F32);
    psImageInit(readout->image, 0.0);
    for (int i = 0; i < numFakes; i++) {
        pmSource *source = fakes->data[i]; // Fake source
        // Position of source on fake image
        int xFake = source->peak->x + xOffset->data.S32[i];
        int yFake = source->peak->y + yOffset->data.S32[i];
        psRegion region = psRegionSet(xFake - radius, xFake + radius,
                                      yFake - radius, yFake + radius); // PSF region
        psImage *subImage = psImageSubset(envelope, region); // Subimage of fake PSF

        // Position of source on "real" image
        int x0 = source->peak->x - radius;
        int y0 = source->peak->y - radius;

        if (!psImageOverlaySection(readout->image, subImage, x0, y0, "=")) {
            psError(PS_ERR_UNKNOWN, false, "Unable to overlay PSF");
            psFree(readout);
            psFree(xOffset);
            psFree(yOffset);
            psFree(fakes);
            return NULL;
        }
    }
    psFree(xOffset);
    psFree(yOffset);

    readout->weight = (psImage*)psBinaryOp(NULL, readout->image, "+", psScalarAlloc(1.0, PS_TYPE_F32));
    readout->mask = psImageAlloc(numCols, numRows, PS_TYPE_MASK);
    psImageInit(readout->mask, 0);

#ifdef TESTING
    {
        // Write out the envelope
        psFits *fits = psFitsOpen("psf_field_full.fits", "w");
        psFitsWriteImage(fits, NULL, readout->image, 0, NULL);
        psFitsClose(fits);
    }
#endif

    // Reset the sources to point to the new pixels, and measure the moments in preparation for PSF fitting
    for (int i = 0; i < numFakes; i++) {
        pmSource *source = fakes->data[i]; // Fake source
        float x = source->peak->xf;     // x coordinates of source
        float y = source->peak->yf;     // y coordinates of source

        psFree(source->pixels);
        psFree(source->weight);
        psFree(source->maskView);
        psFree(source->maskObj);
        source->pixels = NULL;
        source->weight = NULL;
        source->maskView = NULL;
        source->maskObj = NULL;

        if (!pmSourceDefinePixels(source, readout, x, y, (float)radius)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to define pixels for source.");
            psFree(readout);
            psFree(fakes);
            return NULL;
        }

        if (!pmSourceMoments(source, radius)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to measure moments for source.");
            psFree(readout);
            psFree(fakes);
            return NULL;
        }
    }

    // Don't assume Poisson errors
    pmPSFOptions *options = pmPSFOptionsAlloc(); // Options for fitting a PSF
    options->poissonErrorsPhotLMM = true;
    options->poissonErrorsPhotLin = false;
    options->poissonErrorsParams = false;
    options->stats = psStatsAlloc(PSF_STATS);
    options->radius = radius;
    options->psfTrendMode = PM_TREND_MAP;
    options->psfTrendNx = xOrder;
    options->psfTrendNy = yOrder;
    options->psfFieldNx = numCols;
    options->psfFieldNy = numRows;
    options->psfFieldXo = 0;
    options->psfFieldYo = 0;

    pmSourceFitModelInit (15, 0.01, WEIGHT_VAL, options->poissonErrorsPhotLMM);

    pmPSFtry *try = pmPSFtryModel(fakes, modelName, options, 0, 0xff);
    psFree(options);
    if (!try) {
        psError(PS_ERR_UNKNOWN, false, "Unable to fit PSF model to PSF envelope.");
        psFree(readout);
        psFree(fakes);
        return NULL;
    }

    pmPSF *psf = psMemIncrRefCounter(try->psf); // Output PSF
    psFree(try);

#ifdef TESTING
    {
        // Need to translate peak flux --> integrated flux
        pmModel *fakeModel = pmModelFromPSFforXY(psf, (float)numCols / 2.0, (float)numRows / 2.0,
                                                 1.0); // Fake model, with central intensity of 1.0
        float flux0 = fakeModel->modelFlux(fakeModel->params); // Flux for central intensity of 1.0
        for (int i = 0; i < numFakes; i++) {
            pmSource *source = fakes->data[i]; // Fake source
            source->psfMag -= 2.5 * log10(flux0);
        }

        pmReadout *generated = pmReadoutAlloc(NULL); // Generated image
        pmReadoutFakeFromSources(generated, numCols, numRows, fakes, NULL, NULL, psf, NAN, radius,
                                 false);
        {
            psFits *fits = psFitsOpen("psf_field_model.fits", "w");
            psFitsWriteImage(fits, NULL, generated->image, 0, NULL);
            psFitsClose(fits);
        }
        psBinaryOp(generated->image, generated->image, "-", readout->image);
        {
            psFits *fits = psFitsOpen("psf_field_resid.fits", "w");
            psFitsWriteImage(fits, NULL, generated->image, 0, NULL);
            psFitsClose(fits);
        }
        psFree(generated);
    }
#endif

    psFree(readout);

    return psf;
}
