# include "psphotInternal.h" 

bool psphotMatchSourcesAddMissing (pmConfig *config, const pmFPAview *view, const char *filerule, psArray *objects);
bool psphotMatchSourcesSetIDs (psArray *objects);
 
psArray *psphotMatchSources (pmConfig *config, const pmFPAview *view, const char *filerule) 
{
    psArray *objects = psArrayAllocEmpty(100);

    int num = psphotFileruleCount(config, filerule);

    // loop over the available readouts
    for (int i = 0; i < num; i++) {
        if (!psphotMatchSourcesReadout (objects, config, view, filerule, i)) {
	    psError (PSPHOT_ERR_CONFIG, false, "failed to merge sources for %s entry %d", filerule, i);
	    psFree (objects);
            return NULL;
        }
    }

    // create sources for images where an object has been detected in the other images
    psphotMatchSourcesAddMissing (config, view, filerule, objects);

    // choose a consistent position; set common sequence values
    psphotMatchSourcesSetIDs (objects);

    return objects;
}

bool psphotMatchSourcesReadout (psArray *objects, pmConfig *config, const pmFPAview *view, const char *filerule, int index) { 
 
    bool status = false;

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

    int RADIUS = psMetadataLookupF32 (&status, recipe, "PSPHOT.STACK.MATCH.RADIUS");
    psAssert (status, "programming error: must define PSPHOT.STACK.MATCH.RADIUS");

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, filerule, 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?");
    psAssert (detections->allSources, "all sources not defined?");

    psphotMatchSourcesToObjects(objects, detections->allSources, RADIUS);

    return true;
}

# define NEXT1 { i++; continue; } 
# define NEXT2 { j++; continue; } 
bool psphotMatchSourcesToObjects (psArray *objects, psArray *sources, float RADIUS) { 
 
    float dx, dy; 
 
    float RADIUS2 = RADIUS*RADIUS;

    // sort the source list by X 
    sources = psArraySort (sources, pmSourceSortByX); 
    objects = psArraySort (objects, pmPhotObjSortByX); 
 
    psVector *foundSrc = psVectorAlloc(sources->n, PS_TYPE_U8);
    psVectorInit (foundSrc, 0);

    psVector *foundObj = psVectorAlloc(objects->n, PS_TYPE_U8);
    psVectorInit (foundObj, 0);

    // match sources to existing objects

    psLogMsg ("psphot", PS_LOG_DETAIL, "attempt to match sources (%ld vs %ld)", sources->n, objects->n);

    int i, j; 
    for (i = j = 0; (i < sources->n) && (j < objects->n); ) { 
 
        pmSource  *src = sources->data[i]; 
        pmPhotObj *obj = objects->data[j]; 
 
        if (!src) NEXT1; 
        if (!src->peak) NEXT1; 
        if (!isfinite(src->peak->xf)) NEXT1; 
        if (!isfinite(src->peak->yf)) NEXT1; 
 
        if (!obj) NEXT2; 
        if (!isfinite(obj->x)) NEXT2; 
        if (!isfinite(obj->y)) NEXT2; 
 
        dx = src->peak->xf - obj->x; 
        if (dx < -1.02*RADIUS) NEXT1; 
        if (dx > +1.02*RADIUS) NEXT2; 
 
	/* this block will match a given detection to the closest object within range of that detection.
	   XXX note that this matches ALL detections within range of the single object to that same object 
	   this is bad, but I cannot just go in linear order (ie, mark off each object as they are
	   used).  I should make a list of all Nobj * Ndet pairs in range and choose the matches
	   based on their separations.  UGH
	*/
    
        // we are within match range, look for matches: 
	int Jmin = -1;
	float Rmin = RADIUS2;
        for (int J = j; (dx > -1.02*RADIUS) && (J < objects->n); J++) { 
 
	    // skip objects that are already assigned:
	    if (foundObj->data.U8[J]) continue;
	    obj = objects->data[J]; 
	    
	    dx = src->peak->xf - obj->x; 
            dy = src->peak->yf - obj->y; 
 
            float dr = dx*dx + dy*dy; 
            if (dr > RADIUS2) continue; 
	    if (dr < Rmin) {
		Rmin = dr;
		Jmin  = J;
	    }
	}

	// no match, try next source
	if (Jmin == -1) {
	    i++;
	    continue;
	}
	obj = objects->data[Jmin]; 
	foundObj->data.U8[Jmin] = 1;

	// add to object
	pmPhotObjAddSource (obj, src);
	foundSrc->data.U8[i] = 1;
        i++; 
    } 

    // add missed sources to new objects

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

	if (foundSrc->data.U8[i]) continue;

        pmSource *src = sources->data[i]; 

	pmPhotObj *obj = pmPhotObjAlloc();
	pmPhotObjAddSource(obj, src);
	psArrayAdd (objects, 100, obj);
	psFree (obj);
    }
    psLogMsg ("psphot", PS_LOG_DETAIL, "matched sources (%ld vs %ld)", sources->n, objects->n);

    psFree (foundSrc);
    psFree (foundObj);
    return true;
} 

