# 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

// A guess for when the moments aren't available
static pmModel *wildGuess(pmSource *source, // Source for which to guess
                          pmPSF *psf    // The point-spread function
    )
{
    pmModel *model = pmModelAlloc(psf->type);
    psF32 *PAR = model->params->data.F32;
    PAR[PM_PAR_SKY]  = 0;
    // XXX get this from the image pixels
    PAR[PM_PAR_I0]   = source->peak->flux;
    PAR[PM_PAR_XPOS] = source->peak->xf;
    PAR[PM_PAR_YPOS] = source->peak->yf;
    return model;
}

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

# if (0)		
		if (!psphotGuessModel_Unthreaded (readout, cells->data[j], psf, maskVal, markVal)) {
		    psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		    return false;
		}
# endif
	}

	// 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 ++;
	
	// XXX if a source is faint, it will not have moments measured.
	// it must be modelled as a PSF.  In this case, we need to use
	// the peak centroid to get the coordinates and get the peak flux
	// from the image?
	pmModel *modelEXT;
	if (!source->moments) {
	    modelEXT = wildGuess(source, psf);
	} else {
	    // use the source moments, etc to guess basic model parameters
	    modelEXT = pmSourceModelGuess (source, psf->type); // ALLOC X5
	    if (!modelEXT) {
		modelEXT = wildGuess(source, psf);
	    }
	    // these valuse are set in pmSourceModelGuess, should this rule be in there as well?
	    if (source->mode &  PM_SOURCE_MODE_SATSTAR) {
		modelEXT->params->data.F32[PM_PAR_XPOS] = source->moments->Mx;
		modelEXT->params->data.F32[PM_PAR_YPOS] = source->moments->My;
	    } else {
		modelEXT->params->data.F32[PM_PAR_XPOS] = source->peak->xf;
		modelEXT->params->data.F32[PM_PAR_YPOS] = source->peak->yf;
	    }
	}

	// set PSF parameters for this model (apply 2D shape model)
	pmModel *modelPSF = pmModelFromPSF (modelEXT, psf); // ALLOC X5
	if (modelPSF == NULL) {
	    psWarning ("Failed to determine PSF model at r,c = (%d,%d); trying centre of image",
		    source->peak->y, source->peak->x);

	    // Try the center of the image
	    modelEXT->params->data.F32[PM_PAR_XPOS] = 0.5*readout->image->numCols;
	    modelEXT->params->data.F32[PM_PAR_YPOS] = 0.5*readout->image->numRows;
	    modelPSF = pmModelFromPSF (modelEXT, psf);
	    if (modelPSF == NULL) {
		psError(PSPHOT_ERR_PSF, false, "Failed to determine PSF model at center of image");
		psFree(modelEXT);
		return false;
	    }
	    source->mode |= PM_SOURCE_MODE_BADPSF;
	}
	psFree (modelEXT); // FREE (x3)

	// XXX need to define the guess flux?
	// 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;
}

# if (0)
// construct models only for sources in the specified region
bool psphotGuessModel_Unthreaded (pmReadout *readout, psArray *sources, pmPSF *psf, psImageMaskType maskVal, psImageMaskType markVal) {

    int nSrc = 0;

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

	// XXXX this is just for a test: use this to mark sources for which the model is measured
	// check later that all are used.
	source->mode |= PM_SOURCE_MODE_EXT_LIMIT;

	// 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 ++;
	
	// XXX if a source is faint, it will not have moments measured.
	// it must be modelled as a PSF.  In this case, we need to use
	// the peak centroid to get the coordinates and get the peak flux
	// from the image?
	pmModel *modelEXT;
	if (!source->moments) {
	    modelEXT = wildGuess(source, psf);
	} else {
	    // use the source moments, etc to guess basic model parameters
	    modelEXT = pmSourceModelGuess (source, psf->type); // ALLOC
	    if (!modelEXT) {
		modelEXT = wildGuess(source, psf);
	    }
	    // these valuse are set in pmSourceModelGuess, should this rule be in there as well?
	    if (source->mode &  PM_SOURCE_MODE_SATSTAR) {
		modelEXT->params->data.F32[PM_PAR_XPOS] = source->moments->Mx;
		modelEXT->params->data.F32[PM_PAR_YPOS] = source->moments->My;
	    } else {
		modelEXT->params->data.F32[PM_PAR_XPOS] = source->peak->xf;
		modelEXT->params->data.F32[PM_PAR_YPOS] = source->peak->yf;
	    }
	}

	// set PSF parameters for this model (apply 2D shape model)
	pmModel *modelPSF = pmModelFromPSF (modelEXT, psf); // ALLOC
	if (modelPSF == NULL) {
	    psError(PSPHOT_ERR_PSF, false,
		    "Failed to determine PSF model at r,c = (%d,%d); trying centre of image",
		    source->peak->y, source->peak->x);
	    //
	    // Try the centre of the image
	    //
	    modelEXT->params->data.F32[PM_PAR_XPOS] = 0.5*readout->image->numCols;
	    modelEXT->params->data.F32[PM_PAR_YPOS] = 0.5*readout->image->numRows;
	    modelPSF = pmModelFromPSF (modelEXT, psf);
	    if (modelPSF == NULL) {
		psError(PSPHOT_ERR_PSF, false,
			"Failed to determine PSF model at centre of image");
		psFree(modelEXT);
		return false;
	    }

	    source->mode |= PM_SOURCE_MODE_BADPSF;
	}
	psFree (modelEXT);

	// XXX need to define the guess flux?
	// 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);

    }

    return true;
}
# endif
