/** @file  pmPSFtry.c
 *
 *  XXX: need description of file purpose
 *
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.69 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-01-27 06:39:38 $
 *  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 "pmTrend2D.h"
#include "pmResiduals.h"
#include "pmGrowthCurve.h"
#include "pmSpan.h"
#include "pmFootprintSpans.h"
#include "pmFootprint.h"
#include "pmPeaks.h"
#include "pmMoments.h"
#include "pmModelFuncs.h"
#include "pmModelClass.h"
#include "pmModel.h"
#include "pmModelUtils.h"
#include "pmSourceMasks.h"
#include "pmSourceExtendedPars.h"
#include "pmSourceDiffStats.h"
#include "pmSourceSatstar.h"
#include "pmSourceLensing.h"
#include "pmSource.h"
#include "pmSourceFitModel.h"
#include "pmPSF.h"
#include "pmPSFtry.h"
#include "pmDetections.h"

#include "pmSourceVisual.h"

// 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:

// XXX some test code (delete eventually)
// bool countMaskedSources(pmPSFtry *try, char *msg) {
// 
//     int N1 = 0;
//     int N2 = 0;
//     for (int i = 0; i < try->sources->n; i++) {
//         pmSource *source = try->sources->data[i];
//     	// fprintf (stderr, "%llx : %d\n", (long long int) source, (source->mode & PM_SOURCE_MODE_PSFSTAR));
//     	if (source->mode & PM_SOURCE_MODE_PSFSTAR) {
//     	    N1 ++;
//         }
//     	if (try->mask->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
//     	    N2 ++;
//         }
//     }
//     fprintf (stderr, "%s : masked: %d or %d of %ld\n", msg, N1, N2, try->sources->n);
//     return true;
// }

// generate a pmPSFtry with a copy of the test PSF sources
pmPSFtry *pmPSFtryModel (const psArray *sources, const char *modelName, pmPSFOptions *options, psImageMaskType maskVal, psImageMaskType markVal)
{
    assert (options->fitOptions);

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

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

    // maskVal is used to test for rejected pixels, and must include markVal
    maskVal |= markVal;

    // stage 1:  fit an EXT model to all candidates PSF sources -- this is independent of the modeled 2D variations in the PSF
    if (!pmPSFtryFitEXT(psfTry, options, maskVal, markVal)) {
        psError(PS_ERR_UNKNOWN, false, "failed to fit EXT models to sources for psf model");
        psFree(psfTry);
        return NULL;
    }

    // hard limit on minimum number of stars
    if ((sources->n <  3)) {
        psError (PS_ERR_UNKNOWN, true, "failed to determine PSF parameters");
        return NULL;
    }

    // this is a bit tricky, because we have two cases (MAP vs POLY), and they have a different
    // definition for 'order' (order_MAP = order_POLY + 1).  in addition, we have a
    // user-specified MAX order, which we should respect, regardless of the mode

    // set the max order (0 = constant) which the number of psf stars can support:
    int MaxOrderForStars = 0;

    // we require only 3 stars for n = 0, increase stars / cell for higher order
    if (sources->n >=  16) MaxOrderForStars = 1; //  4 cells, 4 per cell
    if (sources->n >=  54) MaxOrderForStars = 2; //  9 cells, 6 per cell
    if (sources->n >= 128) MaxOrderForStars = 3; // 16 cells, 8 per cell
    if (sources->n >= 300) MaxOrderForStars = 4; // 25 cells, 12 per cell
    if (sources->n >  576) MaxOrderForStars = 5; // 36 cells, 16 per cell

    // rule of thumb: require 3 stars per 'cell' (order+1)^2
    // if (sources->n >= 12) MaxOrderForStars = 1; // 4 cells
    // if (sources->n >= 27) MaxOrderForStars = 2; // 9 cells
    // if (sources->n >= 48) MaxOrderForStars = 3; // 16 cells
    // if (sources->n >= 75) MaxOrderForStars = 4; // 25 cells
    // if (sources->n > 108) MaxOrderForStars = 5; // 36 cells

    int orderMax = PS_MAX (options->psfTrendNx, options->psfTrendNy);
    int orderMin = 0;
    if (options->psfTrendMode == PM_TREND_MAP) {
        MaxOrderForStars ++;
        orderMin ++;
    }
    orderMax = PS_MIN (orderMax, MaxOrderForStars);

    // save the raw source mask (generated by pmPSFtryFitEXT)
    psVector *srcMask = psVectorCopy (NULL, psfTry->mask, PS_TYPE_VECTOR_MASK);

    // we will save the PSF with the best fit (min systematic error)
    pmPSF *minPSF = NULL;
    psVector *minMask = NULL;

    // as we loop over orders, we need to refer to the initial selection, but we modify the
    // option values to match the current guess: save the max values here:
    int Nx = (options->psfTrendMode == PM_TREND_MAP) ? options->psfTrendNx : options->psfTrendNx + 1;
    int Ny = (options->psfTrendMode == PM_TREND_MAP) ? options->psfTrendNy : options->psfTrendNy + 1;
    for (int i = orderMin; i <= orderMax; i++) {

        if (Nx > Ny) {
            options->psfTrendNx = i;
            options->psfTrendNy = PS_MAX (orderMin, (int)(i * (Ny / (float) Nx) + 0.5));
        } else {
            options->psfTrendNy = i;
            options->psfTrendNx = PS_MAX (orderMin, (int)(i * (Nx / (float) Ny) + 0.5));
        }

	fprintf (stderr, "fitting %d x %d model for PSF\n", options->psfTrendNx, options->psfTrendNy);

        // free existing data, if any
        psFree(psfTry->psf);
        psFree(psfTry->mask);

        // allocate a mask and a psf model using the current Nx,Ny order values;
        psfTry->psf = pmPSFAlloc (options);
        psfTry->mask = psVectorCopy (NULL, srcMask, PS_TYPE_VECTOR_MASK);

        // stage 2: construct a psf (pmPSF) from this collection of model fits, including the 2D variation
	bool goodFit = false;
        if (!pmPSFtryMakePSF (&goodFit, psfTry)) {
            psError(PS_ERR_UNKNOWN, false, "failed to construct a psf model from collection of sources");
            psFree(psfTry);
            return NULL;
        }
	if (!goodFit) {
	    psWarning ("poor psf fit for order %d, skipping\n", i);
	    continue;
	}

        // stage 3: refit with fixed shape parameters, measure pmPSFtry->metric
        if (!pmPSFtryFitPSF (psfTry, options, maskVal, markVal)) {
            psError(PS_ERR_UNKNOWN, false, "failed to construct a psf model from collection of sources");
            psFree(psfTry);
            return NULL;
        }

        // stage 4: measure systematic error in pmPSFtry->metric
        if (!pmPSFtryMetric (psfTry)) {
            psError(PS_ERR_UNKNOWN, false, "failed to measure systematic error of metric");
            psFree(psfTry);
            return NULL;
        }

        if (!minPSF) {
            minPSF = psMemIncrRefCounter(psfTry->psf);
            minMask = psMemIncrRefCounter(psfTry->mask);
        }

        if (psfTry->psf->dApResid < minPSF->dApResid) {
            psFree (minPSF);
            psFree (minMask);
            minPSF = psMemIncrRefCounter(psfTry->psf);
            minMask = psMemIncrRefCounter(psfTry->mask);
        }
    }
    psFree (srcMask);

    if (!minPSF) {
	psError(PS_ERR_UNKNOWN, false, "failed to construct a valid psf model from the sources");
	psFree(psfTry);
	return NULL;
    }

    // keep the ones matching the min systematic error:
    psFree (psfTry->psf);
    psFree (psfTry->mask);
    psfTry->psf = minPSF;
    psfTry->mask = minMask;

    // XXXXX this is probably not used any more.  Are the chisq of the fits so bad? can we
    // fix them by softening the errors on the brightest pixels?

    // 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_F32);
    psVector *chisq = psVectorAlloc (psfTry->sources->n, PS_TYPE_F32);
    psVector *mask  = psVectorAlloc (psfTry->sources->n, PS_TYPE_VECTOR_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.F32[i] = 0.0;
            chisq->data.F32[i] = 0.0;
            mask->data.PS_TYPE_VECTOR_MASK_DATA[i] = 0xff;
            psTrace ("psModules.objects", 4, "dropping %d (%d,%d) : no PSF model\n", i, source->peak->x, source->peak->y);
        } else {
            flux->data.F32[i] = source->modelPSF->params->data.F32[PM_PAR_I0];
            chisq->data.F32[i] = source->modelPSF->chisq / source->modelPSF->nDOF;
            mask->data.PS_TYPE_VECTOR_MASK_DATA[i] = 0;
        }
    }

    // use 3hi/3lo sigma clipping on the chisq fit
    psStats *stats = options->stats;

    // linear clipped fit of chisq trend vs flux
    if (options->chiFluxTrend) {

	if (0) {
	    FILE *f = fopen ("test.psf.dat", "w");
	    int fd = fileno (f);
	    p_psVectorPrint (fd, flux, "flux");
	    p_psVectorPrint (fd, chisq, "chisq");
	    p_psVectorPrint (fd, mask, "mask");
	    fclose (f);
	}

        bool result = psVectorClipFitPolynomial1D(psfTry->psf->ChiTrend, options->stats,
                                                  mask, 0xff, chisq, NULL, flux);
        psStatsOptions meanStat = psStatsMeanOption(options->stats->options); // Statistic for mean
        psStatsOptions stdevStat = psStatsStdevOption(options->stats->options); // Statistic for stdev

        psLogMsg ("pmPSFtry", 4, "chisq vs flux fit: %f +/- %f\n",
                  psStatsGetValue(stats, meanStat), psStatsGetValue(stats, stdevStat));

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

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

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

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