bool psphotMatchSourcesAddMissing (pmConfig *config, const pmFPAview *view, const char *filerule, psArray *objects) {

    bool status = false;

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

    // determine properties (sky, moments) of initial sources
    float OUTER = psMetadataLookupF32 (&status, recipe, "SKY_OUTER_RADIUS");
    psAssert (status, "missing SKY_OUTER_RADIUS in recipe?");

    int nImages = psphotFileruleCount(config, filerule);

    // generate look-up arrays for detections and readouts
    psArray *detArrays = psArrayAlloc(nImages);
    psArray *readouts = psArrayAlloc(nImages);

    for (int i = 0; i < nImages; i++) {

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

	detArrays->data[i] = psMemIncrRefCounter(detections);
	readouts->data[i] = psMemIncrRefCounter(readout);
    }

    // vector to track if source for an image is found
    psVector *found = psVectorAlloc(nImages, PS_TYPE_U8);

    for (int i = 0; i < objects->n; i++) { 
        pmPhotObj *obj = objects->data[i]; 

	// we will find the input source with the max number of spans and reproduce that footprint
	int nSpansMax = 0;
	int iSpansMax = -1;

	// mark the images for which sources have been found
	psVectorInit (found, 0);
	for (int j = 0; j < obj->sources->n; j++) {

	    pmSource *src = obj->sources->data[j]; 
	    int index = src->imageID;
	    psAssert (index >= 0, "invalid index");
	    psAssert (index < found->n, "invalid index");

	    if (src->peak && src->peak->footprint && src->peak->footprint->nspans > nSpansMax) {
		nSpansMax = src->peak->footprint->nspans;
		iSpansMax = j;
	    }

	    found->data.U8[index] = 1;
	}

	// we make a copy of the largest footprint; this will be used for all new sources associated with this object
	pmFootprint *footprint = NULL;
	if (iSpansMax != -1) { // copy the footprint info
	    pmSource *src = obj->sources->data[iSpansMax]; 
	    psAssert(src->peak, "source does not exist?");
	    psAssert(src->peak->footprint, "footprint does not exist");
	    psAssert(src->peak->footprint->nspans == nSpansMax, "wrong footprint?");
	    
	    // we only care about the spans, do not worry about the image of this footprint
	    footprint = pmFootprintCopyData(src->peak->footprint, NULL);
	}

	// generate new sources for the image that are missing
	for (int index = 0; index < found->n; index++) {
	    if (found->data.U8[index]) continue;

	    pmDetections *detections = detArrays->data[index];
	    pmReadout *readout = readouts->data[index];
	    int row0 = readout->image->row0;
	    int col0 = readout->image->col0;

	    // The peak type is not used in psphot. PM_PEAK_LONE may be wrong, but irrelevant
	    float peakFlux = readout->image->data.F32[(int)(obj->y-row0-0.5)][(int)(obj->x-col0-0.5)];
	    pmPeak *peak = pmPeakAlloc(obj->x, obj->y, peakFlux, PM_PEAK_LONE);
	    peak->xf = obj->x;
	    peak->yf = obj->y;
	    peak->dx = NAN;
	    peak->dy = NAN;
	    
	    // assign to a footprint on this readout->image
	    if (footprint) {
		peak->footprint = pmFootprintCopyData(footprint, readout->image);

		// the peak does not claim ownership of the footprint (it does not free it). save a copy of this 
		// footprint on detections->footprints so we can free it later
		psArrayAdd(detections->footprints, 100, peak->footprint); 
		psFree (peak->footprint);
	    }
	    
	    // create a new source
	    pmSource *source = pmSourceAlloc();
	    source->imageID = index;
	    source->mode2 |= PM_SOURCE_MODE2_MATCHED; // source is generated based on another image
	    source->type = PM_SOURCE_TYPE_STAR; // until we know more, assume a PSF fit

	    // add the peak
	    source->peak = 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;
	    pmPhotObjAddSource(obj, source);
	    psArrayAdd (detections->allSources, 100, source);
	    psFree (source);
	}
	psFree (footprint);
    }

    // how many sources do we have now?
    int nSources = 0;
    for (int i = 0; i < objects->n; i++) { 
        pmPhotObj *obj = objects->data[i]; 
	nSources += obj->sources->n;
	psAssert (obj->sources->n == nImages, "failed to match sources?");
    }
    psLogMsg ("psphot", PS_LOG_DETAIL, "total of %d sources for %d images", nSources, nImages);


    psFree (found);
    psFree (detArrays);
    psFree (readouts);
    return true;
}

bool psphotMatchSourcesSetIDs (psArray *objects) {

    for (int i = 0; i < objects->n; i++) { 
        pmPhotObj *obj = objects->data[i]; 

	// set the source->seq values 
	for (int j = 0; j < obj->sources->n; j++) {
	    pmSource *src = obj->sources->data[j]; 
	    src->seq = i;
	}
    }
    return true;
}
