# include "psphotInternal.h"

bool psphotChooseCellSizes (int *Cx, int *Cy, pmReadout *readout, int nThreads);
psphotGuessModelForRegionArgs *psphotGuessModelForRegionArgsAlloc();
bool psphotGuessModelForRegion (psphotGuessModelForRegionArgs *args);

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

# define NFILL 4

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

    // int nThreads = psMetadataLookupS32(&status, config->arguments, "NTHREADS"); // Number of threads
    // if (!status) {
    // nThreads = 0;
    // }
    int 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);

    // the strategy here is to divide the image into 2x2 blocks of cells and cycle through
    // the four discontiguous sets of cells, threading all within a set and blocking between
    // sets 

    // we divide the image region into 2*2 blocks of size Nx*Ny, the image will have 
    // Cx*Cy blocks so that (2Nx)Cx = numCols, (2Ny)Cy = numRows.  We want to choose Cx and
    // Cy so that (2Nx)Cx * (2Ny)Cy = 4 * NFILL * nThreads -- each of the four sets of cells
    // has enough cells to allow NFILL cells for each thread (to better distribute heavy and
    // light load cells
    
    // choose Cx, Cy:
    int Cx = 1, Cy = 1;
    psphotChooseCellSizes (&Cx, &Cy, readout, nThreads);

    // the array runs from readout->image->col0 to readout->image->col0 + readout->image->numCols 
    int Xo = readout->image->col0;
    int Yo = readout->image->row0;
    int Nx = readout->image->numCols / (2*Cx);
    int Ny = readout->image->numRows / (2*Cy);

    // we can thread all of the cells within a set, but need to block between sets
    for (int jy = 0; jy < 2; jy++) {
	for (int jx = 0; jx < 2; jx++) {

	    // generate jobs for all of the Cx*Cy cells in this set
	    for (int ix = 0; ix < Cx; ix++) {
		for (int iy = 0; iy < Cy; iy++) {
		
		    int x0 = (2*ix + jx)*Nx + Xo;
		    int x1 = x0 + Nx;

		    int y0 = (2*iy + jy)*Ny + Yo;
		    int y1 = y0 + Ny;

		    if (readout->image->numCols + Xo - x1 < Nx) {
			x1 = readout->image->numCols + Xo;
		    }
		    if (readout->image->numRows + Yo - y1 < Ny) {
			y1 = readout->image->numRows + Yo;
		    }

		    // fprintf (stderr, "launch %d,%d - %d,%d\n", x0, y0, x1, y1);
		    
		    psRegion *cellRegion = psRegionAlloc (x0, x1, y0, y1);
		    *cellRegion = psRegionForImage (readout->image, *cellRegion);
		    // if we got the math above right, this should be a NOP
									     
		    psphotGuessModelForRegionArgs *args = psphotGuessModelForRegionArgsAlloc();
		    args->readout = psMemIncrRefCounter(readout);
		    args->sources = psMemIncrRefCounter(sources);
		    args->psf     = psMemIncrRefCounter(psf);
		    args->region  = psMemIncrRefCounter(cellRegion);

		    args->maskVal = maskVal;
		    args->markVal = markVal;

		    // allocate a job -- if threads are not defined, this just runs the job
		    psThreadJob *job = psThreadJobAlloc ("PSPHOT_GUESS_MODEL");
		    psArrayAdd(job->args, 1, args);
		    if (!psThreadJobAddPending(job)) {
			psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
			return false;
		    }
		    psFree(cellRegion);
		    psFree(job);
		    psFree(args);
		}
	    }

	    // 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);
	    }
	}
    }
    
    psLogMsg ("psphot.models", 4, "built models for %ld objects: %f sec\n", sources->n, psTimerMark ("psphot.models"));
    return true;
}

static void psphotGuessModelForRegionArgsFree(psphotGuessModelForRegionArgs *args)
{
    psFree(args->readout);
    psFree(args->sources);
    psFree(args->psf);
    psFree(args->region);
    return;
}

psphotGuessModelForRegionArgs *psphotGuessModelForRegionArgsAlloc()
{
    psphotGuessModelForRegionArgs *args = psAlloc(sizeof(psphotGuessModelForRegionArgs));
    psMemSetDeallocator(args, (psFreeFunc)psphotGuessModelForRegionArgsFree);

    args->readout = NULL;
    args->sources = NULL;
    args->psf = NULL;
    args->region = NULL;

    args->maskVal = 0;
    args->markVal = 0;
    return args;
}

// construct models only for sources in the specified region
bool psphotGuessModelForRegion (psphotGuessModelForRegionArgs *args) {

    pmReadout *readout = args->readout;
    psArray *sources = args->sources;
    pmPSF *psf = args->psf;
    psRegion *region = args->region;
    psImageMaskType maskVal = args->maskVal;
    psImageMaskType markVal = args->markVal;

    int nSrc = 0;

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

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

	if (source->peak->xf <  region->x0) continue;
	if (source->peak->yf <  region->y0) continue;
	if (source->peak->xf >= region->x1) continue;
	if (source->peak->yf >= region->y1) 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);

    }

    // fprintf (stderr, "%d for region %lf,%lf - %lf,%lf\n", nSrc, region->x0, region->y0, region->x1, region->y1);
    return true;
}

bool psphotChooseCellSizes (int *Cx, int *Cy, pmReadout *readout, int nThreads) {

    int nCells = nThreads * NFILL; // number of cells in a single set
    int C = sqrt(nCells) + 0.5;
    
    // we need to assign Cx and Cy based on the dimensionality of the image
    // crude way to find most evenly balanced factors of nCells:
    for (int i = C; i >= 1; i--) {
	int C1 = nCells / C;
	int C2 = nCells / C1;
	if (C1*C2 != nCells) continue;

	if (readout->image->numRows > readout->image->numCols) {
	    *Cx = PS_MAX (C1, C2);
	    *Cy = PS_MIN (C1, C2);
	} else {
	    *Cx = PS_MAX (C1, C2);
	    *Cy = PS_MIN (C1, C2);
	}
	return true;
    }
    *Cx = 1;
    *Cy = 1; 

    return true;
}
