/* @file  pmSourceFitPCM.c
 * structures and functions to support PSF-convolved model fitting
 *
 * @author EAM, IfA
 *
 * @version $Revision: 1.29 $
 * @date $Date: 2009-02-16 22:30:50 $
 * Copyright 2010 Institute for Astronomy, University of Hawaii
 */

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

#include <stdio.h>
#include <math.h>
#include <string.h>
#include <strings.h>
#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 "pmModel.h"
#include "pmModelUtils.h"
#include "pmModelClass.h"
#include "pmSourceMasks.h"
#include "pmSourceExtendedPars.h"
#include "pmSourceDiffStats.h"
#include "pmSourceSatstar.h"
#include "pmSourceLensing.h"
#include "pmSource.h"
#include "pmSourcePhotometry.h"
#include "pmSourceFitModel.h"
#include "pmPCMdata.h"

# define FACILITY "psModules.objects"

// input source has both modelPSF and modelEXT.  on successful exit, we set the
// modelConv to contain the fitted parameters, and the modelFlux to contain the 
// convolved model image.

# define TIMING 0

bool pmSourceChisqModelFlux (pmSource *source, pmModel *model, psImageMaskType maskVal);

bool pmSourceFitPCM (pmPCMdata *pcm, pmSource *source, pmSourceFitOptions *fitOptions, psImageMaskType maskVal, psImageMaskType markVal, int psfSize) {
    
    if (TIMING) { psTimerStart ("pmSourceFitPCM"); }

    psVector *params  = pcm->modelConv->params;
    psVector *dparams  = pcm->modelConv->dparams;

    // force the floating parameters to fall within the contraint ranges
    for (int i = 0; i < params->n; i++) {
	pcm->modelConv->modelLimits (PS_MINIMIZE_PARAM_MIN, i, params->data.F32, NULL);
	pcm->modelConv->modelLimits (PS_MINIMIZE_PARAM_MAX, i, params->data.F32, NULL);
    }

    // set up the minimization process
    psMinimization *myMin = psMinimizationAlloc (fitOptions->nIter, fitOptions->minTol, fitOptions->maxTol);
    myMin->chisqConvergence = fitOptions->chisqConvergence;
    myMin->gainFactorMode = fitOptions->gainFactorMode;
    myMin->isInteractive = fitOptions->isInteractive;

    psImage *covar = psImageAlloc (params->n, params->n, PS_TYPE_F32);
    // NOTE : 4 allocs to here

    float t1, t2, t3, t4, t5;
    if (TIMING) { t1 = psTimerMark ("pmSourceFitPCM"); }

    // NOTE : 996 allocs in here
    bool fitStatus = pmPCM_MinimizeChisq (myMin, covar, params, source, pcm);
    if (TIMING) { t2 = psTimerMark ("pmSourceFitPCM"); }

    for (int i = 0; i < dparams->n; i++) {
        if ((pcm->constraint->paramMask != NULL) && pcm->constraint->paramMask->data.PS_TYPE_VECTOR_MASK_DATA[i])
            continue;
        dparams->data.F32[i] = sqrt(covar->data.F32[i][i]);
        psTrace ("psModules.objects", 4, "%f +/- %f", params->data.F32[i], dparams->data.F32[i]);
    }
    if (fitOptions->saveCovariance) {
	psFree(pcm->modelConv->covar);
	pcm->modelConv->covar = psMemIncrRefCounter(covar);
    }
    psTrace ("psphot", 4, "niter: %d, chisq: %f", myMin->iter, myMin->value);
    if (TIMING) { t3 = psTimerMark ("pmSourceFitPCM"); }

    if (fitStatus && source->modelFlux) {
        // renormalize output model image (generated by fitting process)
        float Io = params->data.F32[PM_PAR_I0];
        for (int iy = 0; iy < source->modelFlux->numRows; iy++) {
            for (int ix = 0; ix < source->modelFlux->numCols; ix++) {
                source->modelFlux->data.F32[iy][ix] /= Io;
            }
        }
    }
    pcm->modelConv->nIter = myMin->iter;
    pcm->modelConv->nPar = pcm->nPar;

    // save the resulting chisq, nDOF, nIter
    if (fitOptions->poissonErrors) {
	pcm->modelConv->chisq = myMin->value;
	pcm->modelConv->nPix = pcm->nPix;
	pcm->modelConv->nDOF = pcm->nDOF;
	pcm->modelConv->chisqNorm = pcm->modelConv->chisq / pcm->modelConv->nDOF;
    } else {
	// xxx this is wrong because it does not convolve with the psf
	pmPCMMakeModel (source, pcm->modelConv, pcm->nsigma, maskVal, psfSize);
	pmSourceChisqModelFlux (source, pcm->modelConv, maskVal);
    }
    if (TIMING) { t4 = psTimerMark ("pmSourceFitPCM"); }

    // set the model success or failure status
    pcm->modelConv->flags |= PM_MODEL_STATUS_FITTED;

    if (!fitStatus) {
	if (isnan(myMin->value)) {
	    pcm->modelConv->flags |= PM_MODEL_STATUS_NAN_CHISQ;
	} else {
	    pcm->modelConv->flags |= PM_MODEL_STATUS_NONCONVERGE;
	}
    }

    if (myMin->chisqConvergence) {
      if (myMin->lastDelta > myMin->minTol) pcm->modelConv->flags |= PM_MODEL_STATUS_WEAK_FIT;
    } else {
      if (myMin->rParSigma > myMin->minTol*pcm->nPar) pcm->modelConv->flags |= PM_MODEL_STATUS_WEAK_FIT;
    }

    // once we have fitted a model, we need to record that this model is a PCM model:
    pcm->modelConv->isPCM = true;

    // models can go insane: reject these
    bool onPic = true;
    onPic &= (params->data.F32[PM_PAR_XPOS] >= source->pixels->col0);
    onPic &= (params->data.F32[PM_PAR_XPOS] <  source->pixels->col0 + source->pixels->numCols);
    onPic &= (params->data.F32[PM_PAR_YPOS] >= source->pixels->row0);
    onPic &= (params->data.F32[PM_PAR_YPOS] <  source->pixels->row0 + source->pixels->numRows);
    if (!onPic) pcm->modelConv->flags |= PM_MODEL_STATUS_OFFIMAGE;

    source->mode |= PM_SOURCE_MODE_FITTED; // XXX is this needed?
    if (TIMING) { t5 = psTimerMark ("pmSourceFitPCM"); }

     if (TIMING) {
 	fprintf (stderr, "nIter: %2d, npix: %5d, t1: %6.4f, t2: %6.4f, t3: %6.4f, t4: %6.4f, t5: %6.4f\n", myMin->iter, pcm->nPix, t1, t2, t3, t4, t5);
     }

    psFree(myMin);
    psFree(covar);

    return true;
}

