# include "psphotInternal.h"

// XXX I don't like this name
bool psphotBlendFit (pmConfig *config, pmReadout *readout, psArray *sources, pmPSF *psf) {

    int Nfit = 0;
    int Npsf = 0;
    int Next = 0;
    int Nfail = 0;
    bool status;

    psTimerStart ("psphot.fit.nonlinear");

    // select the appropriate recipe information
    psMetadata *recipe  = psMetadataLookupPtr (&status, config->recipes, PSPHOT_RECIPE);
    assert (recipe);

    // determine the number of allowed threads
    int nThreads = psMetadataLookupS32(&status, config->arguments, "NTHREADS"); // Number of threads
    if (!status) {
        nThreads = 0;
    }

    // source fitting parameters for extended source fits
    int fitIter = psMetadataLookupS32(&status, recipe, "EXT_FIT_ITER"); // Max number of fit iterations
    assert (status && fitIter > 0);

    float fitTol = psMetadataLookupF32 (&status, recipe, "EXT_FIT_TOL"); // Fit tolerance
    assert (status && isfinite(fitTol) && fitTol > 0);

    bool poisson = psMetadataLookupBool(&status, recipe, "POISSON.ERRORS.PHOT.LMM"); // Poisson errors?
    assert (status);

    float skySig = psMetadataLookupF32(&status, recipe, "SKY_SIG");
    assert (status && isfinite(skySig) && skySig > 0);

    // Define source fitting parameters for extended source fits
    pmSourceFitModelInit(fitIter, fitTol, PS_SQR(skySig), poisson);

    psphotInitLimitsPSF (recipe, readout);
    psphotInitLimitsEXT (recipe);
    psphotInitRadiusPSF (recipe, psf->type);

    // starts the timer, sets up the array of fitSets
    psphotFitInit (nThreads);

    // source analysis is done in S/N order (brightest first)
    sources = psArraySort (sources, pmSourceSortBySN);

    // choose Cx, Cy (see psphotThreadTools.c for overview of the concepts)
    int Cx = 1, Cy = 1;
    psphotChooseCellSizes (&Cx, &Cy, readout, nThreads);

    psArray *cellGroups = psphotAssignSources (Cx, Cy, sources);

    for (int i = 0; i < cellGroups->n; i++) {

        psArray *cells = cellGroups->data[i];

        for (int j = 0; j < cells->n; j++) {

	    // allocate a job -- if threads are not defined, this just runs the job
	    psThreadJob *job = psThreadJobAlloc ("PSPHOT_BLEND_FIT");
	    psArray *newSources = psArrayAllocEmpty(16);

	    psArrayAdd(job->args, 1, readout);
	    psArrayAdd(job->args, 1, recipe);
	    psArrayAdd(job->args, 1, cells->data[j]); // sources
	    psArrayAdd(job->args, 1, psf);
	    psArrayAdd(job->args, 1, newSources); // return for new sources
	    psFree (newSources);

	    PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfit
	    PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Npsf
	    PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Next
	    PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfail

	    if (!psThreadJobAddPending(job)) {
		psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		psFree (job);
		return NULL;
	    }
	    psFree(job);

# if (0)
            {
                int nfit = 0;
                int npsf = 0;
                int next = 0;
                int nfail = 0;
                psArray *newSources = psArrayAllocEmpty(16);

                if (!psphotBlendFit_Unthreaded (&nfit, &npsf, &next, &nfail, readout, recipe, cells->data[j], psf, newSources)) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
                    return NULL;
                }
                Nfit += nfit;
                Npsf += npsf;
                Next += next;
                Nfail += nfail;

                // add these back onto sources
                for (int k = 0; k < newSources->n; k++) {
                    psArrayAdd (sources, 16, newSources->data[k]);
                }
                psFree (newSources);
            }
# endif
        }

	// wait for the threads to finish and manage results
	if (!psThreadPoolWait (false)) {
	    psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
	    return NULL;
	}

	// we have only supplied one type of job, so we can assume the types here
	psThreadJob *job = NULL;
	while ((job = psThreadJobGetDone()) != NULL) {
	    if (job->args->n < 1) {
		fprintf (stderr, "error with job\n");
	    } else {
		psScalar *scalar = NULL;
		scalar = job->args->data[5];
		Nfit += scalar->data.S32;
		scalar = job->args->data[6];
		Npsf += scalar->data.S32;
		scalar = job->args->data[7];
		Next += scalar->data.S32;
		scalar = job->args->data[8];
		Nfail += scalar->data.S32;

		// add these back onto sources
		psArray *newSources = job->args->data[4];
		for (int j = 0; j < newSources->n; j++) {
		    psArrayAdd (sources, 16, newSources->data[j]);
		}
	    }
	    psFree(job);
            }
    }
    psFree (cellGroups);

    if (psTraceGetLevel("psphot") >= 6) {
      psphotSaveImage (NULL, readout->image,  "image.v2.fits");
    }

    psLogMsg ("psphot.psphotBlendFit", PS_LOG_INFO, "fit models: %f sec for %d objects (%d psf, %d ext, %d failed, %ld skipped)\n", psTimerMark ("psphot.fit.nonlinear"), Nfit, Npsf, Next, Nfail, sources->n - Nfit);

    psphotVisualShowResidualImage (readout);
    psphotVisualShowFlags (sources);

    return true;
}

