# include "psphotInternal.h"

bool psphotSetMomentsWindow (psMetadata *recipe, psArray *sources);

psArray *psphotSourceStats (pmConfig *config, pmReadout *readout, pmDetections *detections, bool setWindow) {

    bool status = false;
    psArray *sources = NULL;

    psTimerStart ("psphot.stats");

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

    // determine properties (sky, moments) of initial sources
    float OUTER    = psMetadataLookupF32 (&status, recipe, "SKY_OUTER_RADIUS");
    if (!status) return NULL;

    OUTER = PS_MAX(OUTER, 20.0); // XXX Guarantee that we can encompass the max moments radius

    char *breakPt  = psMetadataLookupStr (&status, recipe, "BREAK_POINT");
    if (!status) return NULL;

    psArray *peaks = detections->peaks;
    if (!peaks) {
        psError(PS_ERR_UNEXPECTED_NULL, false, "No peaks found!");
        return NULL;
    }

    // generate the array of sources, define the associated pixel
    sources = psArrayAllocEmpty (peaks->n);

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

        pmPeak *peak = peaks->data[i];
        if (peak->assigned) continue;

        // create a new source
        pmSource *source = pmSourceAlloc();

        // add the peak
        source->peak = psMemIncrRefCounter(peak);

        // allocate space for moments
        source->moments = pmMomentsAlloc();

        // allocate image, weight, mask arrays for each peak (square of radius OUTER)
        pmSourceDefinePixels (source, readout, source->peak->x, source->peak->y, OUTER);

        peak->assigned = true;
        psArrayAdd (sources, 100, source);
        psFree (source);
    }

    if (!strcasecmp (breakPt, "PEAKS")) {
        psLogMsg ("psphot", PS_LOG_INFO, "%ld sources : %f sec\n", sources->n, psTimerMark ("psphot.stats"));
        psLogMsg ("psphot", PS_LOG_INFO, "break point PEAKS, skipping MOMENTS\n");
        psphotVisualShowMoments (sources);
        return sources;
    }

    if (setWindow) {
	if (!psphotSetMomentsWindow(recipe, sources)) {
	    psError(PS_ERR_UNEXPECTED_NULL, false, "Failed to determine Moments Window!");
	    return NULL;
	}
    }

    // threaded measurement of the source magnitudes
    int Nfail = 0;
    int Nmoments = 0;
    int Nfaint = 0;

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

            psArrayAdd(job->args, 1, cells->data[j]); // sources
            psArrayAdd(job->args, 1, recipe);
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nmoments
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfail
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfaint

            if (!psThreadJobAddPending(job)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to launch thread job PSPHOT_SOURCE_STATS");
                psFree (job);
                return NULL;
            }
            psFree(job);
        }

        // wait for the threads to finish and manage results
        if (!psThreadPoolWait (false)) {
            psError(PS_ERR_UNKNOWN, false, "Failure in thread job PSPHOT_SOURCE_STATS");
            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[2];
                Nmoments += scalar->data.S32;
                scalar = job->args->data[3];
                Nfail += scalar->data.S32;
                scalar = job->args->data[4];
                Nfaint += scalar->data.S32;
            }
            psFree(job);
        }
    }

    psFree (cellGroups);

    psLogMsg ("psphot", PS_LOG_INFO, "%ld sources, %d moments, %d faint, %d failed: %f sec\n", sources->n, Nmoments, Nfaint, Nfail, psTimerMark ("psphot.stats"));

    psphotVisualShowMoments (sources);

    return (sources);
}

