# include "psphotInternal.h"

#define PSF_SOURCE_MASK (PM_SOURCE_MODE_FAIL | PM_SOURCE_MODE_SATSTAR | PM_SOURCE_MODE_BLEND | \
                         PM_SOURCE_MODE_BADPSF | PM_SOURCE_MODE_DEFECT | PM_SOURCE_MODE_SATURATED | \
                         PM_SOURCE_MODE_CR_LIMIT | PM_SOURCE_MODE_EXT_LIMIT) // Mask to apply for PSF sources

// for now, let's store the detections on the readout->analysis for each readout
bool psphotMergeSources (pmConfig *config, const pmFPAview *view)
{
    bool status = true;

    int num = psMetadataLookupS32 (&status, config->arguments, "PSPHOT.INPUT.NUM");
    psAssert (status, "programming error: must define PSPHOT.INPUT.NUM");

    // loop over the available readouts
    for (int i = 0; i < num; i++) {
        if (!psphotMergeSourcesReadout (config, view, "PSPHOT.INPUT", i)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed to merge sources for PSPHOT.INPUT entry %d", i);
            return false;
        }
    }
    return true;
}

// add newly selected sources to the existing list of sources
bool psphotMergeSourcesReadout (pmConfig *config, const pmFPAview *view, const char *filename, int index) {

    bool status;

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

    psArray *newSources = detections->newSources;
    psAssert (newSources, "missing sources?");

    // XXX TEST:
    if (detections->allSources) {
        psphotMaskCosmicRayFootprintCheck(detections->allSources);
    }
    if (detections->newSources) {
        psphotMaskCosmicRayFootprintCheck(detections->newSources);
    }

    if (!detections->allSources) {
        detections->allSources = psArrayAllocEmpty(newSources->n);
    }
    psArray *allSources = detections->allSources;

    for (int i = 0; i < newSources->n; i++) {
        pmSource *source = newSources->data[i];
        psArrayAdd (allSources, 100, source);
    }

    psFree (detections->newSources);
    detections->newSources = NULL;

    return true;
}

// Merge the externally supplied sources with the existing sources.  Mark them as having mode
// PM_SOURCE_MODE_EXTERNAL.

// XXX This function needs to be updated to loop over set of input files.  At the moment, we
// only expect a single entry for PSPHOT.INPUT.CMF and PSPHOT.SOURCES.TEXT, so we can only
// associate input sources with a single entry for PSPHOT.INPUT
bool psphotLoadExtSources (pmConfig *config, const pmFPAview *view) {

    bool status;
    pmDetections *extCMF = NULL;
    psArray *extSourcesTXT = NULL;
    int index = 0;

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PSPHOT.INPUT", 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");
    if (!detections) {
	detections = pmDetectionsAlloc();
	detections->newSources = psArrayAllocEmpty (100);
	// save detections on the readout->analysis
	if (!psMetadataAddPtr (readout->analysis, PS_LIST_TAIL, "PSPHOT.DETECTIONS", PS_META_REPLACE | PS_DATA_UNKNOWN, "psphot detections", detections)) {
	    psError (PSPHOT_ERR_CONFIG, false, "problem saving detections on readout");
	    return false;
	}
    } else {
	psMemIncrRefCounter(detections); // so we can free the detections below
    }

    psArray *sources = detections->newSources;
    psAssert (sources, "missing sources?");

    // load data from input CMF file:
    {
        pmReadout *readoutCMF = pmFPAfileThisReadout (config->files, view, "PSPHOT.INPUT.CMF");
        if (!readoutCMF) goto loadTXT;

        extCMF = psMetadataLookupPtr (NULL, readoutCMF->analysis, "PSPHOT.DETECTIONS");
        if (extCMF) {
            for (int i = 0; i < extCMF->allSources->n; i++) {
                pmSource *source = extCMF->allSources->data[i];
                source->mode |= PM_SOURCE_MODE_EXTERNAL;

                // the supplied peak flux needs to be re-normalized
                source->peak->flux = 1.0;
                source->peak->value = 1.0;

                // drop the loaded source modelPSF
                psFree (source->modelPSF);
                source->modelPSF = NULL;
		source->imageID = index;

                psArrayAdd (detections->newSources, 100, source);
            }
        }
    }

