# include "psphotInternal.h"

void psphotCountPSFStars (psArray *sources) {

    int nPSF = 0;

    // select the candidate PSF stars (pointers to original sources)
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];
        if (source->mode & PM_SOURCE_MODE_PSFSTAR) {
            nPSF ++;
        }
    }

    fprintf (stderr, "N PSF: %d\n", nPSF);
}

// try PSF models and select best option
pmPSF *psphotChoosePSF (pmReadout *readout, psArray *sources, psMetadata *recipe) {

    bool status;

    psTimerStart ("psphot.choose.psf");

    // bit-masks to test for good/bad pixels
    psMaskType maskVal = psMetadataLookupU8(&status, recipe, "MASK.PSPHOT");
    assert (maskVal);

    // bit-mask to mark pixels not used in analysis
    psMaskType markVal = psMetadataLookupU8(&status, recipe, "MARK.PSPHOT");
    assert (markVal);

    // maskVal is used to test for rejected pixels, and must include markVal
    maskVal |= markVal;

    // examine PSF sources in S/N order (brightest first)
    sources = psArraySort (sources, pmSourceSortBySN);

    // structure to store user options defining the psf
    pmPSFOptions *options = pmPSFOptionsAlloc ();

    // load user options from the recipe. no need to check existence -- they are
    // array to store candidate PSF stars
    int NSTARS = psMetadataLookupS32 (&status, recipe, "PSF_MAX_NSTARS");
    assert (status);

    // supply the measured sky variance for optional constant errors (non-poissonian)
    float SKY_SIG = psMetadataLookupF32 (&status, recipe, "SKY_SIG");
    assert (status);

    // use poissonian errors or local-sky errors
    options->poissonErrorsPhotLMM = psMetadataLookupBool (&status, recipe, "POISSON.ERRORS.PHOT.LMM");
    assert (status);

    // use poissonian errors or local-sky errors
    options->poissonErrorsPhotLin = psMetadataLookupBool (&status, recipe, "POISSON.ERRORS.PHOT.LIN");
    assert (status);

    // use poissonian errors or local-sky errors
    options->poissonErrorsParams = psMetadataLookupBool (&status, recipe, "POISSON.ERRORS.PARAMS");
    assert (status);

    // how to model the PSF variations across the field
    options->psfTrendMode = pmTrend2DModeFromString (psMetadataLookupStr (&status, recipe, "PSF.TREND.MODE"));
    assert (status);

    options->psfTrendNx = psMetadataLookupS32 (&status, recipe, "PSF.TREND.NX");
    assert (status);

    options->psfTrendNy = psMetadataLookupS32 (&status, recipe, "PSF.TREND.NY");
    assert (status);

    // get the fixed PSF fit radius
    // XXX check that PSF_FIT_RADIUS < SKY_OUTER_RADIUS
    options->radius = psMetadataLookupF32 (&status, recipe, "PSF_FIT_RADIUS");
    assert (status);

    // XXX ROBUST seems to be too agressive given the small numbers.
    // options->stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    options->stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);

    // dimensions of the field for which the PSF is defined
    options->psfFieldNx = readout->image->numCols;
    options->psfFieldNy = readout->image->numRows;
    options->psfFieldXo = readout->image->col0;
    options->psfFieldYo = readout->image->row0;

    int fitIter = psMetadataLookupS32(&status, recipe, "PSF_FIT_ITER"); // Maximum number of fit iterations
    if (!status || fitIter <= 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "PSF_FIT_ITER is not positive");
        return false;
    }
    float fitTol = psMetadataLookupF32 (&status, recipe, "PSF_FIT_TOL"); // Fit tolerance
    if (!status || !isfinite(fitTol) || fitTol <= 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "PSF_FIT_TOL is not positive");
        return false;
    }
    pmSourceFitModelInit(fitIter, fitTol, PS_SQR(SKY_SIG), options->poissonErrorsPhotLMM);

    psArray *stars = psArrayAllocEmpty (sources->n);

    // psphotCountPSFStars (sources);

    // select the candidate PSF stars (pointers to original sources)
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];
        if (source->mode & PM_SOURCE_MODE_PSFSTAR) {
            // keep NSTARS PSF stars, unmark the rest
            if (stars->n < NSTARS) {
                psArrayAdd (stars, 200, source);
            } else {
                source->mode &= ~PM_SOURCE_MODE_PSFSTAR;
            }
        }
    }

    // psphotCountPSFStars (sources);

    // check that the identified psf stars sufficiently cover the region; if not, extend the
    // limits somewhat
    psphotCheckStarDistribution (stars, sources, options);

    // psphotCountPSFStars (sources);

    psLogMsg ("psphot.pspsf", PS_LOG_DETAIL, "selected candidate %ld PSF objects\n", stars->n);

    // get the list pointers for the PSF_MODEL entries
    psArray *modelNames = NULL;
    psMetadataItem *item = psMetadataLookup (recipe, "PSF_MODEL");
    if (item == NULL) psAbort("missing PSF_MODEL selection");

    if (item->type == PS_DATA_STRING) {
        modelNames = psArrayAlloc(1);
        modelNames->data[0] = psStringCopy (item->data.V);
    } else {
        if (item->type != PS_DATA_METADATA_MULTI) psAbort("missing PSF_MODEL selection");
        modelNames = psListToArray (item->data.list);
    }

    // generate a psf model using the first selection
    if (stars->n == 0) {
        psError(PSPHOT_ERR_PSF, false, "Failed to fit any models when choosing PSF");
        bool mdok;                      // Status of MD lookup
        if (!psMetadataLookupBool(&mdok, recipe, PSPHOT_RECIPE_PSF_FAKE_ALLOW)) {
            return false;
        }
        psErrorStackPrint(stderr, "Using guess PSF model");
        psErrorClear();

        psFree(options);

        // unset the PSFSTAR flags (none are used):
        for (int i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i];
            source->mode &= ~PM_SOURCE_MODE_PSFSTAR;
        }

        // XXX set sxx, etc from FWHM in recipe
        pmPSF *psf = pmPSFBuildSimple (modelNames->data[0], 1.0, 1.0, 0.0, 1.0);
        psf->fieldNx = readout->image->numCols;
        psf->fieldNy = readout->image->numRows;
        psFree (modelNames);

        bool status = true;
        status &= psphotMakeFluxScale (readout->image, recipe, psf);
        status &= psphotPSFstats (readout, recipe, psf);
        if (!status) {
            psError(PSPHOT_ERR_PSF, false, "Failed to fit any models when choosing PSF");
            psFree (psf);
            return NULL;
        }

        // XXX set DSX_MEAN, etc?
        return psf;
    }

    // set up an array to store the results
    psArray *models = psArrayAlloc (modelNames->n);

    // try each model option listed in config
    for (int i = 0; i < modelNames->n; i++) {
        char *modelName = modelNames->data[i];
        pmPSFtry *try = pmPSFtryModel (stars, modelName, options, maskVal, markVal); // Attempt at fit
        if (!try) {
            // No big deal --- we'll try another model
            psErrorClear();
            continue;
        }
        models->data[i] = try;
    }

    // select the best of the models
    // here we are using the clippedStdev on the metric as the indicator
    int bestN = -1;
    float bestM = 0.0;
    for (int i = 0; i < models->n; i++) {
        pmPSFtry *try = models->data[i];
        if (try == NULL) {
            psTrace ("psphot", 3, "PSF model %d is NULL", i);
            continue;
        }
        float M = try->psf->dApResid;
        if (bestN < 0 || M < bestM) {
            bestM = M;
            bestN = i;
        }
    }

    // use the best model:
    if (bestN < 0) {
        psFree (models);
        psFree (stars);

        psErrorStackPrint (stderr, "Failed to fit any models when choosing PSF");
        psErrorClear();

        bool mdok;                      // Status of MD lookup
        if (!psMetadataLookupBool(&mdok, recipe, PSPHOT_RECIPE_PSF_FAKE_ALLOW)) {
            psFree (modelNames);
            psFree(options);
            return NULL;
        }

        // generate a psf model using the first selection
        psLogMsg ("psphot.pspsf", PS_LOG_INFO, "Using guess PSF model");

        // unset the PSFSTAR flags (none are used):
        for (int i = 0; i < sources->n; i++) {
            pmSource *source = sources->data[i];
            source->mode &= ~PM_SOURCE_MODE_PSFSTAR;
        }

        // XXX set sxx, etc from FWHM in recipe
        pmPSF *psf = pmPSFBuildSimple (modelNames->data[0], 1.0, 1.0, 0.0, 1.0);
        psf->fieldNx = readout->image->numCols;
        psf->fieldNy = readout->image->numRows;
        psFree (modelNames);

        bool status = true;
        status &= psphotMakeFluxScale (readout->image, recipe, psf);
        status &= psphotPSFstats (readout, recipe, psf);
        if (!status) {
            psError(PSPHOT_ERR_PSF, false, "Failed to fit any models when choosing PSF");
            psFree (psf);
            return NULL;
        }

        // XXX set DSX_MEAN, etc?
        return psf;
    }

    psFree (modelNames);
    psFree (stars);

    pmPSFtry *try = models->data[bestN];

    // unset the PSFSTAR flag for stars not used for PSF model
    for (int i = 0; i < try->sources->n; i++) {
        pmSource *source = try->sources->data[i];
        if (try->mask->data.U8[i]) {
            source->mode &= ~PM_SOURCE_MODE_PSFSTAR;
        }
    }

    // build a PSF residual image
    if (!psphotMakeResiduals (try->sources, recipe, try->psf, maskVal)) {
        psError(PSPHOT_ERR_PSF, false, "Unable to construct residual table for PSF");
        psFree (models);
        psFree(options);
        return NULL;
    }

    // build the flux-to-magnitude conversion information
    if (!psphotMakeFluxScale (readout->image, recipe, try->psf)) {
        psError(PSPHOT_ERR_PSF, false, "Unable to construct flux scale for PSF");
        psFree (models);
        psFree(options);
        return NULL;
    }

    // build curve-of-growth vector for this psf
    if (!psphotMakeGrowthCurve (readout, recipe, try->psf)) {
        psError(PSPHOT_ERR_PSF, false, "Unable to construct flux scale for PSF");
        psFree (models);
        psFree(options);
        return NULL;
    }

    // XXX test dump of psf star data and psf-subtracted image
    if (psTraceGetLevel("psphot.psfstars") > 5) {
        psphotDumpPSFStars (readout, try, options->radius, maskVal, markVal);
    }

    // save only the best model;
    // XXX we are not saving the fitted sources
    // XXX do we want to keep them so we may optionally write them out?
    pmPSF *psf = psMemIncrRefCounter(try->psf);
    psFree (models);

    if (!psphotPSFstats (readout, recipe, psf)) {
        psError(PSPHOT_ERR_PSF, false, "cannot measure PSF shape terms");
        psFree(options);
        psFree(psf);
        return NULL;
    }

    // psphotCountPSFStars (sources);

    char *modelName = pmModelClassGetName (psf->type);
    psLogMsg ("psphot.pspsf", PS_LOG_INFO, "select psf model: %f sec\n", psTimerMark ("psphot.choose.psf"));
    psLogMsg ("psphot.pspsf", PS_LOG_INFO, "psf model %s, ApResid: %f +/- %f\n", modelName, psf->ApResid, psf->dApResid);

    psFree (options);
    return (psf);
}

