# include <pslib.h>
# include "psLibUtils.h"
# include "pmObjects.h"
# include "psModulesUtils.h"
# include "pmPSF.h"
# include "pmPSFtry.h"
# include "pmModelGroup.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->modelFLT);
    psFree (test->modelPSF);
    psFree (test->metric);
    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)
{

    pmModelType type;

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

    // XXX probably need to increment ref counter
    type           = pmModelSetType (modelName);
    test->psf      = pmPSFAlloc (type);
    test->sources  = psMemCopy(sources);
    test->modelFLT = psArrayAlloc (sources->n);
    test->modelPSF = psArrayAlloc (sources->n);
    test->metric   = psVectorAlloc (sources->n, PS_TYPE_F64);
    test->fitMag   = psVectorAlloc (sources->n, PS_TYPE_F64);
    test->mask     = psVectorAlloc (sources->n, PS_TYPE_U8);

    for (int i = 0; i < test->modelFLT->n; i++) {
        test->mask->data.U8[i]  = 0;
        test->modelFLT->data[i] = NULL;
        test->modelPSF->data[i] = NULL;
        test->metric->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:

pmPSFtry *pmPSFtryModel (psArray *sources, char *modelName, float RADIUS)
{
    bool status;
    float obsMag;
    float fitMag;
    float x;
    float y;
    int Nflt = 0;
    int Npsf = 0;

    pmPSFtry *try = pmPSFtryAlloc (sources, modelName);

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

	pmSource *source = try->sources->data[i];
	pmModel  *model  = pmSourceModelGuess (source, try->psf->type);
	x = source->peak->x;
	y = source->peak->y;

	// set temporary object mask and fit object
	// fit model as FLT, not PSF
	psImageKeepCircle (source->mask, x, y, RADIUS, "OR", PSPHOT_MASK_MARKED);
	status = pmSourceFitModel (source, model, false);
	psImageKeepCircle (source->mask, x, y, RADIUS, "AND", ~PSPHOT_MASK_MARKED);

	// exclude the poor fits
	if (!status) {
	    try->mask->data.U8[i] = PSFTRY_MASK_FLT_FAIL;
	    psFree (model);
	    continue;
	}
	try->modelFLT->data[i] = model;
	Nflt ++;
    }
    psLogMsg ("psphot.psftry", 4, "fit flt:   %f sec for %d sources\n", psTimerMark ("fit"), sources->n);
    psTrace ("psphot.psftry", 3, "keeping %d of %d PSF candidates (FLT)\n", Nflt, sources->n);

    // make this optional?
    // DumpModelFits (try->modelFLT, "modelsFLT.dat");

    // stage 2: construct a psf (pmPSF) from this collection of model fits
    pmPSFFromModels (try->psf, try->modelFLT, try->mask);

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

	pmSource *source = try->sources->data[i];
	pmModel  *modelFLT = try->modelFLT->data[i];

	// set shape for this model based on PSF
	pmModel *modelPSF = pmModelFromPSF (modelFLT, try->psf);
	x = source->peak->x;
	y = source->peak->y;

	psImageKeepCircle (source->mask, x, y, RADIUS, "OR", PSPHOT_MASK_MARKED);
	status = pmSourceFitModel (source, modelPSF, true);

	// skip poor fits
	if (!status) {
	    try->mask->data.U8[i] = PSFTRY_MASK_PSF_FAIL;
	    psFree (modelPSF);
	    goto next_source;
	}

	// otherwise, save the resulting model
	try->modelPSF->data[i] = modelPSF;

	// XXX : use a different aperture radius from the fit radius?
	// XXX : use a different estimator for the local sky?
	// XXX : pass 'source' as input?
	if (!pmSourcePhotometry (&fitMag, &obsMag, modelPSF, source->pixels, source->mask)) {
	    try->mask->data.U8[i] = PSFTRY_MASK_BAD_PHOT;
	    goto next_source;
	}

	try->metric->data.F64[i] = obsMag - fitMag;
	try->fitMag->data.F64[i] = fitMag;
	Npsf ++;

    next_source:
	psImageKeepCircle (source->mask, x, y, RADIUS, "AND", ~PSPHOT_MASK_MARKED);

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

    // make this optional
    // DumpModelFits (try->modelPSF, "modelsPSF.dat");

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

    return (try);
}


bool pmPSFtryMetric (pmPSFtry *try
                     , float RADIUS)
{

    float dBin;
    int   nKeep, nSkip;

    // the measured (aperture - fit) magnitudes (dA == try->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

    // rflux = ten(0.4*fitMag);
    psVector *rflux = psVectorAlloc (try->sources->n, PS_TYPE_F64);
    for (int i = 0; i < try->sources->n; i++) {
	if (try->mask->data.U8[i] & PSFTRY_MASK_ALL) continue;
	rflux->data.F64[i] = pow(10.0, 0.4*try->fitMag->data.F64[i]);
    }

    // find min and max of (1/flux):
    psStats *stats = psStatsAlloc (PS_STAT_MIN | PS_STAT_MAX);
    psVectorStats (stats, rflux, NULL, try->mask, PSFTRY_MASK_ALL);

    // build binned versions of rflux, metric
    dBin = (stats->max - stats->min) / 10.0;
    psVector *rfBin = psVectorCreate(NULL, stats->min, stats->max, dBin, PS_TYPE_F64);
    psVector *daBin = psVectorAlloc (rfBin->n, PS_TYPE_F64);
    psVector *maskB = psVectorAlloc (rfBin->n, PS_TYPE_U8);
    psFree (stats);

    psTrace ("psphot.metricmodel", 3, "rflux max: %g, min: %g, delta: %g\n", stats->max, stats->min, dBin);

    // group data in daBin bins, measure lower 50% mean
    for (int i = 0; i < daBin->n; i++) {

        psVector *tmp = psVectorAlloc (try->sources->n, PS_TYPE_F64);
        tmp->n = 0;

        // accumulate data within bin range
        for (int j = 0; j < try->sources->n; j++) {
	    // masked for: bad model fit, outlier in parameters
	    if (try->mask->data.U8[j] & PSFTRY_MASK_ALL) continue;

	    // skip points with extreme dA values
	    if (fabs(try->metric->data.F64[j]) > 0.5) continue;

	    // skip points outside of this bin
	    if (rflux->data.F64[j] < rfBin->data.F64[i] - 0.5*dBin)
		continue;
	    if (rflux->data.F64[j] > rfBin->data.F64[i] + 0.5*dBin)
		continue;

	    tmp->data.F64[tmp->n] = try->metric->data.F64[j];
	    tmp->n ++;
	}

        // is this a valid point?
        maskB->data.U8[i] = 0;
        if (tmp->n < 2) {
            maskB->data.U8[i] = 1;
            psFree (tmp);
            continue;
        }

        // dA values are contaminated with low outliers
        // measure statistics only on upper 50% of points
        // this would be easier if we could sort in reverse:
        //
        // psVectorSort (tmp, tmp);
        // tmp->n = 0.5*tmp->n;
        // stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN);
        // psVectorStats (stats, tmp, NULL, NULL, 0);
        // psTrace ("psphot.metricmodel", 4, "rfBin %d (%g): %d pts, %g\n", i, rfBin->data.F64[i], tmp->n, stats->sampleMedian);

        psVectorSort (tmp, tmp);
        nKeep = 0.5*tmp->n;
        nSkip = tmp->n - nKeep;

        psVector *tmp2 = psVectorAlloc (nKeep, PS_TYPE_F64);
        for (int j = 0; j < tmp2->n; j++) {
            tmp2->data.F64[j] = tmp->data.F64[j + nSkip];
        }

        stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN);
        psVectorStats (stats, tmp2, NULL, NULL, 0);
        psTrace ("psphot.metricmodel", 4, "rfBin %d (%g): %d pts, %g\n", i, rfBin->data.F64[i], tmp->n, stats->sampleMedian);

        daBin->data.F64[i] = stats->sampleMedian;

        psFree (stats);
        psFree (tmp);
        psFree (tmp2);
    }

    // linear fit to rfBin, daBin
    psPolynomial1D *poly = psPolynomial1DAlloc(1, PS_POLYNOMIAL_ORD);
    psStats *fitstat = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);
    poly = psVectorClipFitPolynomial1D (poly, fitstat, maskB, 1, daBin, NULL, rfBin);

    // XXX EAM : this is the intended API (cycle 7? cycle 8?)
    // poly = psVectorFitPolynomial1D(poly, maskB, 1, daBin, NULL, rfBin);

    // XXX EAM : replace this when the above version is implemented
    // poly = psVectorFitPolynomial1DOrd(poly, maskB, rfBin, daBin, NULL);

    psVector *daBinFit = psPolynomial1DEvalVector (poly, rfBin);
    psVector *daResid  = (psVector *) psBinaryOp (NULL, (void *) daBin, "-", (void *) daBinFit);

    stats = psStatsAlloc (PS_STAT_CLIPPED_STDEV);
    stats = psVectorStats (stats, daResid, NULL, maskB, 1);

    try->psf->ApResid = poly->coeff[0];
    try->psf->dApResid = stats->clippedStdev;
    try->psf->skyBias = poly->coeff[1] / (M_PI * PS_SQR(RADIUS));

    psFree (rflux);
    psFree (rfBin);
    psFree (daBin);
    psFree (maskB);
    psFree (daBinFit);
    psFree (daResid);
    psFree (poly);
    psFree (stats);
    psFree (fitstat);

    return true;
}

