/** @file  pmPSF.c
 *
 * This file contains typedefs for the Point-Spread Function and prototypes
 * for functions that calculate the PSF.
 *
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.35 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2008-01-02 20:41:55 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

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

/*****************************************************************************/
/* INCLUDE FILES                                                             */
/*****************************************************************************/

#include <strings.h>  // for strcasecmp
#include <pslib.h>
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmPeaks.h"
#include "pmMoments.h"
#include "pmResiduals.h"
#include "pmGrowthCurve.h"
#include "pmTrend2D.h"
#include "pmPSF.h"
#include "pmModel.h"
#include "pmSource.h"
#include "pmModelClass.h"
#include "pmModelUtils.h"
#include "pmSourcePhotometry.h"
#include "pmFPAMaskWeight.h"
#include "psVectorBracket.h"
#include "pmErrorCodes.h"

/*****************************************************************************/
/* FUNCTION IMPLEMENTATION - PUBLIC                                          */
/*****************************************************************************/

static void pmPSFOptionsFree (pmPSFOptions *options) {

    if (!options) return;

    psFree (options->stats);
    return;
}

pmPSFOptions *pmPSFOptionsAlloc () {

    pmPSFOptions *options = (pmPSFOptions *) psAlloc(sizeof(pmPSFOptions));
    psMemSetDeallocator(options, (psFreeFunc) pmPSFOptionsFree);

    options->type          = 0;

    options->stats         = NULL;

    options->psfTrendMode  = PM_TREND_NONE;
    options->psfTrendNx    = 0;
    options->psfTrendNy    = 0;
    options->psfFieldNx    = 0;
    options->psfFieldNy    = 0;
    options->psfFieldXo    = 0;
    options->psfFieldYo    = 0;

    options->poissonErrorsPhotLMM = true;
    options->poissonErrorsPhotLin = false;
    options->poissonErrorsParams  = true;

    return options;
}

bool psMemCheckPSFOptions(psPtr ptr)
{
    PS_ASSERT_PTR(ptr, false);
    return ( psMemGetDeallocator(ptr) == (psFreeFunc) pmPSFOptionsFree);
}

/*****************************************************************************
pmPSFFree(psf): function to free a pmPSF structure
 *****************************************************************************/
static void pmPSFFree (pmPSF *psf)
{
    if (psf == NULL) {
        return;
    }

    psFree (psf->ChiTrend);
    psFree (psf->psfTrendStats);
    psFree (psf->ApTrend);
    psFree (psf->FluxScale);
    psFree (psf->growth);
    psFree (psf->params);
    psFree (psf->residuals);
    return;
}

/*****************************************************************************
 pmPSFAlloc (type): allocate a pmPSF.

 NOTE: PSF model parameters which are not modeled on an image are set to NULL
 in psf->params.

 These are normally:

 X-center
 Y-center
 Sky background value
 Object Normalization
 *****************************************************************************/
