# include "psphotInternal.h"

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

    bool status;

    psTimerStart ("psphot");

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

    float PSF_MIN_DS = psMetadataLookupF32 (&status, recipe, "PSF_MIN_DS");
    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);

    // 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;
            }
        }
    }
    // check that the identified psf stars sufficiently cover the region; if not, extend the
    // limits somewhat
    psphotCheckStarDistribution (stars, sources, options);

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

    // measure and save the median value of dparams[PM_PAR_SXX],dparams[PM_PAR_SYY]
    // these are used by psphotEvalPSF to identify extended sources
    // XXX is this still needed?
    psVector *Sx = psVectorAllocEmpty (try->sources->n, PS_TYPE_F32);
    psVector *Sy = psVectorAllocEmpty (try->sources->n, PS_TYPE_F32);
    psVector *dSx = psVectorAllocEmpty (try->sources->n, PS_TYPE_F32);
    psVector *dSy = psVectorAllocEmpty (try->sources->n, PS_TYPE_F32);
    psVector *dSN = psVectorAllocEmpty (try->sources->n, PS_TYPE_F32);
    for (int i = 0; i < try->sources->n; i++) {
        // masked for: bad model fit, outlier in parameters
        if (try->mask->data.U8[i] & PSFTRY_MASK_ALL)
            continue;

        pmSource *source = try->sources->data[i];
        Sx->data.F32[Sx->n] = source->modelPSF->params->data.F32[PM_PAR_SXX];
        Sy->data.F32[Sy->n] = source->modelPSF->params->data.F32[PM_PAR_SYY];
        dSN->data.F32[dSN->n] = source->modelPSF->dparams->data.F32[PM_PAR_I0] / source->modelPSF->params->data.F32[PM_PAR_I0];
        dSx->data.F32[dSx->n] = source->modelPSF->dparams->data.F32[PM_PAR_SXX];
        dSy->data.F32[dSy->n] = source->modelPSF->dparams->data.F32[PM_PAR_SYY];
        Sx->n ++;
        Sy->n ++;
        dSN->n ++;
        dSx->n ++;
        dSy->n ++;
    }
    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);

    psVectorStats (stats, dSx, NULL, NULL, 0);
    float dSxo = stats->sampleMean;
    psLogMsg ("psphot.choosePSF", PS_LOG_INFO, "dSxo: mean: %f, median: %f, stdev: %f, npts: %ld", stats->sampleMean, stats->sampleMedian, stats->sampleStdev, dSx->n);

    psVectorStats (stats, dSy, NULL, NULL, 0);
    float dSyo = stats->sampleMean;
    psLogMsg ("psphot.choosePSF", PS_LOG_INFO, "dSyo: mean: %f, median: %f, stdev: %f, npts: %ld", stats->sampleMean, stats->sampleMedian, stats->sampleStdev, dSy->n);

    for (int i = 0; i < dSN->n; i++) {
        dSx->data.F32[i] = (dSx->data.F32[i] - dSxo) / PS_MAX (PSF_MIN_DS, Sx->data.F32[i]*dSN->data.F32[i]);
        dSy->data.F32[i] = (dSy->data.F32[i] - dSyo) / PS_MAX (PSF_MIN_DS, Sy->data.F32[i]*dSN->data.F32[i]);
    }

    psVectorStats (stats, dSx, NULL, NULL, 0);
    float nSx = stats->sampleStdev;
    psLogMsg ("psphot.choosePSF", PS_LOG_INFO, "ndSx: mean: %f, median: %f, stdev: %f, npts: %ld", stats->sampleMean, stats->sampleMedian, stats->sampleStdev, dSx->n);
    psVectorStats (stats, dSy, NULL, NULL, 0);
    float nSy = stats->sampleStdev;
    psLogMsg ("psphot.choosePSF", PS_LOG_INFO, "ndSy: mean: %f, median: %f, stdev: %f, npts: %ld", stats->sampleMean, stats->sampleMedian, stats->sampleStdev, dSy->n);

    psFree (Sx);
    psFree (Sy);
    psFree (dSx);
    psFree (dSy);
    psFree (dSN);
    psFree (stats);

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "DSX_MEAN", PS_META_REPLACE, "mean offset of dSXX", dSxo);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "DSY_MEAN", PS_META_REPLACE, "mean offset of dSYY", dSyo);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "DSX_STDV", PS_META_REPLACE, "stdev of mean-corrected dSXX", nSx);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "DSY_STDV", PS_META_REPLACE, "stdev of mean-corrected dSYY", nSy);

    // 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 a PSF residual image
    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;
    }

    // XXX test dump of psf star data and psf-subtracted image
    if (psTraceGetLevel("psphot.psfstars") > 5) {
        psphotSaveImage (NULL, readout->image,  "rawstars.fits");

        for (int i = 0; i < try->sources->n; i++) {
            // masked for: bad model fit, outlier in parameters
            if (try->mask->data.U8[i] & PSFTRY_MASK_ALL)
                continue;

            pmSource *source = try->sources->data[i];
            float x = source->modelPSF->params->data.F32[PM_PAR_XPOS];
            float y = source->modelPSF->params->data.F32[PM_PAR_YPOS];

            // set the mask and subtract the PSF model
            // XXX should we be using maskObj? should we be unsetting the mask?
            // use pmModelSub because modelFlux has not been generated
            assert (source->maskObj);
            psImageKeepCircle (source->maskObj, x, y, options->radius, "OR", markVal);
            pmModelSub (source->pixels, source->maskObj, source->modelPSF, PM_MODEL_OP_FULL, maskVal);
            psImageKeepCircle (source->maskObj, x, y, options->radius, "AND", PS_NOT_U8(markVal));
        }

        FILE *f = fopen ("shapes.dat", "w");
        for (int i = 0; i < try->sources->n; i++) {
            psF32 inPar[10];  // must be psF32 to pmPSF_FitToModel

            // masked for: bad model fit, outlier in parameters
            if (try->mask->data.U8[i] & PSFTRY_MASK_ALL) continue;

            pmSource *source = try->sources->data[i];
            psF32 *outPar = source->modelEXT->params->data.F32;

            psEllipseShape shape;

            shape.sx  = outPar[PM_PAR_SXX] / M_SQRT2;
            shape.sy  = outPar[PM_PAR_SYY] / M_SQRT2;
            shape.sxy = outPar[PM_PAR_SXY];

            psEllipsePol pol = pmPSF_ModelToFit (outPar);
            inPar[PM_PAR_E0] = pol.e0;
            inPar[PM_PAR_E1] = pol.e1;
            inPar[PM_PAR_E2] = pol.e2;
            pmPSF_FitToModel (inPar, 0.1);

            psEllipseAxes axes1 = psEllipseShapeToAxes (shape, 20.0);
            psEllipseAxes axes2;
            (void)psEllipsePolToAxes(pol, 0.1, &axes2);
            psEllipsePol pol2 = psEllipseAxesToPol (axes1);

            fprintf (f, "%3d  %7.2f %7.2f  %7.4f %7.4f %7.4f  --  %7.4f %7.4f %7.4f  :  %7.4f %7.4f %7.4f  --  %7.4f %7.4f %7.4f : %7.4f %7.4f %6.1f : %7.4f %7.4f %6.1f\n",
                     i, outPar[PM_PAR_XPOS], outPar[PM_PAR_YPOS],
                     outPar[PM_PAR_SXX], outPar[PM_PAR_SXY], outPar[PM_PAR_SYY],
                     pol.e0, pol.e1, pol.e2,
                     pol2.e0, pol2.e1, pol2.e2,
                     inPar[PM_PAR_SXX], inPar[PM_PAR_SXY], inPar[PM_PAR_SYY],
                     axes1.major, axes1.minor, axes1.theta*PM_DEG_RAD,
                     axes2.major, axes2.minor, axes2.theta*PM_DEG_RAD
                     );
        }
        fclose (f);

        psphotSaveImage (NULL, readout->image,  "psfstars.fits");
        pmSourcesWritePSFs (try->sources, "psfstars.dat");
        pmSourcesWriteEXTs (try->sources, "extstars.dat", false);
        // XXX need alternative output function
        // psMetadata *psfData = pmPSFtoMetadata (NULL, try->psf);
        // psMetadataConfigWrite (psfData, "psfmodel.dat");
        psLogMsg ("psphot.choosePSF", PS_LOG_INFO, "wrote out psf-subtracted image, psf data, exiting\n");
        exit (0);
    }

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

    char *modelName = pmModelClassGetName (psf->type);
    psLogMsg ("psphot.pspsf", PS_LOG_INFO, "select psf model: %f sec\n", psTimerMark ("psphot"));
    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);

    // construct a normalized PSF model at this coordinate (Io = 1.0)
    pmModel *modelPSF = pmModelFromPSFforXY (psf, 0.5*image->numCols, 0.5*image->numRows, 1.0);
    if (modelPSF == NULL) {
        psError(PSPHOT_ERR_PSF, false, "Failed to estimate PSF model at image centre");
        return false;
    }

    // get the model full-width at half-max
    double FWHM_X = 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);

    double FWHM_Y = FWHM_X * (axes.minor / axes.major);

    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_X",   PS_META_REPLACE, "PSF FWHM Major axis", FWHM_X);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_Y",   PS_META_REPLACE, "PSF FWHM Minor axis", FWHM_Y);
    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 (modelPSF);

    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;

        // moments->Sx,Sy are roughly sigma_x,y
        moments.x2 = PS_SQR(source->moments->Sx);
        moments.y2 = PS_SQR(source->moments->Sy);
        moments.xy = source->moments->Sxy;

        // 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_X",   PS_META_REPLACE, "PSF FWHM Major axis (from moments)", FWHM_X);
    psMetadataAddF32 (recipe, PS_LIST_TAIL, "FWHM_Y",   PS_META_REPLACE, "PSF FWHM Minor axis (from moments)", FWHM_Y);
    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;
}
