# 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 () {
    psTimerStart ("psphot.fits");
    return true;
}

bool psphotFitSummary () {

    fprintf (stderr, "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, psMaskType maskVal) {

    float x, y, dR;

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

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

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

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

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

    // 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];

        // 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_SUBTRACTED;
    }
    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_SUBTRACTED;
    return true;
}

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

    double chiTrend;

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

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

    // 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);
    source->mode |=  PM_SOURCE_MODE_SUBTRACTED;
    return true;
}

static float EXT_MIN_SN;
static float EXT_MOMENTS_RAD;
static pmModelType modelTypeEXT;

bool psphotInitLimitsEXT (psMetadata *recipe) {

    bool status;

    // extended source model parameters
    EXT_MIN_SN       = psMetadataLookupF32 (&status, recipe, "EXT_MIN_SN");
    EXT_MOMENTS_RAD  = psMetadataLookupF32 (&status, recipe, "EXT_MOMENTS_RADIUS");

    // 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 *sources, pmPSF *psf, psMaskType maskVal) {

    bool okEXT, okDBL;
    float chiEXT, chiDBL;
    double chiTrend;
    pmModel *ONE = 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;
    if (source->moments->SN < EXT_MIN_SN) return false;

    // recalculate the source moments using the larger extended-source moments radius
    if (!pmSourceMoments (source, EXT_MOMENTS_RAD)) return false;

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

    // psTraceSetLevel("psModules.objects.pmSourceFitSet", 5);

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

    pmModel *EXT = psphotFitEXT (readout, source, modelTypeEXT, maskVal);
    okEXT = psphotEvalEXT (tmpSrc, EXT);
    chiEXT = EXT->chisq / EXT->nDOF;

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

    // correct first model chisqs for flux trend
    ONE = DBL->data[0];
    chiTrend = psPolynomial1DEval (psf->ChiTrend, ONE->params->data.F32[1]);
    ONE->chisqNorm = ONE->chisq / chiTrend;

    // save chisq for double-star/galaxy comparison
    chiDBL = ONE->chisq / ONE->nDOF;

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

    psFree (tmpSrc);

    // 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
    psFree (EXT);
    psFree (DBL);
    return false;

keepEXT:
    // sub EXT
    psFree (DBL);

    // save new model
    source->modelEXT = EXT;
    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);
    source->mode |=  PM_SOURCE_MODE_SUBTRACTED;
    psTrace ("psphot", 5, "blob as EXT: %f %f\n", EXT->params->data.F32[PM_PAR_XPOS], EXT->params->data.F32[PM_PAR_YPOS]);
    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_SUBTRACTED;
    source->mode    |= PM_SOURCE_MODE_PAIR;

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

    // 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);
    psTrace ("psphot", 5, "blob as DBL: %f %f\n", ONE->params->data.F32[PM_PAR_XPOS], ONE->params->data.F32[PM_PAR_YPOS]);

    psArrayAdd (sources, 100, newSrc);
    psFree (newSrc);
    psFree (DBL);
    return true;
}

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

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

    NfitDBL ++;

    // make a guess at the position of the two sources
    moments.x2 = source->moments->Sx;
    moments.y2 = source->moments->Sy;
    moments.xy = source->moments->Sxy;
    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;
    psphotCheckRadiusPSFBlend (readout, source, PSF, 8.0);
    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, psMaskType maskVal) {

    NfitEXT ++;

    // use the source moments, etc to guess basic model parameters
    pmModel *EXT = pmSourceModelGuess (source, modelType);
    PS_ASSERT (EXT, NULL);

    // if (isnan(EXT->params->data.F32[1])) psAbort("nan in ext fit");

    psphotCheckRadiusEXT (readout, source, EXT);

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

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

