# include "psphotInternal.h"

// given a source with an existing modelPSF, attempt a full PSF fit, subtract if successful
// XXX this function does not call pmSourceFitModelInit : fix this?

static int NfitPSF = 0;
static int NfitBlend = 0;
static int NfitDBL = 0;
static int NfitEXT = 0;

bool psphotFitInit (int nThreads) {
    psTimerStart ("psphot.fits");
    pmSourceFitSetInit (nThreads);
    return true;
}

bool psphotFitSummary () {

    psLogMsg ("psphot.pspsf", PS_LOG_INFO, "fitted %5d psf, %5d blend, %5d ext, %5d dbl : %6.2f sec\n",
             NfitPSF, NfitBlend, NfitEXT, NfitDBL, psTimerMark ("psphot.fits"));
    return true;
}

bool psphotFitBlend (pmReadout *readout, pmSource *source, pmPSF *psf, psImageMaskType maskVal, psImageMaskType markVal) {

    float x, y, dR;

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

    // if this source is not a possible blend, just fit as PSF
    if ((source->blends == NULL) || (source->mode & PM_SOURCE_MODE_SATSTAR)) {
        bool status = psphotFitPSF (readout, source, psf, maskVal, markVal);
        return status;
    }

    // save the PSF model from the Ensemble fit
    pmModel *PSF = pmModelCopy (source->modelPSF);

    if (isnan(PSF->params->data.F32[PM_PAR_I0])) psAbort("nan in blend fit primary");

    x = PSF->params->data.F32[PM_PAR_XPOS];
    y = PSF->params->data.F32[PM_PAR_YPOS];

    psArray *modelSet = psArrayAllocEmpty (source->blends->n + 1);
    psArrayAdd (modelSet, 16, PSF);

    psArray *sourceSet = psArrayAllocEmpty (source->blends->n + 1);
    psArrayAdd (sourceSet, 16, source);

    psTrace ("psphot", 4, "fitting blended source at %f %f : %f\n", source->peak->xf, source->peak->yf, source->peak->flux);

    // we need to include all blends in the fit (unless primary is saturated?)
    dR = 0;
    for (int i = 0; i < source->blends->n; i++) {
        pmSource *blend = source->blends->data[i];

        // find the blend which is furthest from source
        dR = PS_MAX (dR, hypot (blend->peak->xf - x, blend->peak->yf - y));

        // create the model and guess parameters for this blend
        pmModel *model = pmModelAlloc (PSF->type);
        for (int j = 0; j < model->params->n; j++) {
            model->params->data.F32[j] = PSF->params->data.F32[j];
            model->dparams->data.F32[j] = PSF->dparams->data.F32[j];
        }

        // XXX assume local sky is 0.0?
        model->params->data.F32[PM_PAR_I0] = blend->peak->flux;
        model->params->data.F32[PM_PAR_XPOS] = blend->peak->xf;
        model->params->data.F32[PM_PAR_YPOS] = blend->peak->yf;

        // these should never be invalid values
        // XXX drop these tests eventually
        if (isnan(model->params->data.F32[PM_PAR_I0]))   psAbort("nan in blend fit");
        if (isnan(model->params->data.F32[PM_PAR_XPOS])) psAbort("nan in blend fit");
        if (isnan(model->params->data.F32[PM_PAR_YPOS])) psAbort("nan in blend fit");

        // add this blend to the list
        psArrayAdd (modelSet, 16, model);
        psArrayAdd (sourceSet, 16, blend);

	psTrace ("psphot", 5, "adding source at %f %f : %f\n", blend->peak->xf, blend->peak->yf, blend->peak->flux);

        // free to avoid double counting model
        psFree (model);
    }

    // extend source radius as needed
    psphotCheckRadiusPSFBlend (readout, source, PSF, markVal, dR);

    // fit PSF model
    pmSourceFitSet (source, modelSet, PM_SOURCE_FIT_PSF, maskVal);

    // clear the circular mask
    psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal)); 

    if (!isfinite(PSF->params->data.F32[PM_PAR_I0])) psAbort("nan in fit");

    // correct model chisq for flux trend
    double chiTrend = psPolynomial1DEval (psf->ChiTrend, PSF->params->data.F32[PM_PAR_I0]);
    PSF->chisqNorm = PSF->chisq / chiTrend;

    // evaluate the blend objects, subtract if good, free otherwise
    for (int i = 1; i < modelSet->n; i++) {
        pmSource *blend = sourceSet->data[i];
        pmModel *model  = modelSet->data[i];

	if (!isfinite(model->params->data.F32[PM_PAR_I0])) psAbort("nan in fit");

        // correct model chisq for flux trend
        chiTrend = psPolynomial1DEval (psf->ChiTrend, model->params->data.F32[PM_PAR_I0]);
        model->chisqNorm = model->chisq / chiTrend;

        // if this one failed, skip it
        if (!psphotEvalPSF (blend, model)) {
            psTrace ("psphot", 4, "failed on blend %d of %ld\n", i, modelSet->n);
            continue;
        }

        // otherwise, supply the resulting model to the corresponding blend
        psFree(blend->modelPSF);
        blend->modelPSF = psMemIncrRefCounter (model);
        psTrace ("psphot", 5, "fitted blend as PSF\n");

        // build cached model and subtract
        pmSourceCacheModel (blend, maskVal);
        pmSourceSub (blend, PM_MODEL_OP_FULL, maskVal);
        blend->mode |=  PM_SOURCE_MODE_BLEND_FIT;
    }
    NfitBlend += modelSet->n;

    // evaluate the primary object
    if (!psphotEvalPSF (source, PSF)) {
        psTrace ("psphot", 4, "failed on blend 0 of %ld\n", modelSet->n);
        psFree (PSF);
        psFree (modelSet);
        psFree (sourceSet);
        return false;
    }
    psFree (modelSet);
    psFree (sourceSet);

    // save the new, successful model
    psFree (source->modelPSF);
    source->modelPSF = PSF;
    psTrace ("psphot", 5, "fitted primary as PSF\n");

    // build cached model and subtract
    pmSourceCacheModel (source, maskVal);
    pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    source->mode |=  PM_SOURCE_MODE_BLEND_FIT;
    return true;
}