bool pmPSFtryMetric_Alt (pmPSFtry *try, float RADIUS) {

    // the measured (aperture - fit) magnitudes (dA == try->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

    // rflux = ten(0.4*fitMag);
    psVector *rflux = psVectorAlloc (try->sources->n, PS_TYPE_F64);
    for (int i = 0; i < try->sources->n; i++) {
	if (try->mask->data.U8[i] & PSFTRY_MASK_ALL) continue;
	rflux->data.F64[i] = pow(10.0, 0.4*try->fitMag->data.F64[i]);
    }

    // XXX EAM : try 3hi/1lo sigma clipping on the rflux vs metric fit
    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);

    // XXX EAM 
    stats->min = 1.0;
    stats->max = 3.0;
    stats->clipIter = 3;

    // linear clipped fit to rfBin, daBin
    psPolynomial1D *poly = psPolynomial1DAlloc (1, PS_POLYNOMIAL_ORD);
    poly = psVectorClipFitPolynomial1D (poly, stats, try->mask, PSFTRY_MASK_ALL, try->metric, NULL, rflux);
    fprintf (stderr, "fit stats: %f +/- %f\n", stats->sampleMedian, stats->sampleStdev);

    try->psf->ApResid = poly->coeff[0];
    try->psf->dApResid = stats->sampleStdev;
    try->psf->skyBias = poly->coeff[1] / (M_PI * PS_SQR(RADIUS));

    fprintf (stderr, "*******************************************************************************\n");

    FILE *f;
    f = fopen ("apresid.dat", "w");
    if (f == NULL) psAbort ("pmPSFtry", "can't open output file");

    for (int i = 0; i < try->sources->n; i++) {
	fprintf (f, "%3d %8.4f %12.5e %8.4f %3d\n", i, try->fitMag->data.F64[i], rflux->data.F64[i], try->metric->data.F64[i], try->mask->data.U8[i]);
    }
    fclose (f);

    psFree (rflux);
    psFree (poly);
    psFree (stats);

    // psFree (daFit);
    // psFree (daResid);

    return true;
}