/** @file  pmPSFtry.c
 *
 *  XXX: need description of file purpose
 *
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.34.2.1 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2007-03-08 23:27:24 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

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

# include <pslib.h>
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmFPAMaskWeight.h"
#include "pmPeaks.h"
#include "pmMoments.h"
#include "pmModel.h"
#include "pmSource.h"
#include "pmGrowthCurve.h"
#include "pmPSF.h"
#include "pmPSFtry.h"
#include "pmModelGroup.h"
#include "pmSourceFitModel.h"
#include "pmSourcePhotometry.h"

// ********  pmPSFtry functions  **************************************************
// * pmPSFtry holds a single pmPSF model test, with the input sources, the freely
// * fitted version of the model, the pmPSF fit to the fitted model parameters,
// * and the PSF fits to the source. It also includes the statistics from the
// * fits, both the individual sources, and the collection

// free a pmPSFtry structure
static void pmPSFtryFree (pmPSFtry *test)
{

    if (test == NULL)
        return;

    psFree (test->psf);
    psFree (test->sources);
    psFree (test->metric);
    psFree (test->metricErr);
    psFree (test->fitMag);
    psFree (test->mask);
    return;
}

// allocate a pmPSFtry based on the desired sources and the model (identified by name)
pmPSFtry *pmPSFtryAlloc (psArray *sources, char *modelName, bool poissonErrors, psPolynomial2D *psfTrendMask)
{

    // validate the requested model name
    pmModelType type = pmModelSetType (modelName);
    if (type == -1) {
        psError (PS_ERR_UNKNOWN, true, "invalid model name %s", modelName);
        return NULL;
    }

    pmPSFtry *test = (pmPSFtry *) psAlloc(sizeof(pmPSFtry));

    test->psf       = pmPSFAlloc (type, poissonErrors, psfTrendMask);
    test->metric    = psVectorAlloc (sources->n, PS_TYPE_F64);
    test->metricErr = psVectorAlloc (sources->n, PS_TYPE_F64);
    test->fitMag    = psVectorAlloc (sources->n, PS_TYPE_F64);
    test->mask      = psVectorAlloc (sources->n, PS_TYPE_U8);

    test->sources   = psArrayAlloc (sources->n);
    for (int i = 0; i < sources->n; i++) {
        test->sources->data[i] = pmSourceCopy (sources->data[i]);
        test->mask->data.U8[i]  = 0;
        test->metric->data.F64[i] = 0;
        test->metricErr->data.F64[i] = 0;
        test->fitMag->data.F64[i] = 0;
    }

    psMemSetDeallocator(test, (psFreeFunc) pmPSFtryFree);
    return (test);
}

// build a pmPSFtry for the given model:
// - fit each source with the free-floating model
// - construct the pmPSF from the collection of models
// - fit each source with the PSF-parameter models
// - measure the pmPSF quality metric (dApResid)

// sources used in for pmPSFtry may be masked by the analysis
// mask values indicate the reason the source was rejected:

// generate a pmPSFtry with a copy of the test PSF sources
pmPSFtry *pmPSFtryModel (psArray *sources, char *modelName, float RADIUS, bool poissonErrors, psPolynomial2D *psfTrendMask, bool applyWeights)
{
    bool status;
    float x;
    float y;
    int Next = 0;
    int Npsf = 0;

    pmPSFtry *psfTry = pmPSFtryAlloc (sources, modelName, poissonErrors, psfTrendMask);
    if (psfTry == NULL) {
        psError (PS_ERR_UNKNOWN, false, "failed to allocate psf model");
        return NULL;
    }

    // stage 1:  fit an independent model (freeModel) to all sources
    psTimerStart ("fit");
    for (int i = 0; i < psfTry->sources->n; i++) {

        pmSource *source = psfTry->sources->data[i];
        source->modelEXT = pmSourceModelGuess (source, psfTry->psf->type);
        if (source->modelEXT == NULL) {
            psError(PS_ERR_UNKNOWN, false, "failed to build model");
            return NULL;
        }
        x = source->peak->x;
        y = source->peak->y;

        // set temporary object mask and fit object
        // fit model as EXT, not PSF

        psImageKeepCircle (source->mask, x, y, RADIUS, "OR", PM_MASK_MARK);
        status = pmSourceFitModel (source, source->modelEXT, PM_SOURCE_FIT_EXT);
        psImageKeepCircle (source->mask, x, y, RADIUS, "AND", PS_NOT_U8(PM_MASK_MARK));

        // exclude the poor fits
        if (!status) {
            psfTry->mask->data.U8[i] = PSFTRY_MASK_EXT_FAIL;
            continue;
        }
        Next ++;
    }
    psLogMsg ("psphot.psftry", 4, "fit ext:   %f sec for %d of %ld sources\n", psTimerMark ("fit"), Next, sources->n);
    psTrace ("psphot.psftry", 3, "keeping %d of %ld PSF candidates (EXT)\n", Next, sources->n);

    // stage 2: construct a psf (pmPSF) from this collection of model fits
    if (!pmPSFFromPSFtry (psfTry, applyWeights)) {
        psError(PS_ERR_UNKNOWN, false, "failed to construct a psf model from collection of sources");
        return NULL;
    }

    // stage 3: refit with fixed shape parameters
    psTimerStart ("fit");
    for (int i = 0; i < psfTry->sources->n; i++) {
        // masked for: bad model fit, outlier in parameters
        if (psfTry->mask->data.U8[i] & PSFTRY_MASK_ALL)
            continue;

        pmSource *source = psfTry->sources->data[i];

        // set shape for this model based on PSF
        source->modelPSF = pmModelFromPSF (source->modelEXT, psfTry->psf);
        source->modelPSF->radiusFit = RADIUS;
        x = source->peak->x;
        y = source->peak->y;

        // set the mask and fit the PSF model
        psImageKeepCircle (source->mask, x, y, RADIUS, "OR", PM_MASK_MARK);
        status = pmSourceFitModel (source, source->modelPSF, PM_SOURCE_FIT_PSF);
        psImageKeepCircle (source->mask, x, y, RADIUS, "AND", PS_NOT_U8(PM_MASK_MARK));

        // skip poor fits
        if (!status) {
            psfTry->mask->data.U8[i] = PSFTRY_MASK_PSF_FAIL;
            continue;
        }

        // XXX : use a different aperture radius from the fit radius?
        if (!pmSourceMagnitudes (source, psfTry->psf, PM_SOURCE_PHOT_INTERP)) {
            psfTry->mask->data.U8[i] = PSFTRY_MASK_BAD_PHOT;
            continue;
        }

        psfTry->fitMag->data.F64[i] = source->psfMag;
        psfTry->metric->data.F64[i] = source->apMag - source->psfMag;
        psfTry->metricErr->data.F64[i] = source->errMag;
        Npsf ++;
    }
    psfTry->psf->nPSFstars = Npsf;

    psLogMsg ("psphot.psftry", 4, "fit psf:   %f sec for %d of %ld sources\n", psTimerMark ("fit"), Npsf, sources->n);
    psTrace ("psphot.psftry", 3, "keeping %d of %ld PSF candidates (PSF)\n", Npsf, sources->n);

    // measure the chi-square trend as a function of flux (PAR[PM_PAR_I0])
    // this should be linear for Poisson errors and quadratic for constant sky errors
    psVector *flux  = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);
    psVector *chisq = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);
    psVector *mask  = psVectorAlloc (psfTry->sources->n, PS_TYPE_MASK);

    // generate the x and y vectors, and mask missing models
    for (int i = 0; i < psfTry->sources->n; i++) {
        pmSource *source = psfTry->sources->data[i];
        if (source->modelPSF == NULL) {
            flux->data.F64[i] = 0.0;
            chisq->data.F64[i] = 0.0;
            mask->data.U8[i] = 1;
        } else {
            flux->data.F64[i] = source->modelPSF->params->data.F32[PM_PAR_I0];
            chisq->data.F64[i] = source->modelPSF->chisq / source->modelPSF->nDOF;
            mask->data.U8[i] = 0;
        }
    }

    // use 3hi/3lo sigma clipping on the chisq fit
    psStats *stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    stats->clipSigma = 3.0;
    stats->clipIter = 3;

    // linear clipped fit of chisq trend vs flux
    bool result = psVectorClipFitPolynomial1D(psfTry->psf->ChiTrend, stats, mask, 1, chisq, NULL, flux);
    psLogMsg ("pmPSFtry", 4, "chisq vs flux fit: %f +/- %f\n", stats->robustMedian, stats->robustStdev);
    if (!result) {
	result = psVectorFitPolynomial1D(psfTry->psf->ChiTrend, mask, 1, chisq, NULL, flux);
    }

    psFree (flux);
    psFree (stats);
    psFree (mask);
    psFree (chisq);

    if (!result) {
        psError(PS_ERR_UNKNOWN, false, "Failed to fit psf->ChiTrend");
        psFree(psfTry);
        return NULL;
    }

    for (int i = 0; i < psfTry->psf->ChiTrend->nX + 1; i++) {
        psLogMsg ("pmPSFtry", 4, "chisq vs flux fit term %d: %f +/- %f\n", i,
                  psfTry->psf->ChiTrend->coeff[i]*pow(10000, i),
                  psfTry->psf->ChiTrend->coeffErr[i]*pow(10000,i));
    }

    // XXX this function wants aperture radius for pmSourcePhotometry
    pmPSFtryMetric (psfTry, RADIUS);
    psLogMsg ("psphot.pspsf", 3, "try model %s, ap-fit: %f +/- %f : sky bias: %f\n",
              modelName, psfTry->psf->ApResid, psfTry->psf->dApResid, psfTry->psf->skyBias);

    return (psfTry);
}

bool pmPSFtryMetric (pmPSFtry *psfTry, float RADIUS)
{
    // the measured (aperture - fit) magnitudes (dA == psfTry->metric)
    //   depend on both the true ap-fit (dAo) and the bias in the sky measurement:
    //     dA = dAo + dsky/flux
    //   where flux is the flux of the star
    // we fit this trend to find the infinite flux aperture correction (dAo),
    //   the nominal sky bias (dsky), and the error on dAo
    // the values of dA are contaminated by stars with close neighbors in the aperture
    //   we use an outlier rejection to avoid this bias

    // r2rflux = radius^2 * ten(0.4*fitMag);
    psVector *r2rflux = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);

    for (int i = 0; i < psfTry->sources->n; i++) {
        if (psfTry->mask->data.U8[i] & PSFTRY_MASK_ALL)
            continue;
        r2rflux->data.F64[i] = PS_SQR(RADIUS) * pow(10.0, 0.4*psfTry->fitMag->data.F64[i]);
    }

    // use 3hi/1lo sigma clipping on the r2rflux vs metric fit
    psStats *stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    stats->min = 1.0;
    stats->max = 3.0;
    stats->clipIter = 3;

    // fit ApTrend only to r2rflux, ignore x,y,flux variations for now
    // linear clipped fit of ApResid to r2rflux
    psPolynomial1D *poly = psPolynomial1DAlloc (PS_POLYNOMIAL_ORD, 1);

    // XXX test to only fit a constant offset (no SKYBIAS)
    poly->mask[1] = 1;

    bool result = psVectorClipFitPolynomial1D (poly, stats, psfTry->mask, PSFTRY_MASK_ALL, psfTry->metric, psfTry->metricErr, r2rflux);
    if (!result) {
        psLogMsg(__func__, 3, "Failed to fit clipped poly");
	if (!psVectorFitPolynomial1D(poly, psfTry->mask, PSFTRY_MASK_ALL,
				     psfTry->metric, psfTry->metricErr, r2rflux)) {
	    return false;
	}
    }
    psLogMsg ("pmPSFtryMetric", 4, "fit stats: %f +/- %f\n", stats->robustMedian, stats->robustStdev);
    psLogMsg ("pmPSFtryMetric", 4, "apresid: %f +/- %f\n", poly->coeff[0], poly->coeffErr[0]);

    // XXX test dump of fitted model (dump when tracing?)
    if (psTraceGetLevel("psModules.objects") >= 4) {
        FILE *f = fopen ("resid.dat", "w");
        psVector *apfit = psPolynomial1DEvalVector (poly, r2rflux);
        for (int i = 0; i < psfTry->sources->n; i++) {
            int keep = (psfTry->mask->data.U8[i] & PSFTRY_MASK_ALL);

            pmSource *source = psfTry->sources->data[i];
            float x = source->peak->x;
            float y = source->peak->y;

            fprintf (f, "%d  %d, %f %f %f  %f %f %f  %f\n",
                     i, keep, x, y,
                     psfTry->fitMag->data.F64[i],
                     r2rflux->data.F64[i],
                     psfTry->metric->data.F64[i],
                     psfTry->metricErr->data.F64[i],
                     apfit->data.F64[i]);
        }
        fclose (f);
        psFree (apfit);
    }

    pmPSFMaskApTrend (psfTry->psf->ApTrend, PM_PSF_APTREND_SKYBIAS);
    psfTry->psf->ApTrend->coeff[0][0][0][0] = poly->coeff[0];
    psfTry->psf->ApTrend->coeff[0][0][1][0] = 0;

    psfTry->psf->skyBias  = poly->coeff[1];
    psfTry->psf->ApResid  = poly->coeff[0];
    psfTry->psf->dApResid = stats->robustStdev;

    psFree (r2rflux);
    psFree (poly);
    psFree (stats);

    return true;
}

/*
  (aprMag' - fitMag) = rflux*skyBias + ApTrend(x,y)
  (aprMag - rflux*skyBias) - fitMag = ApTrend(x,y)
  (aprMag - rflux*skyBias) = fitMag + ApTrend(x,y)
*/