// XXX deprecate this function or merge with the empirical model
bool pmSourceModelGuessPCM (pmPCMdata *pcm, pmSource *source, psImageMaskType maskVal, psImageMaskType markVal) {

    if (!pcm->modelConv->modelGuess(pcm->modelConv, source, maskVal, markVal)) {
	return false;
    }
    return true;

    // the following was an attempt to make analytical modifications to the shape terms based on the psf
    // this has been replaced with a more empirical approach
# if (0)
    // generate copy of the model
    // XXX we could modify the parameter values or even the model 
    // here based on the observed seeing (some lookup table...)

    // XXX test : modify the Io, SXX, SYY terms based on the psf SXX, SYY terms:
    // SXX,SYY in model parameters are sqrt(2) * shape.Sxx,Syy
    psEllipseShape psfShape;
    psfShape.sx  = source->modelPSF->params->data.F32[PM_PAR_SXX] / M_SQRT2;
    psfShape.sxy = source->modelPSF->params->data.F32[PM_PAR_SXY];
    psfShape.sy  = source->modelPSF->params->data.F32[PM_PAR_SYY] / M_SQRT2;

    psEllipseAxes psfAxes = psEllipseShapeToAxes (psfShape, 20.0);
    if (!isfinite(psfAxes.major)) return false;
    if (!isfinite(psfAxes.minor)) return false;
    if (!isfinite(psfAxes.theta)) return false;

    // XXX test : modify the Io, SXX, SYY terms based on the psf SXX, SYY terms:
    psEllipseShape extShape;
    extShape.sx  = pcm->modelConv->params->data.F32[PM_PAR_SXX] / M_SQRT2;
    extShape.sxy = pcm->modelConv->params->data.F32[PM_PAR_SXY];
    extShape.sy  = pcm->modelConv->params->data.F32[PM_PAR_SYY] / M_SQRT2;

    psEllipseAxes extAxes = psEllipseShapeToAxes (extShape, 20.0);
    if (!isfinite(extAxes.major)) return false;
    if (!isfinite(extAxes.minor)) return false;
    if (!isfinite(extAxes.theta)) return false;

    // decrease the initial guess ellipse by psf_minor axis:
    psEllipseAxes extAxesMod;
    extAxesMod.major = sqrt (PS_MAX (0.25, PS_SQR(extAxes.major) - PS_SQR(psfAxes.minor)));
    extAxesMod.minor = sqrt (PS_MAX (0.25, PS_SQR(extAxes.minor) - PS_SQR(psfAxes.minor)));
    extAxesMod.theta = extAxes.theta;

    psEllipseShape extShapeMod = psEllipseAxesToShape (extAxesMod);
    if (!isfinite(extShapeMod.sx))  return false;
    if (!isfinite(extShapeMod.sy))  return false;
    if (!isfinite(extShapeMod.sxy)) return false;

    pcm->modelConv->params->data.F32[PM_PAR_SXX] = extShapeMod.sx * M_SQRT2;
    pcm->modelConv->params->data.F32[PM_PAR_SXY] = extShapeMod.sxy;
    pcm->modelConv->params->data.F32[PM_PAR_SYY] = extShapeMod.sy * M_SQRT2;

    // increase the initial guess central intensity by 2pi r^2:
    pcm->modelConv->params->data.F32[PM_PAR_I0] *= (1.0 + PS_SQR(psfAxes.minor) / PS_SQR(extAxesMod.minor));

    return true;
# endif
}