pmPSF *pmPSFAlloc (const pmPSFOptions *options)
{
    PS_ASSERT_PTR_NON_NULL(options, NULL);
    int Nparams;

    pmPSF *psf = (pmPSF *) psAlloc(sizeof(pmPSF));
    psMemSetDeallocator(psf, (psFreeFunc) pmPSFFree);

    psf->type     = options->type;
    psf->chisq    = 0.0;
    psf->ApResid  = 0.0;
    psf->dApResid = 0.0;
    psf->skyBias  = 0.0;
    psf->skySat   = 0.0;
    psf->nPSFstars  = 0;
    psf->nApResid   = 0;

    psf->poissonErrorsPhotLMM = options->poissonErrorsPhotLMM;
    psf->poissonErrorsPhotLin = options->poissonErrorsPhotLin;
    psf->poissonErrorsParams = options->poissonErrorsParams;

    // the ApTrend components are (x, y).  It may be represented with a polynomial or with a
    // psImageMap.  We set it initially to NULL. the user must allocate it before using.
    // psf->ApTrend = pmTrend2DAlloc (PM_TREND_MAP, Nx, Ny, 1, 1, stats);
    psf->ApTrend = NULL;

    // the flux scale is the relationship between the integrated flux and the peak flux for a
    // psf model.  this is a 2D function of position, and is modeled with pmTrend2D after the
    // psf model is determined for an image.  until it is determined, the flux calculation
    // integrates the sources
    psf->FluxScale = NULL;

    if (psf->poissonErrorsPhotLMM) {
        psf->ChiTrend = psPolynomial1DAlloc (PS_POLYNOMIAL_ORD, 1);
    } else {
        psf->ChiTrend = psPolynomial1DAlloc (PS_POLYNOMIAL_ORD, 2);
    }

    // don't define a growth curve : user needs to choose radius bins
    psf->growth = NULL;

    // by default, we do not construct the residual image
    psf->residuals = NULL;

    Nparams = pmModelClassParameterCount (options->type);
    if (!Nparams) {
        psError(PS_ERR_UNKNOWN, true, "Undefined pmModelType");
        return(NULL);
    }
    psf->params = psArrayAlloc(Nparams);

    // save the trend stats on the psf for use in pmPSFFromPSFtry
    psf->psfTrendStats = psMemIncrRefCounter (options->stats);

    // the psf parameters may have 2D variations represented as either a polynomial (ordinary
    // or chebychev) or as an image map.  The size of the image map is determined by pmPSFtry
    // by minimizing the scatter in the fitted data and the rms error in the complete image
    // map.  In this case, the user-supplied options of psfTrendNx and psfTrendNy are the
    // maximum value used for these axes.  For the polynomial terms, the order is not currently
    // set, dynamically.  The value of psfTrendNx and psfTrendNy are used.

    psImageBinning *binning = psImageBinningAlloc();
    binning->nXruff = options->psfTrendNx;
    binning->nYruff = options->psfTrendNy;
    binning->nXfine = options->psfFieldNx;
    binning->nYfine = options->psfFieldNy;

    // for polynomial representations, nXruff, nYruff are the order number, and may be 0.
    // in this case, we cannot set the psImageBinning scale because 0 would be invalid.  these
    // elements are only used by the image map representation
    if (options->psfTrendMode == PM_TREND_MAP) {
        psImageBinningSetScale (binning, PS_IMAGE_BINNING_CENTER);
        psImageBinningSetSkipByOffset (binning, options->psfFieldXo, options->psfFieldYo);
    }

    // trendNx & trendNy are used in pmPSFtry as the max for these values
    psf->psfTrendMode = options->psfTrendMode;
    psf->trendNx      = options->psfTrendNx;
    psf->trendNy      = options->psfTrendNy;
    psf->fieldNx      = options->psfFieldNx;
    psf->fieldNy      = options->psfFieldNy;
    psf->fieldXo      = options->psfFieldXo;
    psf->fieldYo      = options->psfFieldYo;

    // define the parameter trends
    if (options->psfTrendMode != PM_TREND_NONE) {
        for (int i = 0; i < psf->params->n; i++) {
            if (i == PM_PAR_SKY) continue;
            if (i == PM_PAR_I0) continue;
            if (i == PM_PAR_XPOS) continue;
            if (i == PM_PAR_YPOS) continue;

            psf->params->data[i] = pmTrend2DNoImageAlloc (options->psfTrendMode, binning, options->stats);
        }
    }
    psFree (binning);
    return psf;
}

bool psMemCheckPSF(psPtr ptr)
{
    PS_ASSERT_PTR(ptr, false);
    return ( psMemGetDeallocator(ptr) == (psFreeFunc) pmPSFFree);
}

// the PSF models the \sigma_{xy} variation of the elliptical contour as a function of position in the image with a
// polynomial.  an individual object has a contour of the form (x^2/2sx^2) + (y^2/2sy^2) + sxy*x*y
// these are the values of the model->params.  the psf->params term for sxy is actually fitted
// to sxy/(sxx^-2 + syy^-2)^2

// XXX this is only an approximate solution.  A better solution would be to fit the second moment, Mxy. the
// problem here is that converting from Mxy,SXX,SYY -> SXY is a third order problem:
// Mxy = SXY * (SXX^-4 + SYY^-4 - 2 SXY ^2)

// input: model->param, output: psf->param[PM_PAR_SXY]
double pmPSF_SXYfromModel (psF32 *modelPar)
{
    PS_ASSERT_PTR_NON_NULL(modelPar, NAN);

    double SXX = modelPar[PM_PAR_SXX];
    double SYY = modelPar[PM_PAR_SYY];
    double SXY = modelPar[PM_PAR_SXY];

    double par = SXY / PS_SQR(1.0 / PS_SQR(SXX) + 1.0 / PS_SQR(SYY));
    return (par);
}

// input: fitted psf->param, output: model->param[PM_PAR_SXY]
double pmPSF_SXYtoModel (psF32 *fittedPar)
{
    PS_ASSERT_PTR_NON_NULL(fittedPar, NAN);

    double SXX = fittedPar[PM_PAR_SXX];
    double SYY = fittedPar[PM_PAR_SYY];
    double fit = fittedPar[PM_PAR_SXY];

    double SXY = fit * PS_SQR(1.0 / PS_SQR(SXX) + 1.0 / PS_SQR(SYY));

    assert (!isnan(SXY));

    return SXY;
}