bool psphotSourceStats_Threaded (psThreadJob *job) {

    bool status = false;
    float BIG_RADIUS;
    psScalar *scalar = NULL;

    psArray *sources                = job->args->data[0];
    psMetadata *recipe              = job->args->data[1];

    float INNER        = psMetadataLookupF32 (&status, recipe, "SKY_INNER_RADIUS");
    if (!status) return false;
    float MIN_SN       = psMetadataLookupF32 (&status, recipe, "MOMENTS_SN_MIN");
    if (!status) return false;
    float RADIUS       = psMetadataLookupF32 (&status, recipe, "PSF_MOMENTS_RADIUS");
    if (!status) return false;
    float SIGMA        = psMetadataLookupF32 (&status, recipe, "MOMENTS_GAUSS_SIGMA");
    if (!status) return false;

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

    // threaded measurement of the sources moments
    int Nfail = 0;
    int Nmoments = 0;
    int Nfaint = 0;
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];

        // skip faint sources for moments measurement
        if (source->peak->SN < MIN_SN) {
            source->mode |= PM_SOURCE_MODE_BELOW_MOMENTS_SN;
            Nfaint++;
            continue;
        }

        // measure a local sky value
        // the local sky is now ignored; kept here for reference only
        status = pmSourceLocalSky (source, PS_STAT_SAMPLE_MEDIAN, INNER, maskVal, markVal);
        if (!status) {
            source->mode |= PM_SOURCE_MODE_SKY_FAILURE;
            psErrorClear(); // XXX re-consider the errors raised here
            Nfail ++;
            continue;
        }

        // measure the local sky variance (needed if noise is not sqrt(signal))
        // XXX EAM : this should use ROBUST not SAMPLE median, but it is broken
        status = pmSourceLocalSkyVariance (source, PS_STAT_SAMPLE_MEDIAN, INNER, maskVal, markVal);
        if (!status) {
            source->mode |= PM_SOURCE_MODE_SKYVAR_FAILURE;
            Nfail ++;
            psErrorClear();
            continue;
        }

        // measure basic source moments (no S/N clipping on input pixels)
        status = pmSourceMoments (source, RADIUS, SIGMA, 0.0);
        if (status) {
            Nmoments ++;
            continue;
        }

        // if no valid pixels, or massive swing, likely saturated source,
        // try a much larger box
        BIG_RADIUS = PS_MIN (INNER, 3*RADIUS);
        psTrace ("psphot", 4, "retrying moments for %d, %d\n", source->peak->x, source->peak->y);
        status = pmSourceMoments (source, BIG_RADIUS, 3.0*SIGMA, 0.0);
        if (status) {
            source->mode |= PM_SOURCE_MODE_BIG_RADIUS;
            Nmoments ++;
            continue;
        }

        source->mode |= PM_SOURCE_MODE_MOMENTS_FAILURE;
        Nfail ++;
        psErrorClear();
        continue;
    }

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

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

    scalar = job->args->data[4];
    scalar->data.S32 = Nfaint;

    return true;
}

