# include "psphotInternal.h"

// convert detections to sources and measure their basic properties (moments, local sky, sky
// variance) Note: this function only generates sources for the new peaks (peak->assigned).
// The new sources are added to any existing sources on detections->newSources.  The sources 
// on detections->allSources are ignored.
bool psphotSourceStats (pmConfig *config, const pmFPAview *view, bool setWindow)
{
    bool status = true;

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

    int num = psMetadataLookupS32 (&status, config->arguments, "PSPHOT.INPUT.NUM");
    psAssert (status, "programming error: must define PSPHOT.INPUT.NUM");

    // loop over the available readouts
    for (int i = 0; i < num; i++) {
        if (!psphotSourceStatsReadout (config, view, "PSPHOT.INPUT", i, recipe, setWindow)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed to find initial detections for PSPHOT.INPUT entry %d", i);
            return false;
        }
    }
    return true;
}

bool psphotSourceStatsReadout (pmConfig *config, const pmFPAview *view, const char *filename, int index, psMetadata *recipe, bool setWindow) {

    bool status = false;
    psArray *sources = NULL;

    psTimerStart ("psphot.stats");

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, filename, 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?");

    // XXX TEST:
    if (detections->allSources) {
        psphotMaskCosmicRayFootprintCheck(detections->allSources);
    }
    if (detections->newSources) {
        psphotMaskCosmicRayFootprintCheck(detections->newSources);
    }

    // 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");
    psAssert (status, "missing SKY_OUTER_RADIUS in recipe?");

    // XXX this seems like an arbitrary number...
    OUTER = PS_MAX(OUTER, 20.0); // XXX Guarantee that we can encompass the max moments radius

    char *breakPt  = psMetadataLookupStr (&status, recipe, "BREAK_POINT");
    psAssert (status, "missing BREAK_POINT?");

    float INNER        = psMetadataLookupF32 (&status, recipe, "SKY_INNER_RADIUS"); psAssert (status, "missing SKY_INNER_RADIUS");
    float MIN_SN       = psMetadataLookupF32 (&status, recipe, "MOMENTS_SN_MIN"); psAssert (status, "missing MOMENTS_SN_MIN");

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

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

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

    // generate the array of sources, define the associated pixel
    if (!detections->newSources) {
	detections->newSources = psArrayAllocEmpty (peaks->n);
    }
    sources = detections->newSources;

    // if there are no peaks, we save the empty source array and return
    if (!peaks->n) {
        return true;
    }

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

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

        // create a new source
        pmSource *source = pmSourceAlloc();
	source->imageID = index;

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

    if (setWindow) {
        if (!psphotSetMomentsWindow(recipe, readout->analysis, sources)) {
            psError(PS_ERR_UNEXPECTED_NULL, false, "Failed to determine Moments Window!");
            psFree(detections->newSources);
            return false;
        }
    }

    // if we have measured the window, we will be saving the modified version of these recipe values on readout->analysis
    float SIGMA = psMetadataLookupF32 (&status, readout->analysis, "MOMENTS_GAUSS_SIGMA");
    if (!status) {
        SIGMA = psMetadataLookupF32 (&status, recipe, "MOMENTS_GAUSS_SIGMA");
    }
    float RADIUS = psMetadataLookupF32 (&status, readout->analysis, "PSF_MOMENTS_RADIUS");
    if (!status) {
        RADIUS = psMetadataLookupF32 (&status, recipe, "PSF_MOMENTS_RADIUS");
    }

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

            PS_ARRAY_ADD_SCALAR(job->args, INNER,   PS_TYPE_F32);
            PS_ARRAY_ADD_SCALAR(job->args, MIN_SN,  PS_TYPE_F32);
            PS_ARRAY_ADD_SCALAR(job->args, RADIUS,  PS_TYPE_F32);
            PS_ARRAY_ADD_SCALAR(job->args, SIGMA,   PS_TYPE_F32);

            PS_ARRAY_ADD_SCALAR(job->args, maskVal, PS_TYPE_IMAGE_MASK);
            PS_ARRAY_ADD_SCALAR(job->args, markVal, PS_TYPE_IMAGE_MASK);

            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);
                psFree(detections->newSources);
                return false;
            }
            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");
	    psFree(detections->newSources);
            return false;
        }

        // 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[7];
                Nmoments += scalar->data.S32;
                scalar = job->args->data[8];
                Nfail += scalar->data.S32;
                scalar = job->args->data[9];
                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);

    if (detections->allSources) {
        psphotMaskCosmicRayFootprintCheck(detections->allSources);
    }
    if (detections->newSources) {
        psphotMaskCosmicRayFootprintCheck(detections->newSources);
    }

    return true;
}

