/** @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.25 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2007-06-20 02:22:26 $
 *
 *  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 "pmModel.h"
#include "pmSource.h"
#include "pmGrowthCurve.h"
#include "pmPSF.h"
#include "pmModelGroup.h"
#include "pmSourcePhotometry.h"
#include "pmFPAMaskWeight.h"
#include "psVectorBracket.h"
#include "pmErrorCodes.h"

/*****************************************************************************/
/* DEFINE STATEMENTS                                                         */
/*****************************************************************************/

/*****************************************************************************/
/* TYPE DEFINITIONS                                                          */
/*****************************************************************************/

/*****************************************************************************/
/* GLOBAL VARIABLES                                                          */
/*****************************************************************************/

/*****************************************************************************/
/* FILE STATIC VARIABLES                                                     */
/*****************************************************************************/

/*****************************************************************************/
/* FUNCTION IMPLEMENTATION - LOCAL                                           */
/*****************************************************************************/

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

/*****************************************************************************
pmPSFFree(psf): function to free a pmPSF structure
 *****************************************************************************/
static void pmPSFFree (pmPSF *psf)
{

    if (psf == NULL)
        return;

    psFree (psf->ChiTrend);
    psFree (psf->ApTrend);
    psFree (psf->growth);
    psFree (psf->params_NEW);
    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 (pmModelType type, bool poissonErrors, psPolynomial2D *psfTrendMask)
{
    int Nparams;

    pmPSF *psf = (pmPSF *) psAlloc(sizeof(pmPSF));

    psf->type     = type;
    psf->chisq    = 0.0;
    psf->ApResid  = 0.0;
    psf->dApResid = 0.0;
    psf->skyBias  = 0.0;
    psf->skySat   = 0.0;
    psf->poissonErrors = poissonErrors;

    // the ApTrend components are (x, y, r2rflux, flux)
    // the ApTrend order is defined by the masking
    // initially we set it to SKYBIAS (no X,Y variation)
    // this is adjusted in psphotApResid based on the user's choices
    psf->ApTrend = psPolynomial4DAlloc (PS_POLYNOMIAL_ORD, 2, 2, 1, 1);
    pmPSFMaskApTrend (psf->ApTrend, PM_PSF_APTREND_SKYBIAS);

    if (psf->poissonErrors) {
        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 = pmModelParameterCount (type);
    if (!Nparams) {
        psError(PS_ERR_UNKNOWN, true, "Undefined pmModelType");
        return(NULL);
    }
    psf->params_NEW = psArrayAlloc(Nparams);

    // the order of the PSF parameter (X,Y) fits is determined by the psfTrendMask polynomial
    // (user-specified as in the recipe). the masks of psfTrendMask are applied to each
    // parameter.  if psfTrendMask is NULL, these polynomials are not allocated.  in this case,
    // the user must set them by hand (as in pmPSFfromMD).  the parameters which are not fitted
    // are left as NULL.  these are selected by testing for them by the named values (
    // PM_PAR_XPOS, etc)

    if (psfTrendMask) {
        for (int i = 0; i < psf->params_NEW->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;

            psPolynomial2D *param = psPolynomial2DAlloc(PS_POLYNOMIAL_ORD, psfTrendMask->nX, psfTrendMask->nY);
            for (int nx = 0; nx < param->nX + 1; nx++) {
                for (int ny = 0; ny < param->nY + 1; ny++) {
                    param->mask[nx][ny] = psfTrendMask->mask[nx][ny];
                }
            }
            psf->params_NEW->data[i] = param;
        }
    }

    psMemSetDeallocator(psf, (psFreeFunc) pmPSFFree);
    return(psf);
}

/*****************************************************************************
pmModelFromPSF (*modelEXT, *psf):  use the model position parameters to
construct a realization of the PSF model at the object coordinates
 *****************************************************************************/
pmModel *pmModelFromPSF (pmModel *modelEXT, pmPSF *psf)
{

    // need to define the relationship between the modelEXT and modelPSF ?

    // find function used to set the model parameters
    pmModelFromPSFFunc modelFromPSFFunc = pmModelFromPSFFunc_GetFunction (psf->type);

    // allocate a new pmModel to hold the PSF version
    pmModel *modelPSF = pmModelAlloc (psf->type);

    // set model parameters for this source based on PSF information
    if (!modelFromPSFFunc (modelPSF, modelEXT, psf)) {
        psError(PM_ERR_PSF, false, "Failed to set model params from PSF");
        psFree(modelPSF);
        return NULL;
    }
    // XXX note that model->residuals is just a reference
    modelPSF->residuals = psf->residuals;

    return (modelPSF);
}

// zero and mask out all terms:
static bool maskAllApTerms (psPolynomial4D *trend)
{

    for (int i = 0; i < trend->nX + 1; i++) {
        for (int j = 0; j < trend->nY + 1; j++) {
            for (int k = 0; k < trend->nZ + 1; k++) {
                for (int m = 0; m < trend->nT + 1; m++) {
                    trend->mask[i][j][k][m] = 1;  // mask coeff
                    trend->coeff[i][j][k][m] = 0;  // zero coeff
                }
            }
        }
    }
    return true;
}

// 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)
{

    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)
{

    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)
{
    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)
{
    psEllipseShape shape;

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

    psEllipsePol 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;

    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)
{
    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);

    pmModelType type = pmModelSetType (typeName);
    psPolynomial2D *psfTrend = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, 0, 0);
    pmPSF *psf = pmPSFAlloc (type, true, psfTrend);

    psVector *par = psVectorAlloc (psf->params_NEW->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);

    // set the psf shape parameters
    psPolynomial2D *poly = NULL;
    poly = psf->params_NEW->data[PM_PAR_E0];
    poly->coeff[0][0] = pol.e0;

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

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

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

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

/***********************************************
 * this function masks the psf.ApTrend polynomial
 * to enable the specific subset of the coefficients
 **********************************************/
bool pmPSFMaskApTrend (psPolynomial4D *trend,
                       pmPSFApTrendOptions option)
{

    switch (option) {
    case PM_PSF_APTREND_NONE:
        maskAllApTerms (trend);
        return true;

    case PM_PSF_APTREND_CONSTANT:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        return true;

    case PM_PSF_APTREND_SKYBIAS:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[0][0][1][0] = 0;  // unmask skybias
        return true;

    case PM_PSF_APTREND_SKYSAT:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[0][0][1][0] = 0;  // unmask skybias
        trend->mask[0][0][0][1] = 0;  // unmask skysat
        return true;

    case PM_PSF_APTREND_XY_LIN:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[1][0][0][0] = 0;  // unmask x
        trend->mask[0][1][0][0] = 0;  // unmask y
        return true;

    case PM_PSF_APTREND_XY_QUAD:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[1][0][0][0] = 0;  // unmask x
        trend->mask[2][0][0][0] = 0;  // unmask x^2
        trend->mask[1][1][0][0] = 0;  // unmask x y
        trend->mask[0][1][0][0] = 0;  // unmask y
        trend->mask[0][2][0][0] = 0;  // unmask y^2
        return true;

    case PM_PSF_APTREND_SKY_XY_LIN:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[1][0][0][0] = 0;  // unmask x
        trend->mask[0][1][0][0] = 0;  // unmask y
        trend->mask[0][0][1][0] = 0;  // unmask skybias
        return true;

    case PM_PSF_APTREND_SKYSAT_XY_LIN:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[1][0][0][0] = 0;  // unmask x
        trend->mask[0][1][0][0] = 0;  // unmask y
        trend->mask[0][0][1][0] = 0;  // unmask skybias
        trend->mask[0][0][0][1] = 0;  // unmask skysat
        return true;

    case PM_PSF_APTREND_ALL:
    default:
        maskAllApTerms (trend);
        trend->mask[0][0][0][0] = 0;  // unmask constant
        trend->mask[0][0][1][0] = 0;  // unmask skybias
        trend->mask[0][0][0][1] = 0;  // unmask skysat

        trend->mask[1][0][0][0] = 0;  // unmask x
        trend->mask[2][0][0][0] = 0;  // unmask x^2
        trend->mask[1][1][0][0] = 0;  // unmask x y
        trend->mask[0][1][0][0] = 0;  // unmask y
        trend->mask[0][2][0][0] = 0;  // unmask y^2
        return true;
    }
    return false;
}

pmPSFApTrendOptions pmPSFApTrendOptionFromName (char *name)
{

    // no correction
    if (!strcasecmp (name, "NONE"))
        return PM_PSF_APTREND_NONE;
    if (!strcasecmp (name, "CONSTANT"))
        return PM_PSF_APTREND_CONSTANT;
    if (!strcasecmp (name, "SKYBIAS"))
        return PM_PSF_APTREND_SKYBIAS;
    if (!strcasecmp (name, "SKYSAT"))
        return PM_PSF_APTREND_SKYSAT;
    if (!strcasecmp (name, "XY_LIN"))
        return PM_PSF_APTREND_XY_LIN;
    if (!strcasecmp (name, "XY_QUAD"))
        return PM_PSF_APTREND_XY_QUAD;
    if (!strcasecmp (name, "SKY_XY_LIN"))
        return PM_PSF_APTREND_SKY_XY_LIN;
    if (!strcasecmp (name, "SKYSAT_XY_LIN"))
        return PM_PSF_APTREND_SKYSAT_XY_LIN;
    if (!strcasecmp (name, "ALL"))
        return PM_PSF_APTREND_ALL;

    psError(PS_ERR_UNKNOWN, true, "Unknown APTREND value: %s", name);
    return PM_PSF_APTREND_ERROR;
}

// 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)
{

    // 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);
            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;
}