// New Concept: the PSF modelling function fits the polarization terms e0, e1, e2:

// convert the parameters used in the fitted source model
// to the parameters used in the 2D PSF model
bool pmPSF_FitToModel (psF32 *fittedPar, float minMinorAxis)
{
    PS_ASSERT_PTR_NON_NULL(fittedPar, false);

    psEllipsePol pol;

    pol.e0 = fittedPar[PM_PAR_E0];
    pol.e1 = fittedPar[PM_PAR_E1];
    pol.e2 = fittedPar[PM_PAR_E2];

    psEllipseAxes axes;
    if (psEllipsePolToAxes (pol, minMinorAxis, &axes) != PS_ERR_NONE) {
        psError(PM_ERR_PSF, false, "Failed to convert e[012] (%g,%g,%g) to axes",
                pol.e0, pol.e1, pol.e2);
        return false;
    }
    psEllipseShape shape = psEllipseAxesToShape (axes);

    fittedPar[PM_PAR_SXX] = shape.sx * M_SQRT2;
    fittedPar[PM_PAR_SYY] = shape.sy * M_SQRT2;
    fittedPar[PM_PAR_SXY] = shape.sxy;

    return true;
}

// convert the PSF parameters used in the 2D PSF model fit into the
// parameters used in the source model
psEllipsePol pmPSF_ModelToFit (psF32 *modelPar)
{
    // must assert non-NULL input parameter
    psEllipsePol pol;
    pol.e0 = NAN;
    pol.e1 = NAN;
    pol.e2 = NAN;
    PS_ASSERT_PTR_NON_NULL(modelPar, pol);

    psEllipseShape shape;

    shape.sx  = modelPar[PM_PAR_SXX] / M_SQRT2;
    shape.sy  = modelPar[PM_PAR_SYY] / M_SQRT2;
    shape.sxy = modelPar[PM_PAR_SXY];

    pol = psEllipseShapeToPol (shape);

    return pol;
}

// convert the parameters used in the fitted source model to the psEllipseAxes representation
// (major,minor,theta)
psEllipseAxes pmPSF_ModelToAxes (psF32 *modelPar, double maxAR)
{
    psEllipseShape shape;
    psEllipseAxes axes;
    axes.major = NAN;
    axes.minor = NAN;
    axes.theta = NAN;
//   XXX: must assert non-NULL input parameter
    PS_ASSERT_PTR_NON_NULL(modelPar, axes);

    shape.sx  = modelPar[PM_PAR_SXX] / M_SQRT2;
    shape.sy  = modelPar[PM_PAR_SYY] / M_SQRT2;
    shape.sxy = modelPar[PM_PAR_SXY];

    if ((shape.sx == 0) || (shape.sy == 0)) {
        axes.major = 0.0;
        axes.minor = 0.0;
        axes.theta = 0.0;
    } else {
        // XXX this is not really consistent with the model fit range above
        axes = psEllipseShapeToAxes (shape, maxAR);
    }

    return axes;
}

// convert the psEllipseAxes representation (major,minor,theta) to the parameters used in the
// fitted source model
bool pmPSF_AxesToModel (psF32 *modelPar, psEllipseAxes axes)
{
    PS_ASSERT_PTR_NON_NULL(modelPar, false);

    if ((axes.major <= 0) || (axes.minor <= 0)) {
        modelPar[PM_PAR_SXX] = 0.0;
        modelPar[PM_PAR_SYY] = 0.0;
        modelPar[PM_PAR_SXY] = 0.0;
        return true;
    }

    psEllipseShape shape = psEllipseAxesToShape (axes);

    modelPar[PM_PAR_SXX] = shape.sx * M_SQRT2;
    modelPar[PM_PAR_SYY] = shape.sy * M_SQRT2;
    modelPar[PM_PAR_SXY] = shape.sxy;

    return true;
}