loadTXT:

    // load data from input TXT file:
    {
        pmChip *chipTXT = pmFPAfileThisChip (config->files, view, "PSPHOT.INPUT");
        if (!chipTXT) goto finish;

        extSourcesTXT = psMetadataLookupPtr (NULL, chipTXT->analysis, "PSPHOT.SOURCES.TEXT");
        if (extSourcesTXT) {
            for (int i = 0; i < extSourcesTXT->n; i++) {
                pmSource *source = extSourcesTXT->data[i];
                source->mode |= PM_SOURCE_MODE_EXTERNAL;

                // drop the loaded source modelPSF
                psFree (source->modelPSF);
                source->modelPSF = NULL;
		source->imageID = index;

                psArrayAdd (detections->newSources, 100, source);
            }
        }
    }

finish:

    psFree (detections);

    if (!extCMF && !extSourcesTXT) {
        psLogMsg ("psphot", 3, "no external sources for this readout");
        return true;
    }

    int nCMF = extCMF        ? extCMF->allSources->n        : 0;
    int nTXT = extSourcesTXT ? extSourcesTXT->n : 0;

    psLogMsg ("psphot", 3, "%d external sources (%d cmf, %d text) merged to yield %ld total sources",
              nCMF + nTXT, nCMF, nTXT, sources->n);
    return true;
}

// extract the input sources corresponding to this readout
// XXX this function needs to be updated to work with the new context of pshot inputs
psArray *psphotLoadPSFSources (pmConfig *config, const pmFPAview *view) {

    bool status;

    // find the currently selected readout
    pmReadout  *readout = pmFPAfileThisReadout (config->files, view, "PSPHOT.INPUT.CMF");
    if (!readout) {
        psLogMsg ("psphot", 3, "readout not found");
        return NULL;
    }

    pmDetections *detections = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.DETECTIONS");
    if (!detections) {
        psLogMsg ("psphot", 3, "no psf sources for this readout");
        return NULL;
    }

    psArray *sources = detections->allSources;
    if (!sources) {
        psLogMsg ("psphot", 3, "no psf sources for this readout");
        return NULL;
    }

    return sources;
}

// this function is used to fix sources which were loaded externally, but have passed from
// psphotDetectionsFromSources to psphotSourceStats and are now stored on
// detections->newSources.
bool psphotRepairLoadedSources (pmConfig *config, const pmFPAview *view) {

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PSPHOT.INPUT", 0); // File of interest
    psAssert (file, "missing file?");

    pmReadout *readout = pmFPAviewThisReadout(view, file->fpa);
    psAssert (readout, "missing readout?");

    pmDetections *detections = psMetadataLookupPtr (NULL, readout->analysis, "PSPHOT.DETECTIONS");
    if (!detections) {
        psError(PSPHOT_ERR_CONFIG, false, "missing detections");
        return false;
    }

    psArray *sources = detections->newSources;
    psAssert (sources, "missing sources?");

    // peak flux is wrong : set based on previous image
    // use the peak measured in the moments analysis:
    for (int i = 0; i < sources->n; i++) {
      pmSource *source = sources->data[i];
      source->peak->flux = source->moments->Peak;
    }

    return true;
}

// generate the detection structure for the supplied array of sources
// XXX this currently assumes there is a single input file
bool psphotDetectionsFromSources (pmConfig *config, const pmFPAview *view, psArray *sources) {

    // find the currently selected readout
    pmFPAfile *file = pmFPAfileSelectSingle(config->files, "PSPHOT.INPUT", 0); // File of interest
    psAssert (file, "missing file?");

    pmReadout *readout = pmFPAviewThisReadout(view, file->fpa);
    psAssert (readout, "missing readout?");

    pmDetections *detections = pmDetectionsAlloc();

    detections->peaks = psArrayAllocEmpty(100);

    psMetadata *recipe = psMetadataLookupPtr (NULL, config->recipes, PSPHOT_RECIPE);
    if (!recipe) {
        psError(PSPHOT_ERR_CONFIG, false, "missing recipe %s", PSPHOT_RECIPE);
        return false;
    }
    float snMin = psMetadataLookupF32(NULL, recipe, "MOMENTS_SN_MIN");
    if (!isfinite(snMin)) {
        return false;
    }

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

        if (source->mode & PSF_SOURCE_MASK || !isfinite(source->psfMag)) {
            continue;
        }

        // use the existing peak information, otherwise generate a new peak
        if (source->peak) {
            source->peak->assigned = false; // So the moments will be measured
            psArrayAdd (detections->peaks, 100, source->peak);
            continue;
        }

        float flux = powf(10.0, -0.4 * source->psfMag);
        float xpos = model->params->data.F32[PM_PAR_XPOS];
        float ypos = model->params->data.F32[PM_PAR_YPOS];

        pmPeak *peak = pmPeakAlloc(xpos, ypos, 1.0, PM_PEAK_LONE);
        peak->xf = xpos;
        peak->yf = ypos;
        peak->flux = flux; // this are being set wrong, but does it matter?

        if (isfinite (source->errMag) && (source->errMag > 0.0)) {
          peak->SN = 1.0 / source->errMag;
        } else {
          peak->SN = 0.0;
        }

        psArrayAdd (detections->peaks, 100, peak);
        psFree (peak);
    }

    psLogMsg ("psphot", 3, "%ld PSF sources loaded", detections->peaks->n);

    // save detections on the readout->analysis
    if (!psMetadataAddPtr (readout->analysis, PS_LIST_TAIL, "PSPHOT.DETECTIONS", PS_META_REPLACE | PS_DATA_UNKNOWN, "psphot detectinos", detections)) {
        psError (PSPHOT_ERR_CONFIG, false, "problem saving detections on readout");
        return false;
    }
    psFree (detections);
    return true;
}