bool psphotFitPSF (pmReadout *readout, pmSource *source, pmPSF *psf, psImageMaskType maskVal, psImageMaskType markVal) {

    double chiTrend;

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

    NfitPSF ++;

    // save the PSF model from the Ensemble fit
    pmModel *PSF = pmModelCopy (source->modelPSF);
    if (isnan(PSF->params->data.F32[1])) psAbort("nan in psf fit");

    // extend source radius as needed
    psphotCheckRadiusPSF (readout, source, PSF, markVal);

    // fit PSF model (set/unset the pixel mask)
    pmSourceFitModel (source, PSF, PM_SOURCE_FIT_PSF, maskVal);

    if (!isfinite(PSF->params->data.F32[PM_PAR_I0])) psAbort("nan in fit");

    // clear the circular mask
    psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal)); 

    // correct model chisq for flux trend
    chiTrend = psPolynomial1DEval (psf->ChiTrend, PSF->params->data.F32[PM_PAR_I0]);
    PSF->chisqNorm = PSF->chisq / chiTrend;

    // does the PSF model succeed?
    if (!psphotEvalPSF (source, PSF)) {
        psFree (PSF);
        return false;
    }

    // free old model, save new model
    psFree (source->modelPSF);
    source->modelPSF = PSF;
    psTrace ("psphot", 5, "fitted as PSF\n");

    // build cached model and subtract
    pmSourceCacheModel (source, maskVal);
    pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    return true;
}

static pmModelType modelTypeEXT;

bool psphotInitLimitsEXT (psMetadata *recipe) {

    bool status;

    // extended source model descriptions
    char *modelNameEXT = psMetadataLookupStr (&status, recipe, "EXT_MODEL");
    modelTypeEXT = pmModelClassGetType (modelNameEXT);
    psphotInitRadiusEXT (recipe, modelTypeEXT);

    return true;
}

