# include "psphotInternal.h"

// for now, let's store the detections on the readout->analysis for each readout
bool psphotBlendFit (pmConfig *config, const pmFPAview *view, const char *filerule)
{
    bool status = true;

    fprintf (stdout, "\n");
    psLogMsg ("psphot", PS_LOG_INFO, "--- psphot Fit Sources (Non-Linear) ---");

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

    int num = psphotFileruleCount(config, filerule);

    // loop over the available readouts
    for (int i = 0; i < num; i++) {
        if (!psphotBlendFitReadout (config, view, filerule, i, recipe)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed to fit sources (non-linear) for %s entry %d", filerule, i);
            return false;
        }
    }
    return true;
}

// XXX I don't like this name
bool psphotBlendFitReadout (pmConfig *config, const pmFPAview *view, const char *filerule, int index, psMetadata *recipe) {

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

    psTimerStart ("psphot.fit.nonlinear");

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, filerule, index); // File of interest
    psAssert (file, "missing file?");

    pmReadout *readout = pmFPAviewThisReadout(view, file->fpa);
    psAssert (readout, "missing readout?");

    pmDetections *detections = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.DETECTIONS");
    psAssert (detections, "missing detections?");

    psArray *sources = detections->allSources;
    psAssert (sources, "missing sources?");

    if (!sources->n) {
        psLogMsg ("psphot", PS_LOG_INFO, "no sources, skipping blend fit");
        return true;
    }

    pmPSF *psf = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.PSF");
    psAssert (psf, "missing psf?");

    // 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 fitMinTol = psMetadataLookupF32 (&status, recipe, "EXT_FIT_MIN_TOL"); // Fit tolerance
    if (!status || !isfinite(fitMinTol) || fitMinTol <= 0) {
	fitMinTol = psMetadataLookupF32 (&status, recipe, "PSF_FIT_TOL"); // Fit tolerance
	if (!status || !isfinite(fitMinTol) || fitMinTol <= 0) {
	    psAbort("PSF_FIT_MIN_TOL (and PSF_FIT_TOL) not defined or positive");
	}
    }

    float fitMaxTol = psMetadataLookupF32 (&status, recipe, "EXT_FIT_MAX_TOL"); // Fit tolerance
    if (!status || !isfinite(fitMaxTol) || fitMaxTol <= 0) {
	fitMaxTol = 1.0;
    }

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

    float maxChisqDOF = psMetadataLookupF32 (&status, recipe, "EXT_FIT_MAX_CHISQ"); // Fit tolerance

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

    // Define source fitting parameters for extended source fits
    pmSourceFitOptions *fitOptions = pmSourceFitOptionsAlloc();
    fitOptions->nIter         = fitIter;
    fitOptions->minTol        = fitMinTol;
    fitOptions->maxTol        = fitMaxTol;
    fitOptions->maxChisqDOF   = maxChisqDOF;
    fitOptions->poissonErrors = poisson;
    fitOptions->weight        = PS_SQR(skySig);
    fitOptions->mode          = PM_SOURCE_FIT_PSF;
    fitOptions->covarFactor   = psImageCovarianceFactorForAperture(readout->covariance, 10.0); // Covariance matrix

    psphotInitLimitsPSF (recipe, readout);
    psphotInitLimitsEXT (recipe, readout);
    psphotInitRadiusPSF (recipe, readout);

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

    // source analysis is done in S/N order (brightest first)
    sources = psArraySort (sources, pmSourceSortByFlux);
    if (!sources->n) {
        psLogMsg ("psphot", PS_LOG_INFO, "no sources, skipping blend");
	psFree (fitOptions);
        return true;
    }

    // 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
            psArrayAdd(job->args, 1, fitOptions); // default fit options
            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 (fitOptions);
                return NULL;
            }
        }

        // wait for the threads to finish and manage results
        if (!psThreadPoolWait (false, true)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
	    psFree (fitOptions);
            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[6];
                Nfit += scalar->data.S32;
                scalar = job->args->data[7];
                Npsf += scalar->data.S32;
                scalar = job->args->data[8];
                Next += scalar->data.S32;
                scalar = job->args->data[9];
                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");
    }
    psFree (fitOptions);

    psLogMsg ("psphot.psphotBlendFit", PS_LOG_WARN, "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, false);
    psphotVisualShowObjectRegions (readout, recipe, sources);
    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];
    pmSourceFitOptions *fitOptions = job->args->data[5];

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

# if (0)
# define TEST_X 34
# define TEST_Y 28
   
	if ((fabs(source->peak->xf - TEST_X) < 5) && (fabs(source->peak->yf - TEST_Y) < 5)) {
	    fprintf (stderr, "test object\n");
	}

# undef TEST_X
# undef TEST_Y
# endif

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

	// do not include MOMENTS_FAILURES in the fit
        if (source->mode & PM_SOURCE_MODE_MOMENTS_FAILURE) continue;

        // limit selection to some SN limit
        if (source->mode & PM_SOURCE_MODE_EXT_LIMIT) {
	    if (source->moments->KronFlux < FIT_SN_LIM * source->moments->KronFluxErr) continue;
	} else {
	    if (sqrt(source->peak->detValue) < 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 & remove excess noise
        if (source->tmpFlags & PM_SOURCE_TMPF_SUBTRACTED) {
            pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
        }
        Nfit ++;

	if (0) {
	    psF32 Mxx = source->moments->Mxx;
	    psF32 Myy = source->moments->Myy;
	    fprintf (stderr, "1: Mxx: %f, Myy: %f\n", Mxx, Myy);
	}

        // 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, fitOptions, maskVal, markVal)) {
                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, fitOptions, 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;
            }
        }

	if (0) {
	    psF32 Mxx = source->moments->Mxx;
	    psF32 Myy = source->moments->Myy;
	    fprintf (stderr, "2: Mxx: %f, Myy: %f\n", Mxx, Myy);
	}

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

        // re-subtract the object, leave local sky, re-bump noise
        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[6];
    scalar->data.S32 = Nfit;

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

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

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

    return true;
}