/*****************************************************************************
pmPSFFromPSFtry (psfTry): build a PSF model from a collection of
source->modelEXT entires.  The PSF ignores the first 4 (independent) model
parameters and constructs a polynomial fit to the remaining as a function of
image coordinate.
    input: psfTry with fitted source->modelEXT collection, pre-allocated psf
Note: some of the array entries may be NULL (failed fits); ignore them.
 *****************************************************************************/
bool pmPSFFromPSFtry (pmPSFtry *psfTry, bool applyWeight)
{
    pmPSF *psf = psfTry->psf;

    // construct the fit vectors from the collection of objects
    psVector *x  = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);
    psVector *y  = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);
    psVector *z  = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);

    psVector *dz = NULL;
    if (applyWeight) {
        dz = psVectorAlloc (psfTry->sources->n, PS_TYPE_F64);
    }

    // construct the x,y terms
    for (int i = 0; i < psfTry->sources->n; i++) {
        pmSource *source = psfTry->sources->data[i];
        if (source->modelEXT == NULL)
            continue;

        // use F64 for polynomial fitting
        x->data.F64[i] = source->modelEXT->params->data.F32[PM_PAR_XPOS];
        y->data.F64[i] = source->modelEXT->params->data.F32[PM_PAR_YPOS];
    }

    // 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 parameter, so potentially
    // different sets of stars are used for SXX, SYY, SXY, ...

    psStats *stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);

    // skip the unfitted parameters (X, Y, Io, Sky)
    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;

        // select the per-object fitted data for this parameter
        for (int j = 0; j < psfTry->sources->n; j++) {
            pmSource *source = psfTry->sources->data[j];
            if (source->modelEXT == NULL)
                continue;

            z->data.F64[j] = source->modelEXT->params->data.F32[i];
            if (applyWeight) {
                dz->data.F64[j] = source->modelEXT->dparams->data.F32[i];
            }

            // for SXY, we actually fit SXY / (SXX^-2  + SYY^-2)
            if (i == PM_PAR_SXY) {
                z->data.F64[j] = pmPSF_SXYfromModel (source->modelEXT->params->data.F32);
            }
        }

        // fit the collection of measured parameters to the PSF 2D model
        // the mask is carried from previous steps and updated with this operation
        // the weight is either the parameter error or NULL, depending on 'applyWeights'
        if (!psVectorClipFitPolynomial2D(psf->params_NEW->data[i], stats, psfTry->mask, 0xff, z, NULL, x, y)) {
	    // Should save error stack at this point, and add any new errors on the end if
	    // we still have problems; but we don't have an API to do this
	    psLogMsg ("pmPSFtry", 4, "failed to build psf model for parameter %d with clipping", i);

	    if (!psVectorFitPolynomial2D(psf->params_NEW->data[i], psfTry->mask, 0xff, z, NULL, x, y)) {
		psError(PS_ERR_UNKNOWN, false, "failed to build unclipped psf model for parameter %d", i);
		return false;
	    }
        }
    }

    // test dump of star parameters vs position (compare with fitted values)
    if (psTraceGetLevel("psModules.objects") >= 4) {
        FILE *f = fopen ("params.dat", "w");

        for (int j = 0; j < psfTry->sources->n; j++) {
            pmSource *source = psfTry->sources->data[j];
            if (source == NULL)
                continue;

            pmModel *model = source->modelEXT;
            if (model == NULL)
                continue;

            pmModel *modelPSF = pmModelFromPSF (model, psf);

            fprintf (f, "%f %f : ", model->params->data.F32[PM_PAR_XPOS], model->params->data.F32[PM_PAR_YPOS]);

            for (int i = 0; i < psf->params_NEW->n; i++) {
                if (psf->params_NEW->data[i] == NULL)
                    continue;
                fprintf (f, "%f %f : ", model->params->data.F32[i], modelPSF->params->data.F32[i]);
            }
            fprintf (f, "%f %d\n", model->chisq, model->nIter);

            psFree(modelPSF);
        }
        fclose (f);
    }

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