// this function is currently only called by psphotCheckExtSources
bool psphotSourceStatsUpdate (psArray *sources, pmConfig *config, pmReadout *readout) {

    bool status = false;

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

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

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

        pmSource *source = sources->data[i];
        if (!source->peak) continue; // XXX how can we have a peak-less source?

        // allocate space for moments
        if (!source->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);
        source->peak->assigned = true;
    }

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

    // XXX how else could we get the window size in?
    if (!psphotSetMomentsWindow(recipe, readout->analysis, 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");

            // XXX: this must match the above
            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];

    float INNER                     = PS_SCALAR_VALUE(job->args->data[1],F32);
    float MIN_SN                    = PS_SCALAR_VALUE(job->args->data[2],F32);
    float RADIUS                    = PS_SCALAR_VALUE(job->args->data[3],F32);
    float SIGMA                     = PS_SCALAR_VALUE(job->args->data[4],F32);

    psImageMaskType maskVal         = PS_SCALAR_VALUE(job->args->data[5],PS_TYPE_IMAGE_MASK_DATA);
    psImageMaskType markVal         = PS_SCALAR_VALUE(job->args->data[6],PS_TYPE_IMAGE_MASK_DATA);

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

	if (source->tmpFlags & PM_SOURCE_TMPF_MOMENTS_MEASURED) continue;
	source->tmpFlags |= PM_SOURCE_TMPF_MOMENTS_MEASURED;

        // 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, maskVal);
        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, maskVal);
        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[7];
    scalar->data.S32 = Nmoments;

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

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

    return true;
}

// this function attempts to iteratively determine the best value for sigma of the moments weighting Gaussian
// this function modifies the recipe values MOMENTS_SX_MAX, MOMENTS_SY_MAX, and PSF_CLUMP_GRID_SCALE, used by pmSourcePSFClump
bool psphotSetMomentsWindow (psMetadata *recipe, psMetadata *analysis, psArray *sources) {

    bool status;

    float MIN_SN = psMetadataLookupF32 (&status, recipe, "MOMENTS_SN_MIN"); psAssert (status, "missing MOMENTS_SN_MIN");
    float PSF_SN_LIM = psMetadataLookupF32(&status, recipe, "PSF_SN_LIM"); psAssert (status, "missing PSF_SN_LIM");
    psF32 MOMENTS_AR_MAX = psMetadataLookupF32(&status, recipe, "MOMENTS_AR_MAX"); psAssert (status, "missing MOMENTS_AR_MAX");

    // XXX move this to a config file?
    # define NSIGMA 8
    float sigma[NSIGMA] = {1.0, 2.0, 3.0, 4.5, 6.0, 9.0, 12.0, 18.0};
    float Sout[NSIGMA];

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

    // loop over radii:
    for (int i = 0; i < NSIGMA; 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, 0xffff);
        }

        // choose a grid scale that is a fixed fraction of the psf sigma^2
        float PSF_CLUMP_GRID_SCALE = 0.1*PS_SQR(sigma[i]);
        float MOMENTS_SX_MAX = 2.0*PS_SQR(sigma[i]);
        float MOMENTS_SY_MAX = 2.0*PS_SQR(sigma[i]);

        // determine the PSF parameters from the source moment values
        pmPSFClump psfClump = pmSourcePSFClump (NULL, NULL, sources, PSF_SN_LIM, PSF_CLUMP_GRID_SCALE, MOMENTS_SX_MAX, MOMENTS_SY_MAX, MOMENTS_AR_MAX);
        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 (analysis, PS_LIST_TAIL, "PSF.CLUMP.NREGIONS",  PS_META_REPLACE, "psf clump regions", 1);
        psMetadata *regionMD = psMetadataLookupPtr (&status, analysis, "PSF.CLUMP.REGION.000");
        if (!regionMD) {
            regionMD = psMetadataAlloc();
            psMetadataAddMetadata (analysis, 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, analysis, 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 < NSIGMA; i++) {
        minS = PS_MIN(Sout[i], minS);
        maxS = PS_MAX(Sout[i], maxS);
    }
    if (minS > 0.65) Sigma = sigma[NSIGMA-1];
    if (maxS < 0.65) Sigma = sigma[0];

    for (int i = 0; i < NSIGMA - 1 && isnan(Sigma); i++) {
        if (!isfinite(Sout[i]) || !isfinite(Sout[i+1])) continue;
        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(analysis, PS_LIST_TAIL, "PSF_CLUMP_GRID_SCALE", PS_META_REPLACE, "clump grid", 0.1*PS_SQR(Sigma));
    psMetadataAddF32(analysis, PS_LIST_TAIL, "MOMENTS_SX_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(Sigma));
    psMetadataAddF32(analysis, PS_LIST_TAIL, "MOMENTS_SY_MAX", PS_META_REPLACE, "moments limit", 2.0*PS_SQR(Sigma));
    psMetadataAddF32(analysis, PS_LIST_TAIL, "MOMENTS_GAUSS_SIGMA", PS_META_REPLACE, "moments limit", Sigma);
    psMetadataAddF32(analysis, 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;
}

// if we use the footprints, the output peaks list contains both old and new peaks,
// otherwise it only contains the new peaks.