// generate the detection structure for the supplied array of sources
bool psphotSetSourceParams (pmConfig *config, psArray *sources, pmPSF *psf) {

    psMetadata *recipe = psMetadataLookupPtr (NULL, config->recipes, PSPHOT_RECIPE);
    if (!recipe) {
        psError(PSPHOT_ERR_CONFIG, false, "missing recipe %s", PSPHOT_RECIPE);
        return false;
    }

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

        if (source->mode & PSF_SOURCE_MASK || !isfinite(source->psfMag)) {
            continue;
        }

        float flux = powf(10.0, -0.4 * source->psfMag);
        float xpos = model->params->data.F32[PM_PAR_XPOS];
        float ypos = model->params->data.F32[PM_PAR_YPOS];

        pmPeak *peak = pmPeakAlloc(xpos, ypos, 1.0, PM_PEAK_LONE);
        peak->xf = xpos;
        peak->yf = ypos;
        peak->flux = flux; // this are being set wrong, but does it matter?

        if (isfinite (source->errMag) && (source->errMag > 0.0)) {
          peak->SN = 1.0 / source->errMag;
        } else {
          peak->SN = 0.0;
        }

        source->peak = peak;
    }

    psLogMsg ("psphot", 3, "%ld PSF sources loaded", sources->n);

    return true;
}

bool psphotCheckExtSources (pmConfig *config, const pmFPAview *view) {

    bool status;

    psMetadata *recipe = psMetadataLookupPtr (NULL, config->recipes, PSPHOT_RECIPE);
    psAssert (recipe, "missing recipe");

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

    // XXX allSources of newSources?
    psArray *sources = detections->allSources;
    psAssert (sources, "missing sources?");

    if (sources->n) {
        // the user wants to make the psf from these stars; define them as psf stars:
        for (int i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i];
            source->mode |= PM_SOURCE_MODE_PSFSTAR;
        }
        // force psphotChoosePSF to use all loaded sources
        psMetadataAddS32 (recipe, PS_LIST_TAIL, "PSF_MAX_NSTARS", PS_META_REPLACE, "max number of sources for PSF model", sources->n);

        // measure stats of externally specified sources
        if (!psphotSourceStatsUpdate (sources, config, readout)) {
            psError(PSPHOT_ERR_CONFIG, false, "failure to measure stats of existing sources");
            return false;
        }
    } else {

        // find the detections (by peak and/or footprint) in the image.
        if (!psphotFindDetections (config, view, true)) {
            psError(PSPHOT_ERR_CONFIG, false, "unable to find detections in this image");
            return psphotReadoutCleanup (config, view);
        }

        // construct sources and measure basic stats
        psphotSourceStats (config, view, true);

        // find blended neighbors of very saturated stars
        psphotDeblendSatstars (config, view);

        // mark blended peaks PS_SOURCE_BLEND
        if (!psphotBasicDeblend (config, view)) {
            psLogMsg ("psphot", 3, "failed on deblend analysis");
            return psphotReadoutCleanup (config, view);
        }

        // classify sources based on moments, brightness
        if (!psphotRoughClass (config, view)) {
            psLogMsg ("psphot", 3, "failed to find a valid PSF clump for image");
            return psphotReadoutCleanup (config, view);
        }
    }

    return true;
}