// generate a psf model of the requested type, with fixed shape
pmPSF *pmPSFBuildSimple (char *typeName, float sxx, float syy, float sxy, ...)
{

    va_list ap;
    va_start(ap, sxy);

    pmPSFOptions *options = pmPSFOptionsAlloc ();
    options->type = pmModelClassGetType (typeName);
    options->psfTrendMode = PM_TREND_POLY_ORD;
    options->psfTrendNx = 0;
    options->psfTrendNy = 0;

    pmPSF *psf = pmPSFAlloc (options);

    psVector *par = psVectorAlloc (psf->params->n, PS_TYPE_F32);
    par->data.F32[PM_PAR_SXX] = sxx;
    par->data.F32[PM_PAR_SYY] = syy;
    par->data.F32[PM_PAR_SXY] = sxy;

    psEllipsePol pol = pmPSF_ModelToFit(par->data.F32);

    pmTrend2D *trend = NULL;

    // set the psf shape parameters
    trend = psf->params->data[PM_PAR_E0];
    trend->poly->coeff[0][0] = pol.e0;

    trend = psf->params->data[PM_PAR_E1];
    trend->poly->coeff[0][0] = pol.e1;

    trend = psf->params->data[PM_PAR_E2];
    trend->poly->coeff[0][0] = pol.e2;

    for (int i = PM_PAR_SXY + 1; i < psf->params->n; i++) {
        trend = psf->params->data[i];
        trend->poly->coeff[0][0] = (psF32)va_arg(ap, psF64);
    }
    va_end(ap);

    psFree (par);
    psFree (options);
    return psf;
}

// we generate the growth curve for the center of the image with the specified psf model
bool pmGrowthCurveGenerate (pmReadout *readout, pmPSF *psf, bool ignore, psMaskType maskVal, psMaskType mark)
{
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(readout->image, false);

    // bool status;
    float xc, yc, dx, dy;
    float fitMag, apMag;
    float radius;

    // create template model
    pmModel *modelRef = pmModelAlloc(psf->type);

    // use the center of the center pixel of the image
    xc = 0.5*readout->image->numCols + readout->image->col0 + 0.5;
    yc = 0.5*readout->image->numRows + readout->image->row0 + 0.5;
    dx = psf->growth->maxRadius + 1;
    dy = psf->growth->maxRadius + 1;

    // assign the x and y coords to the image center
    // create an object with center intensity of 1000
    modelRef->params->data.F32[PM_PAR_SKY] = 0;
    modelRef->params->data.F32[PM_PAR_I0] = 1000;
    modelRef->params->data.F32[PM_PAR_XPOS] = xc;
    modelRef->params->data.F32[PM_PAR_YPOS] = yc;

    // create modelPSF from this model
    pmModel *model = pmModelFromPSF (modelRef, psf);

    // measure the fitMag for this model
    pmSourcePhotometryModel (&fitMag, model);
    psf->growth->fitMag = fitMag;

    // generate working image for this source
    psRegion region = {xc - dx, xc + dx, yc - dy, yc + dy};

    // force region to stop at dimensions of image
    region = psRegionForImage (readout->image, region);

    // the view, image, and mask retain col0,row0
    psImage *view = psImageSubset (readout->image, region);
    psImage *image = psImageCopy (NULL, view, PS_TYPE_F32);
    psImage *mask = psImageCopy (NULL, view, PS_TYPE_U8);

    psImageInit (image, 0.0);
    psImageInit (mask, 0);

    // place the reference object in the image center
    // no need to mask the source here
    // XXX should we measure this for the analytical model only or the full model?
    pmModelAdd (image, NULL, model, PM_MODEL_OP_FULL, maskVal);

    // loop over a range of source fluxes
    // no need to interpolate since we have forced the object center
    // to 0.5, 0.5 above
    bool completeGrowthCurve = true;            // do we have a complete curve of growth?
    for (int i = 0; i < psf->growth->radius->n; i++) {

        radius = psf->growth->radius->data.F32[i];

        // mask the given aperture and measure the apMag
        psImageKeepCircle (mask, xc, yc, radius, "OR", mark);
        if (!pmSourcePhotometryAper (&apMag, model, image, mask, maskVal)) {
            psError(PM_ERR_PHOTOM, false, "Measuring apMag for radius == %g", radius);
            psErrorStackPrint (stderr, "failure to get growth curve");
            psErrorClear();
            completeGrowthCurve = false;
            break;
        }

        // XXX since we re-mask on each pass, this could be dropped.
        psImageKeepCircle (mask, xc, yc, radius, "AND", PS_NOT_U8(mark));

        // the 'ignore' mode is for testing
        if (ignore) {
            psf->growth->apMag->data.F32[i] = fitMag;
        } else {
            psf->growth->apMag->data.F32[i] = apMag;
        }
    }

    if (completeGrowthCurve) {
        psf->growth->apRef = psVectorInterpolate (psf->growth->radius, psf->growth->apMag, psf->growth->refRadius);
        psf->growth->apLoss = psf->growth->fitMag - psf->growth->apRef;
    } else {
        psf->growth->apRef = NAN;
        psf->growth->apLoss = 0;
    }

    psLogMsg ("psphot.growth", 4, "GrowthCurve : apLoss : %f\n", psf->growth->apLoss);

    psFree (view);
    psFree (image);
    psFree (mask);
    psFree (model);
    psFree (modelRef);

    return completeGrowthCurve ? true : false;
}