bool psphotBlendFit_Threaded (psThreadJob *job) {

    bool status = false;
    int Nfit = 0;
    int Npsf = 0;
    int Next = 0;
    int Nfail = 0;
    psScalar *scalar = NULL;

    pmReadout *readout  = job->args->data[0];
    psMetadata *recipe  = job->args->data[1];
    psArray *sources    = job->args->data[2];
    pmPSF *psf          = job->args->data[3];
    psArray *newSources = job->args->data[4];

    // bit-masks to test for good/bad pixels
    psImageMaskType maskVal = psMetadataLookupImageMask(&status, recipe, "MASK.PSPHOT");
    assert (maskVal);

    // bit-mask to mark pixels not used in analysis
    psImageMaskType markVal = psMetadataLookupImageMask(&status, recipe, "MARK.PSPHOT");
    assert (markVal);

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

    // S/N limit to perform full non-linear fits
    float FIT_SN_LIM = psMetadataLookupF32 (&status, recipe, "FULL_FIT_SN_LIM");

    // option to limit analysis to a specific region
    char *region = psMetadataLookupStr (&status, recipe, "ANALYSIS_REGION");
    psRegion AnalysisRegion = psRegionForImage (readout->image, psRegionFromString (region));
    if (psRegionIsNaN (AnalysisRegion)) psAbort("analysis region mis-defined");

    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];

        // skip non-astronomical objects (very likely defects)
        if (source->mode &  PM_SOURCE_MODE_BLEND) continue;
        if (source->mode &  PM_SOURCE_MODE_CR_LIMIT) continue;
        if (source->type == PM_SOURCE_TYPE_DEFECT) continue;
        if (source->type == PM_SOURCE_TYPE_SATURATED) continue;

        // skip DBL second sources (ie, added by psphotFitBlob
        if (source->mode &  PM_SOURCE_MODE_PAIR) continue;

        // limit selection to some SN limit
        if (source->peak->SN < FIT_SN_LIM) continue;

        // exclude sources outside optional analysis region
        if (source->peak->xf < AnalysisRegion.x0) continue;
        if (source->peak->yf < AnalysisRegion.y0) continue;
        if (source->peak->xf > AnalysisRegion.x1) continue;
        if (source->peak->yf > AnalysisRegion.y1) continue;

        // if model is NULL, we don't have a starting guess
        if (source->modelPSF == NULL) continue;

        // skip sources which are insignificant flux?
        // XXX this is somewhat ad-hoc
        if (source->modelPSF->params->data.F32[1] < 0.1) {
            psTrace ("psphot", 5, "skipping near-zero source: %f, %f : %f\n",
                     source->modelPSF->params->data.F32[1],
                     source->modelPSF->params->data.F32[2],
                     source->modelPSF->params->data.F32[3]);
            continue;
        }

        // replace object in image
        if (source->tmpFlags & PM_SOURCE_TMPF_SUBTRACTED) {
            pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
        }
        Nfit ++;

        // try fitting PSFs or extended sources depending on source->mode
        // these functions subtract the resulting fitted source
        if (source->mode & PM_SOURCE_MODE_EXT_LIMIT) {
            if (psphotFitBlob (readout, source, newSources, psf, maskVal, markVal)) {
                source->type = PM_SOURCE_TYPE_EXTENDED;
                psTrace ("psphot", 5, "source at %7.1f, %7.1f is ext", source->peak->xf, source->peak->yf);
                Next ++;
		source->mode |= PM_SOURCE_MODE_NONLINEAR_FIT;
                continue;
            }
        } else {
            if (psphotFitBlend (readout, source, psf, maskVal, markVal)) {
                source->type = PM_SOURCE_TYPE_STAR;
                psTrace ("psphot", 5, "source at %7.1f, %7.1f is psf", source->peak->xf, source->peak->yf);
                Npsf ++;
		source->mode |= PM_SOURCE_MODE_NONLINEAR_FIT;
                continue;
            }
        }

        psTrace ("psphot", 5, "source at %7.1f, %7.1f failed", source->peak->xf, source->peak->yf);
        Nfail ++;

        // re-subtract the object, leave local sky
        pmSourceCacheModel (source, maskVal);
        pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    }

    // change the value of a scalar on the array (wrap this and put it in psArray.h)
    scalar = job->args->data[5];
    scalar->data.S32 = Nfit;

    scalar = job->args->data[6];
    scalar->data.S32 = Npsf;

    scalar = job->args->data[7];
    scalar->data.S32 = Next;

    scalar = job->args->data[8];
    scalar->data.S32 = Nfail;

    return true;
}

