# include "psphotInternal.h"

// XXX : the threading here is not great.  this may be due to blocks between elements, but
// the selection of the objects in a cell is not optimal.  To fix:
// 1) define the boundaries of the cells up front 
// 2) loop over the sources once and associate them with their cell
// 3) define the threaded function to work with sources for a given cell

// construct an initial PSF model for each object
bool psphotGuessModels (pmConfig *config, pmReadout *readout, psArray *sources, pmPSF *psf) {

    bool status;

    psTimerStart ("psphot.models");

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

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

    // setup the PSF fit radius details
    psphotInitRadiusPSF (recipe, psf->type);

    // 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_GUESS_MODEL");
	    psArrayAdd(job->args, 1, readout);
	    psArrayAdd(job->args, 1, cells->data[j]); // sources
	    psArrayAdd(job->args, 1, psf);

	    // XXX change these to use abstract mask type info
	    PS_ARRAY_ADD_SCALAR(job->args, maskVal,  PS_TYPE_IMAGE_MASK);
	    PS_ARRAY_ADD_SCALAR(job->args, markVal,  PS_TYPE_IMAGE_MASK);

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

	// wait for the threads to finish and manage results
	// wait here for the threaded jobs to finish
	// fprintf (stderr, "wait for threads (%d, %d)\n", jx, jy);
	if (!psThreadPoolWait (false)) {
	    psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
	    return false;
	}

	// we have only supplied one type of job, so we can assume the types here
	psThreadJob *job = NULL;
	while ((job = psThreadJobGetDone()) != NULL) {
	    // we have no returned data from this operation
	    if (job->args->n < 1) {
		fprintf (stderr, "error with job\n");
	    }
	    psFree(job);
	}
    }

    // XXX check that all sources that should have models
    int nMiss = 0;
    for (int i = 0; i < sources->n; i++) {
	pmSource *source = sources->data[i];
	if (source->tmpFlags & PM_SOURCE_TMPF_MODEL_GUESS) {
	    continue;
	}
	nMiss ++;
    }
    psAssert (nMiss == 0, "failed to attempt to build models for %d objects\n", nMiss);

    psFree (cellGroups);

    psLogMsg ("psphot.models", 4, "built models for %ld objects: %f sec\n", sources->n, psTimerMark ("psphot.models"));
    return true;
}

// construct models only for sources in the specified region
bool psphotGuessModel_Threaded (psThreadJob *job) {

    pmReadout *readout = job->args->data[0];
    psArray *sources   = job->args->data[1];
    pmPSF *psf         = job->args->data[2];
    
    psImageMaskType maskVal = PS_SCALAR_VALUE(job->args->data[3],PS_TYPE_IMAGE_MASK_DATA);
    psImageMaskType markVal = PS_SCALAR_VALUE(job->args->data[4],PS_TYPE_IMAGE_MASK_DATA);

    int nSrc = 0;

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

	// this is used to mark sources for which the model is measured. We check later that
	// all are used.
	source->tmpFlags |= PM_SOURCE_TMPF_MODEL_GUESS;

	// skip non-astronomical objects (very likely defects)
	if (source->type == PM_SOURCE_TYPE_DEFECT) continue;
	if (source->type == PM_SOURCE_TYPE_SATURATED) continue;
	if (!source->peak) continue;

	nSrc ++;
	
	// the guess central intensity comes from the peak:
	float Io = source->peak->flux;

	// We have two options to get a guess for the object position: the position from the
	// peak and the position from the moments.  Use the peak position if (a) there are no
	// moments and (b) the sources is not saturated

	bool useMoments = false;
	useMoments = (source->mode & PM_SOURCE_MODE_SATSTAR);  // we only want to try if SATSTAR is set, but..
	useMoments = (useMoments && source->moments);	       // can't if there are no moments
	useMoments = (useMoments && source->moments->nPixels); // can't if the moments were not measured
	useMoments = (useMoments && !(source->mode && PM_SOURCE_MODE_MOMENTS_FAILURE)); // can't if the moments failed...

	float Xo, Yo;
	if (useMoments) {
	    Xo = source->moments->Mx;
	    Yo = source->moments->My;
	} else {
	    Xo = source->peak->xf;
	    Yo = source->peak->yf;
	}

	// set PSF parameters for this model (apply 2D shape model to coordinates Xo, Yo)
	pmModel *modelPSF = pmModelFromPSFforXY(psf, Xo, Yo, Io);

	if (modelPSF == NULL) {
	    psWarning ("Failed to determine PSF model at (%f,%f); trying image center", Xo, Yo);

	    float Xc = 0.5*readout->image->numCols;
	    float Yc = 0.5*readout->image->numRows;
	    pmModel *modelPSF = pmModelFromPSFforXY(psf, Xc, Yc, Io);
	    if (modelPSF == NULL) {
		psError(PSPHOT_ERR_PSF, false, "Failed to determine PSF model at center of image");
		return false;
	    }

	    // Now set the object position at the expected location:
	    modelPSF->params->data.F32[PM_PAR_XPOS] = Xo;
	    modelPSF->params->data.F32[PM_PAR_YPOS] = Yo;
	    source->mode |= PM_SOURCE_MODE_BADPSF;
	}

	// set the fit radius based on the object flux limit and the model
	// this function affects the mask pixels
	psphotCheckRadiusPSF (readout, source, modelPSF, markVal);

	// set the source PSF model
	source->modelPSF = modelPSF;
	source->modelPSF->residuals = psf->residuals;

	pmSourceCacheModel (source, maskVal);  // ALLOC x14 (!)

    }

    return true;
}