bool psphotFitBlob (pmReadout *readout, pmSource *source, psArray *newSources, pmPSF *psf, psImageMaskType maskVal, psImageMaskType markVal) {

    bool okEXT, okDBL;
    float chiEXT, chiDBL;
    double chiTrend;
    pmModel *ONE = NULL;
    pmSource *tmpSrc = NULL;
    pmModel *EXT = NULL;
    psArray *DBL = NULL;

    // skip the source if we don't think it is extended
    if (source->type == PM_SOURCE_TYPE_UNKNOWN) return false;
    if (source->type == PM_SOURCE_TYPE_DEFECT) return false;
    if (source->type == PM_SOURCE_TYPE_SATURATED) return false;

    // set the radius based on the footprint (also sets the mask pixels)
    float radius = psphotSetRadiusEXT (readout, source, markVal);

    // XXX note that this changes the source moments that are published...
    // recalculate the source moments using the larger extended-source moments radius
    // at this stage, skip Gaussian windowing, and do not clip pixels by S/N
    // this uses the footprint to judge both radius and aperture?
    if (!pmSourceMoments (source, radius, 0.0, 0.0, maskVal)) return false;

    psTrace ("psphot", 5, "trying blob...\n");

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

    // this temporary source is used as a place-holder by the psphotEval functions below
    tmpSrc = pmSourceAlloc ();

    // XXX need to handle failures better here
    EXT = psphotFitEXT (readout, source, modelTypeEXT, maskVal, markVal);
    if (!EXT) goto escape;
    if (!isfinite(EXT->params->data.F32[PM_PAR_I0])) goto escape;

    okEXT = psphotEvalEXT (tmpSrc, EXT);
    chiEXT = EXT ? EXT->chisq / EXT->nDOF : NAN;

    // DBL will always be defined, but DBL->data[n] might not
    DBL = psphotFitDBL (readout, source, maskVal, markVal);
    if (!DBL) goto escape;
    if (!DBL->n) goto escape;

    okDBL  = psphotEvalDBL (tmpSrc, DBL->data[0]);
    okDBL &= psphotEvalDBL (tmpSrc, DBL->data[1]);
    // XXX should I keep / save the flags set in the eval functions?

    // clear the circular mask
    psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal)); 

    // correct first model chisqs for flux trend
    chiDBL = NAN;
    ONE = DBL->data[0];
    if (ONE) {
	if (!isfinite(ONE->params->data.F32[PM_PAR_I0])) psAbort("nan in fit");
      chiTrend = psPolynomial1DEval (psf->ChiTrend, ONE->params->data.F32[1]);
      ONE->chisqNorm = ONE->chisq / chiTrend;
      chiDBL = ONE->chisq / ONE->nDOF; // save chisq for double-star/galaxy comparison
    }

    // correct second model chisqs for flux trend
    ONE = DBL->data[1];
    if (ONE) {
	if (!isfinite(ONE->params->data.F32[PM_PAR_I0])) psAbort("nan in fit");
      chiTrend = psPolynomial1DEval (psf->ChiTrend, ONE->params->data.F32[1]);
      ONE->chisqNorm = ONE->chisq / chiTrend;
    }

    psFree (tmpSrc);

    // (void) psTraceSetLevel("psModules.objects.pmSourceFitSet", 0);

    if (okEXT && okDBL) {
        psTrace ("psphot", 5, "blob chisq: %f vs %f\n", chiEXT, chiDBL);
        // XXX EAM : a bogus bias: need to examine this better
        if (3*chiEXT > chiDBL) goto keepDBL;
        goto keepEXT;
    }

    if (okEXT && !okDBL) goto keepEXT;
    if (!okEXT && okDBL) goto keepDBL;

    // both models failed; reject them both
    // XXX -- change type flags to psf in this case and keep original moments?
    psFree (EXT);
    psFree (DBL);
    return false;

keepEXT:
    // sub EXT
    psFree (DBL);

    // save new model
    source->modelEXT = EXT;
    source->modelEXT->fitRadius = radius;
    source->type = PM_SOURCE_TYPE_EXTENDED;
    source->mode |= PM_SOURCE_MODE_EXTMODEL;

    // build cached model and subtract
    pmSourceCacheModel (source, maskVal);
    pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);

# if (PS_TRACE_ON)   
    psTrace ("psphot", 5, "blob as EXT: %f %f\n", EXT->params->data.F32[PM_PAR_XPOS], EXT->params->data.F32[PM_PAR_YPOS]);
    if (psTraceGetLevel("psphot") >= 6) {
	psLogMsg ("psphot", 1, "source 2:\n");
	for (int i = 0; i < source->modelEXT->params->n; i++) {
	    psLogMsg ("psphot", 1, "PAR %d : %f +/- %f\n", i, source->modelEXT->params->data.F32[i], source->modelEXT->dparams->data.F32[i]);
	}
    }
# endif

    return true;

keepDBL:
    // sub DLB
    psFree (EXT);

    // drop old model, save new second model...
    psFree (source->modelPSF);
    source->modelPSF = psMemIncrRefCounter (DBL->data[0]);
    source->mode     |= PM_SOURCE_MODE_PAIR;
    source->modelPSF->fitRadius = radius;

    // copy most data from the primary source (modelEXT, blends stay NULL)
    pmSource *newSrc = pmSourceCopy (source);
    newSrc->modelPSF = psMemIncrRefCounter (DBL->data[1]);
    newSrc->modelPSF->fitRadius = radius;

    // build cached models and subtract
    pmSourceCacheModel (source, maskVal);
    pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    pmSourceCacheModel (newSrc, maskVal);
    pmSourceSub (newSrc, PM_MODEL_OP_FULL, maskVal);