# if (0)
bool psphotBlendFit_Unthreaded (int *nfit, int *npsf, int *next, int *nfail, pmReadout *readout, psMetadata *recipe, psArray *sources, pmPSF *psf, psArray *newSources) {

    bool status = false;
    int Nfit = 0;
    int Npsf = 0;
    int Next = 0;
    int Nfail = 0;

    // bit-masks to test for good/bad pixels
    psImageMaskType maskVal = psMetadataLookupImageMask(&status, recipe, "MASK.PSPHOT");
    assert (maskVal);

    // bit-mask to mark pixels not used in analysis
    psImageMaskType markVal = psMetadataLookupImageMask(&status, recipe, "MARK.PSPHOT");
    assert (markVal);

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

    // S/N limit to perform full non-linear fits
    float FIT_SN_LIM = psMetadataLookupF32 (&status, recipe, "FULL_FIT_SN_LIM");

    // option to limit analysis to a specific region
    char *region = psMetadataLookupStr (&status, recipe, "ANALYSIS_REGION");
    psRegion AnalysisRegion = psRegionForImage (readout->image, psRegionFromString (region));
    if (psRegionIsNaN (AnalysisRegion)) psAbort("analysis region mis-defined");

    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];

        // skip non-astronomical objects (very likely defects)
        if (source->mode &  PM_SOURCE_MODE_BLEND) continue;
        if (source->mode &  PM_SOURCE_MODE_CR_LIMIT) continue;
        if (source->type == PM_SOURCE_TYPE_DEFECT) continue;
        if (source->type == PM_SOURCE_TYPE_SATURATED) continue;

        // skip DBL second sources (ie, added by psphotFitBlob)
        if (source->mode &  PM_SOURCE_MODE_PAIR) continue;

        // limit selection to some SN limit
        if (source->peak->SN < FIT_SN_LIM) continue;

        // exclude sources outside optional analysis region
        if (source->peak->xf < AnalysisRegion.x0) continue;
        if (source->peak->yf < AnalysisRegion.y0) continue;
        if (source->peak->xf > AnalysisRegion.x1) continue;
        if (source->peak->yf > AnalysisRegion.y1) continue;

        // if model is NULL, we don't have a starting guess
        if (source->modelPSF == NULL) continue;

        // skip sources which are insignificant flux?
        // XXX this is somewhat ad-hoc
        if (source->modelPSF->params->data.F32[1] < 0.1) {
            psTrace ("psphot", 5, "skipping near-zero source: %f, %f : %f\n",
                     source->modelPSF->params->data.F32[1],
                     source->modelPSF->params->data.F32[2],
                     source->modelPSF->params->data.F32[3]);
            continue;
        }

        // replace object in image
        if (source->tmpFlags & PM_SOURCE_TMPF_SUBTRACTED) {
            pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
        }
        Nfit ++;

        // try fitting PSFs or extended sources depending on source->mode
        // these functions subtract the resulting fitted source
        if (source->mode & PM_SOURCE_MODE_EXT_LIMIT) {
            if (psphotFitBlob (readout, source, newSources, psf, maskVal, markVal)) {
                source->type = PM_SOURCE_TYPE_EXTENDED;
                psTrace ("psphot", 5, "source at %7.1f, %7.1f is ext", source->peak->xf, source->peak->yf);
                Next ++;
                continue;
            }
        } else {
            if (psphotFitBlend (readout, source, psf, maskVal, markVal)) {
                source->type = PM_SOURCE_TYPE_STAR;
                psTrace ("psphot", 5, "source at %7.1f, %7.1f is psf", source->peak->xf, source->peak->yf);
                Npsf ++;
                continue;
            }
        }

        psTrace ("psphot", 5, "source at %7.1f, %7.1f failed", source->peak->xf, source->peak->yf);
        Nfail ++;

        // re-subtract the object, leave local sky
        pmSourceCacheModel (source, maskVal);
        pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    }

    // change the value of a scalar on the array (wrap this and put it in psArray.h)
    *nfit  = Nfit;
    *npsf  = Npsf;
    *next  = Next;
    *nfail = Nfail;

    // moments are modified by the fit; re-display
    psphotVisualPlotMoments (recipe, sources);
    psphotVisualShowResidualImage (readout);

    return true;
}
# endif
