# include "psphotInternal.h" 

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

static psArray     *psphotMatchFootprintCacheAlloc (int nImages);
static pmFootprint *psphotMatchLookupFootprint (psArray *cache, int footprintID, int imageID);
static pmFootprint *psphotMatchCopyFootprint (psArray *cache, pmFootprint *footprint, int imageID, psImage *image, psArray *footprints);
 
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);
    psArray *footprintCache = psphotMatchFootprintCacheAlloc(nImages);

    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 *largestFootprint = 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?");
	    
            largestFootprint = src->peak->footprint;
	}

	// 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 (largestFootprint) {
                // we save the copies that we make of the of the footprints in a hash so that we can reuse them
                // for all sources that share the fooprint. Without this we had serious memory explosion in
                // dense fields with lots of footprints (spans are small but when you have enough of them ...)
                peak->footprint = psphotMatchLookupFootprint(footprintCache, largestFootprint->id, index);
                if (!peak->footprint) {
                    // the peak does not claim ownership of the footprint (it does not free it). 
                    // psphotMatchCopyFootprint saves a copy of this 
                    // footprint on detections->footprints so we can free it later
                    peak->footprint = psphotMatchCopyFootprint(footprintCache, largestFootprint,
                                            index, readout->image, detections->footprints);
                }
	    }
	    
	    // 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);

#if (0)
            fprintf(stderr, "Add mising source for obj: %5d %5d image: %d flux: %f size: %4d %4d\n",
                                                      i, obj->id, index, peakFlux, source->pixels->numRows, source->pixels->numCols);
#endif

	    peak->assigned = true;
	    pmPhotObjAddSource(obj, source);
	    psArrayAdd (detections->allSources, 100, source);
	    psFree (source);
	}
    }
    psFree(footprintCache);

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

// Cache of footprints created for unmatched sources.
static psArray * psphotMatchFootprintCacheAlloc (int nImages) {
    psArray *cache = psArrayAlloc(nImages);
    for (int i = 0; i < nImages; i++) {
        psHash *hash = psHashAlloc(5000);
        cache->data[i] = hash;
    }
    return cache;
}

// Find copy of footprint with given ID made for a given image
static pmFootprint *psphotMatchLookupFootprint (psArray *cache, int footprintID, int imageID) {
    psHash *hash = (psHash *) cache->data[imageID];

    psAssert(hash != NULL, "missing hash for image %d", imageID);

    // footprintID is the id of the original footprint.
    char key[32];
    sprintf(key, "%d", footprintID); 

    // The footprints in our hashes are the copies of that footprint for the respective images
    pmFootprint *copy = psHashLookup(hash, key);

    return copy;
}

// Create a copy of a given footprint for a given image
static pmFootprint *psphotMatchCopyFootprint (psArray *cache, pmFootprint *footprint, int imageID, psImage *image, psArray *footprints) {

    psAssert((imageID >= 0 && imageID < cache->n), "invalid imageID %d", imageID);

    psHash *hash = (psHash *) cache->data[imageID];

    psAssert(hash != NULL, "missing hash for image %d", imageID);

    char key[32];
    sprintf(key, "%d", footprint->id); 

    pmFootprint *copy = pmFootprintCopyData (footprint, image);

    // save a copy in this image's hash
    psHashAdd(hash, key, copy);

    // save a copy of this footprint on the passed in footprints Array so we can free it later
    psArrayAdd(footprints, 100, copy);

    // drop our ref
    psFree(copy);

    return copy;
}