// this function attempts to iteratively determine the best value for sigma of the moments weighting Gaussian
bool psphotSetMomentsWindow (psMetadata *recipe, psArray *sources) {

    bool status;

    float MIN_SN = psMetadataLookupF32 (&status, recipe, "MOMENTS_SN_MIN");
    if (!status) return false;

    // XXX move this to a config file?
    float sigma[4] = {1.0, 2.0, 3.0, 4.5};
    float Sout[4];

    // this sorts by peak->SN
    sources = psArraySort (sources, pmSourceSortBySN);

    // loop over radii:
    for (int i = 0; i < 4; i++) {

	// XXX move max source number to config
	for (int j = 0; (j < sources->n) && (j < 400); j++) {
 
	    pmSource *source = sources->data[j];
	    psAssert (source->moments, "force moments to exist");
	    source->moments->nPixels = 0;

	    // skip faint sources for moments measurement
	    if (source->peak->SN < MIN_SN) {
		continue;
	    }

	    // measure basic source moments (no S/N clipping on input pixels)
	    status = pmSourceMoments (source, 4*sigma[i], sigma[i], 0.0);
	}

	// choose a grid scale that is a fixed fraction of the psf sigma^2
	psMetadataAddF32(recipe, PS_LIST_TAIL, "PSF_CLUMP_GRID_SCALE", PS_META_REPLACE, "clump grid", 0.1*PS_SQR(sigma[i]));
	psMetadataAddF32(recipe, PS_LIST_TAIL, "MOMENTS_SX_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(sigma[i]));
	psMetadataAddF32(recipe, PS_LIST_TAIL, "MOMENTS_SY_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(sigma[i]));

	// determine the PSF parameters from the source moment values
	pmPSFClump psfClump = pmSourcePSFClump (NULL, sources, recipe);
	psLogMsg ("psphot", 3, "radius %.1f, nStars: %d, nSigma: %5.2f, X,  Y: %f, %f (%f, %f)\n", sigma[i], psfClump.nStars, psfClump.nSigma, psfClump.X, psfClump.Y, sqrt(psfClump.X) / sigma[i], sqrt(psfClump.Y) / sigma[i]);

	psMetadataAddS32 (recipe, PS_LIST_TAIL, "PSF.CLUMP.NREGIONS",  PS_META_REPLACE, "psf clump regions", 1);
	psMetadata *regionMD = psMetadataLookupPtr (&status, recipe, "PSF.CLUMP.REGION.000");
	if (!regionMD) {
	    regionMD = psMetadataAlloc();
	    psMetadataAddMetadata (recipe, PS_LIST_TAIL, "PSF.CLUMP.REGION.000", PS_META_REPLACE, "psf clump region", regionMD);
	    psFree (regionMD);
	}
	psMetadataAddF32 (regionMD, PS_LIST_TAIL, "PSF.CLUMP.X",  PS_META_REPLACE, "psf clump center", psfClump.X);
	psMetadataAddF32 (regionMD, PS_LIST_TAIL, "PSF.CLUMP.Y",  PS_META_REPLACE, "psf clump center", psfClump.Y);
	psMetadataAddF32 (regionMD, PS_LIST_TAIL, "PSF.CLUMP.DX", PS_META_REPLACE, "psf clump center", psfClump.dX);
	psMetadataAddF32 (regionMD, PS_LIST_TAIL, "PSF.CLUMP.DY", PS_META_REPLACE, "psf clump center", psfClump.dY);
	    
	// psphotVisualPlotMoments (recipe, sources);

	Sout[i] = sqrt(0.5*(psfClump.X + psfClump.Y)) / sigma[i];
    }

    // we are looking for sigma for which Sout = 0.65 (or some other value)

    float Sigma = NAN;
    float minS = Sout[0];
    float maxS = Sout[0];
    for (int i = 0; i < 4; i++) {
	minS = PS_MIN(Sout[i], minS);
	maxS = PS_MAX(Sout[i], maxS);
    }
    if (minS > 0.65) Sigma = sigma[3];
    if (maxS < 0.65) Sigma = sigma[0];

    for (int i = 0; i < 3; i++) {
	if ((Sout[i] > 0.65) && (Sout[i+1] > 0.65)) continue;
	if ((Sout[i] < 0.65) && (Sout[i+1] < 0.65)) continue;
	Sigma = sigma[i] + (0.65 - Sout[i])*(sigma[i+1] - sigma[i])/(Sout[i+1] - Sout[i]);
    }
    psAssert (isfinite(Sigma), "did we miss a case?");
	
    // choose a grid scale that is a fixed fraction of the psf sigma^2
    psMetadataAddF32(recipe, PS_LIST_TAIL, "PSF_CLUMP_GRID_SCALE", PS_META_REPLACE, "clump grid", 0.1*PS_SQR(Sigma));
    psMetadataAddF32(recipe, PS_LIST_TAIL, "MOMENTS_SX_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(Sigma));
    psMetadataAddF32(recipe, PS_LIST_TAIL, "MOMENTS_SY_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(Sigma));
    psMetadataAddF32(recipe, PS_LIST_TAIL, "MOMENTS_GAUSS_SIGMA", PS_META_REPLACE, "moments limit", Sigma);
    psMetadataAddF32(recipe, PS_LIST_TAIL, "PSF_MOMENTS_RADIUS", PS_META_REPLACE, "moments limit", 4.0*Sigma);

    psLogMsg ("psphot", 3, "using window function with sigma = %f\n", Sigma);
    return true;
}