// measure average parameters of the PSF model
bool psphotPSFstats (pmReadout *readout, psMetadata *recipe, pmPSF *psf) {

    psEllipseShape shape;
    psEllipseAxes axes;

    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(recipe, false);
    PS_ASSERT_PTR_NON_NULL(psf, false);

    psImage *image = readout->image;
    PS_ASSERT_PTR_NON_NULL(image, false);

    // XXX a minor hack: measure the values for a grid of points across the image, determine mean, UQ, LQ, etc:
    psVector *fwhmMajor = psVectorAllocEmpty (100, PS_DATA_F32);
    psVector *fwhmMinor = psVectorAllocEmpty (100, PS_DATA_F32);

    for (float ix = -0.4; ix <= +0.4; ix += 0.1) {
        for (float iy = -0.4; iy <= +0.4; iy += 0.1) {

            // use the center of the center pixel of the image
            float xc = ix*image->numCols + 0.5*image->numCols + image->col0 + 0.5;
            float yc = iy*image->numRows + 0.5*image->numRows + image->row0 + 0.5;

            // create modelPSF from this model
            pmModel *modelPSF = pmModelFromPSFforXY (psf, xc, yc, 1.0);
            if (!modelPSF) continue;

            // get the model full-width at half-max
            float FWHM_MAJOR = 2*modelPSF->modelRadius (modelPSF->params, 0.5);

            // XXX make sure this is consistent with the re-definition of PM_PAR_SXX
            shape.sx  = modelPSF->params->data.F32[PM_PAR_SXX];
            shape.sy  = modelPSF->params->data.F32[PM_PAR_SYY];
            shape.sxy = modelPSF->params->data.F32[PM_PAR_SXY];
            axes = psEllipseShapeToAxes (shape, 20.0);
            psFree (modelPSF);

            float FWHM_MINOR = FWHM_MAJOR * (axes.minor / axes.major);
            if (!isfinite(FWHM_MAJOR) || !isfinite(FWHM_MINOR)) continue;
            psVectorAppend (fwhmMajor, FWHM_MAJOR);
            psVectorAppend (fwhmMinor, FWHM_MINOR);
        }
    }

    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV | PS_STAT_SAMPLE_QUARTILE);
    psVectorStats (stats, fwhmMajor, NULL, NULL, 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_MAJ",   PS_META_REPLACE, "PSF FWHM Major axis (mean)", stats->sampleMean);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_SG",   PS_META_REPLACE, "PSF FWHM Major axis (sigma)", stats->sampleStdev);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_LQ",   PS_META_REPLACE, "PSF FWHM Major axis (lower quartile)", stats->sampleLQ);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_UQ",   PS_META_REPLACE, "PSF FWHM Major axis (upper quartile)", stats->sampleUQ);

    psVectorStats (stats, fwhmMinor, NULL, NULL, 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_MIN",   PS_META_REPLACE, "PSF FWHM Minor axis (mean)", stats->sampleMean);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_SG",   PS_META_REPLACE, "PSF FWHM Minor axis (sigma)", stats->sampleStdev);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_LQ",   PS_META_REPLACE, "PSF FWHM Minor axis (lower quartile)", stats->sampleLQ);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_UQ",   PS_META_REPLACE, "PSF FWHM Minor axis (upper quartile)", stats->sampleUQ);

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "ANGLE",    PS_META_REPLACE, "PSF angle",           axes.theta);
    psMetadataAddS32 (recipe, PS_LIST_TAIL, "NPSFSTAR", PS_META_REPLACE, "Number of stars used to make PSF", psf->nPSFstars);
    psMetadataAddBool(recipe, PS_LIST_TAIL, "PSFMODEL", PS_META_REPLACE, "Valid PSF Model?", true);

    psFree (fwhmMajor);
    psFree (fwhmMinor);
    psFree (stats);
    return true;
}