# if (PS_TRACE_ON)   
    psTrace ("psphot", 5, "blob as DBL: %f %f\n", ONE->params->data.F32[PM_PAR_XPOS], ONE->params->data.F32[PM_PAR_YPOS]);
    if (psTraceGetLevel("psphot") >= 6) {
	psLogMsg ("psphot", 1, "source 1:\n");
	for (int i = 0; i < newSrc->modelPSF->params->n; i++) {
	    psLogMsg ("psphot", 1, "PAR %d : %f +/- %f\n", i, newSrc->modelPSF->params->data.F32[i], newSrc->modelPSF->dparams->data.F32[i]);
	}
	psLogMsg ("psphot", 1, "source 2:\n");
	for (int i = 0; i < source->modelPSF->params->n; i++) {
	    psLogMsg ("psphot", 1, "PAR %d : %f +/- %f\n", i, source->modelPSF->params->data.F32[i], source->modelPSF->dparams->data.F32[i]);
	}
	psphotVisualShowResidualImage (readout);
    }
# endif

    psArrayAdd (newSources, 100, newSrc);
    psFree (newSrc);
    psFree (DBL);
    return true;

escape:
    psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal)); 
    psFree (tmpSrc);
    psFree (EXT);
    psFree (DBL);
    return false;
}

// fit a double PSF source to an extended blob
psArray *psphotFitDBL (pmReadout *readout, pmSource *source, psImageMaskType maskVal, psImageMaskType markVal) {

    float dx, dy;
    pmModel *DBL;
    pmModel *PSF;
    psEllipseAxes axes;
    psEllipseMoments moments;
    psArray *modelSet;

    NfitDBL ++;

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

    // XXX this is really poor: if we don't have moments for the source, we have no guess.
    // force the measurement?
    psAssert (source->moments, "moments are re-calculated for any extended source");
	
    // make a guess at the position of the two sources
    moments.x2 = source->moments->Mxx;
    moments.xy = source->moments->Mxy;
    moments.y2 = source->moments->Myy;
    axes = psEllipseMomentsToAxes (moments, 20.0);

    // XXX this is really arbitrary: 4 pixel separation?
    dx = 2 * cos (axes.theta);
    dy = 2 * sin (axes.theta);

    // save the PSF model from the Ensemble fit
    PSF = source->modelPSF;
    if (isnan(PSF->params->data.F32[1])) psAbort("nan in dbl fit");

    modelSet = psArrayAlloc (2);

    DBL = pmModelCopy (PSF);
    DBL->params->data.F32[PM_PAR_I0]  *= 0.5;
    DBL->params->data.F32[PM_PAR_XPOS] = source->peak->xf + dx;
    DBL->params->data.F32[PM_PAR_YPOS] = source->peak->yf + dy;
    modelSet->data[0] = DBL;

    DBL = pmModelCopy (PSF);
    DBL->params->data.F32[PM_PAR_I0]  *= 0.5;
    DBL->params->data.F32[PM_PAR_XPOS] = source->peak->xf - dx;
    DBL->params->data.F32[PM_PAR_YPOS] = source->peak->yf - dy;
    modelSet->data[1] = DBL;

    // fit PSF model (set/unset the pixel mask)
    pmSourceFitSet (source, modelSet, PM_SOURCE_FIT_PSF, maskVal);
    return (modelSet);
}

pmModel *psphotFitEXT (pmReadout *readout, pmSource *source, pmModelType modelType, psImageMaskType maskVal, psImageMaskType markVal) {

    NfitEXT ++;

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

    // use the source moments, etc to guess basic model parameters
    pmModel *EXT = pmSourceModelGuess (source, modelType);
    if (!EXT) {
	psTrace ("psphot", 5, "failed to generate a model for source: moments: %f %f\n", source->moments->Mxx, source->moments->Myy);
	return NULL;
    }

    if ((source->moments->Mxx < 1e-3) || (source->moments->Myy < 1e-3)) {
        psTrace ("psphot", 5, "problem source: moments: %f %f\n", source->moments->Mxx, source->moments->Myy);
    }

    // fit EXT (not PSF) model (set/unset the pixel mask)
    pmSourceFitModel (source, EXT, PM_SOURCE_FIT_EXT, maskVal);
    return (EXT);
}
