/** @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.4.4.5 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-04-14 21:43:59 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

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

#include <pslib.h>
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmPeaks.h"
#include "pmMoments.h"
#include "pmModel.h"
#include "pmSource.h"
#include "pmGrowthCurve.h"
#include "pmPSF.h"
#include "pmModelGroup.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->ApTrend);
    psFree (psf->growth);
    psFree (psf->params);
    return;
}



/*****************************************************************************
pmPSFAlloc (type): allocate a pmPSF.
    NOTE: a PSF always has 4 parameters fewer than the equivalent model.
      This is because those 4 parameters are
 X-center
 Y-center
 Sky background value
 Object Normalization
 *****************************************************************************/
pmPSF *pmPSFAlloc (pmModelType type, bool poissonErrors)
{
    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)
    psf->ApTrend = psPolynomial4DAlloc (PS_POLYNOMIAL_ORD, 2, 2, 1, 1);
    pmPSF_MaskApTrend (psf, PM_PSF_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;

    Nparams = pmModelParameterCount (type);
    if (!Nparams) {
        psError(PS_ERR_UNKNOWN, true, "Undefined pmModelType");
        return(NULL);
    }

    psf->params = psArrayAlloc(Nparams - 4);
    for (int i = 0; i < psf->params->n; i++) {
        // XXX EAM : make this a user-defined value?
        psf->params->data[i] = psPolynomial2DAlloc(PS_POLYNOMIAL_ORD, 1, 1);
    }

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



/*****************************************************************************
pmPSFFromModels (*psf, *models, *mask): build a PSF from a collection of
free-fitted models.  The PSF ignores the first 4 (independent) model
parameters and constructs a polynomial fit to the remaining as a function of
image coordinate.
    input: an array of pmModels, pre-allocated psf
Note: some of the array entries may be NULL (failed fits); ignore them.
 *****************************************************************************/
bool pmPSFFromModels (pmPSF *psf, psArray *models, psVector *mask)
{

    // construct the fit vectors from the collection of objects
    // use the mask to ignore missing fits
    psVector *x  = psVectorAlloc (models->n, PS_TYPE_F64);
    psVector *y  = psVectorAlloc (models->n, PS_TYPE_F64);
    psVector *z  = psVectorAlloc (models->n, PS_TYPE_F64);
    psVector *dz = psVectorAlloc (models->n, PS_TYPE_F64);

    for (int i = 0; i < models->n; i++) {
        pmModel *model = models->data[i];
        if (model == NULL)
            continue;

        // XXX EAM : this is fragile: x and y must be stored consistently at 2,3
        // note that the data is provided in the F64 array
        x->data.F64[i] = model->params->data.F32[2];
        y->data.F64[i] = model->params->data.F32[3];
    }

    // we are doing a robust fit.  after each pass, we drop points which are
    // more deviant than three sigma.
    // mask is currently updated for each pass. this is inconsistent?

    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);

    for (int i = 0; i < psf->params->n; i++) {
        for (int j = 0; j < models->n; j++) {
            pmModel *model = models->data[j];
            if (model == NULL)
                continue;
            z->data.F64[j] = model->params->data.F32[i + 4];
            dz->data.F64[j] = 1;
            // XXX EAM : need to use actual errors?
            // XXX EAM : this is fragile: psf(Nparams) = model(Nparams) - 4
        }

        // XXX EAM : this is the expected API (cycle 7? cycle 8?)
        psf->params->data[i] = psVectorClipFitPolynomial2D(psf->params->data[i], stats, mask, 0xff, z, dz, x, y);

        // XXX EAM : drop this when above is implemented...
        // psf->params->data[i] = RobustFit2D (psf->params->data[i], mask, x, y, z, dz);

        // psTrace ("psphot.psftest", 3, "keeping %d of %d PSF candidates (PSF param %d)\n", Nkeep, mask->n, i);
        // psPolynomial2DDump (psf->params->data[i]);
    }

    psFree (stats);
    psFree (x);
    psFree (y);
    psFree (z);
    psFree (dz);
    return (true);
}



/*****************************************************************************
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
    modelFromPSFFunc (modelPSF, modelEXT, psf);

    return (modelPSF);
}

// zero and mask out all terms:
static bool maskAllTerms (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;
}

/***********************************************
 * this function masks the psf.ApTrend polynomial 
 * to enable the specific subset of the coefficients
 **********************************************/
bool pmPSF_MaskApTrend (pmPSF *psf, pmPSF_ApTrendOptions option)
{

    switch (option) {
    case PM_PSF_NONE:
        maskAllTerms (psf->ApTrend);
        return true;

    case PM_PSF_CONSTANT:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        return true;

    case PM_PSF_SKYBIAS:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[0][0][1][0] = 0;  // unmask skybias
        return true;

    case PM_PSF_SKYSAT:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[0][0][1][0] = 0;  // unmask skybias
        psf->ApTrend->mask[0][0][0][1] = 0;  // unmask skybias
        return true;

    case PM_PSF_XY_LIN:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[1][0][0][0] = 0;  // unmask x
        psf->ApTrend->mask[0][1][0][0] = 0;  // unmask y
        return true;

    case PM_PSF_XY_QUAD:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[1][0][0][0] = 0;  // unmask x
        psf->ApTrend->mask[2][0][0][0] = 0;  // unmask x^2
        psf->ApTrend->mask[1][1][0][0] = 0;  // unmask x y
        psf->ApTrend->mask[0][1][0][0] = 0;  // unmask y
        psf->ApTrend->mask[0][2][0][0] = 0;  // unmask y^2
        return true;

    case PM_PSF_SKY_XY_LIN:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[1][0][0][0] = 0;  // unmask x
        psf->ApTrend->mask[0][1][0][0] = 0;  // unmask y
        psf->ApTrend->mask[0][0][1][0] = 0;  // unmask skybias
        return true;

    case PM_PSF_SKYSAT_XY_LIN:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[1][0][0][0] = 0;  // unmask x
        psf->ApTrend->mask[0][1][0][0] = 0;  // unmask y
        psf->ApTrend->mask[0][0][1][0] = 0;  // unmask skybias
        psf->ApTrend->mask[0][0][0][1] = 0;  // unmask skysat
        return true;

    case PM_PSF_ALL:
    default:
        maskAllTerms (psf->ApTrend);
        psf->ApTrend->mask[0][0][0][0] = 0;  // unmask constant
        psf->ApTrend->mask[0][0][1][0] = 0;  // unmask skybias
        psf->ApTrend->mask[0][0][0][1] = 0;  // unmask skysat

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