// determine approximate PSF shape parameters based on the moments
bool psphotMomentsStats (pmReadout *readout, psMetadata *recipe, psArray *sources) {

    // without the PSF model, we can only rely on the source->moments
    // as a measure of the FWHM values

    int FWHM_N = 0;
    double FWHM_X = 0.0;
    double FWHM_Y = 0.0;
    double FWHM_T = 0.0;

    psEllipseMoments moments;
    psEllipseAxes axes;

    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(recipe, false);
    PS_ASSERT_PTR_NON_NULL(sources, false);

    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];
        if (!source) continue;
        if (!(source->mode & PM_SOURCE_MODE_PSFSTAR)) continue;
        if (!source->moments) {
            continue;
        }

        moments.x2 = source->moments->Mxx;
        moments.y2 = source->moments->Myy;
        moments.xy = source->moments->Mxy;

        // limit axis ratio < 20.0
        axes = psEllipseMomentsToAxes (moments, 20.0);

        FWHM_X += axes.major * 2.35;
        FWHM_Y += axes.minor * 2.35;
        FWHM_T += axes.theta;
        FWHM_N ++;
    }

    FWHM_X /= FWHM_N;
    FWHM_Y /= FWHM_N;
    FWHM_T /= FWHM_N;

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_MAJ",   PS_META_REPLACE, "PSF FWHM Major axis (mean)", FWHM_X);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_SG",   PS_META_REPLACE, "PSF FWHM Major axis (sigma)", 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_LQ",   PS_META_REPLACE, "PSF FWHM Major axis (lower quartile)", 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MJ_UQ",   PS_META_REPLACE, "PSF FWHM Major axis (upper quartile)", 0);

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_MIN",   PS_META_REPLACE, "PSF FWHM Minor axis (mean)", FWHM_Y);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_SG",   PS_META_REPLACE, "PSF FWHM Minor axis (sigma)", 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_LQ",   PS_META_REPLACE, "PSF FWHM Minor axis (lower quartile)", 0);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FW_MN_UQ",   PS_META_REPLACE, "PSF FWHM Minor axis (upper quartile)", 0);

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "ANGLE",    PS_META_REPLACE, "PSF angle",           FWHM_T);
    psMetadataAddS32 (recipe, PS_LIST_TAIL, "NPSFSTAR", PS_META_REPLACE, "Number of stars used to make PSF", 0);
    psMetadataAddBool(recipe, PS_LIST_TAIL, "PSFMODEL", PS_META_REPLACE, "Valid PSF Model?", false);

    return true;
}
