# include "psphotInternal.h"

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

    bool status;

    psTimerStart ("psphot");

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

    // array to store candidate PSF stars
    int NSTARS = psMetadataLookupS32 (&status, recipe, "PSF_MAX_NSTARS");
    PS_ASSERT (status, NULL);

    float PSF_MIN_DS = psMetadataLookupF32 (&status, recipe, "PSF_MIN_DS");
    PS_ASSERT (status, NULL);

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

    // use poissonian errors or local-sky errors
    bool POISSON_ERRORS = psMetadataLookupBool (&status, recipe, "POISSON_ERRORS");
    PS_ASSERT (status, NULL);

    // use poissonian errors or local-sky errors
    bool PSF_PARAM_WEIGHTS = psMetadataLookupBool (&status, recipe, "PSF_PARAM_WEIGHTS");
    PS_ASSERT (status, NULL);

    // how to model the PSF variations across the field
    psMetadata *md = psMetadataLookupMetadata (&status, recipe, "PSF.TREND.MASK");
    PS_ASSERT (status, NULL);

    // get the fixed PSF fit radius
    // XXX check that PSF_FIT_RADIUS < SKY_OUTER_RADIUS
    float RADIUS = psMetadataLookupF32 (&status, recipe, "PSF_FIT_RADIUS");
    PS_ASSERT (status, NULL);

    pmSourceFitModelInit (15, 0.01, PS_SQR(SKY_SIG), POISSON_ERRORS);

    psPolynomial2D *psfTrendMask = psPolynomial2DfromMetadata (md);
    if (!psfTrendMask) {
        psError(PSPHOT_ERR_PSF, true, "Unable to construct polynomial from PSF.TREND.MASK in the recipe");
        return NULL;
    }

    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;
            }
        }
    }
    psLogMsg ("psphot.pspsf", PS_LOG_DETAIL, "selected candidate %ld PSF objects\n", stars->n);

    if (stars->n == 0) {
        psLogMsg ("psphot.choosePSF", PS_LOG_WARN, "Failed to find any PSF candidates");
        psFree (stars);
        psFree (psfTrendMask);
        return NULL;
    }

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

    if (mdi->type == PS_DATA_STRING) {
        list = psListAlloc(NULL);
        psListAdd (list, PS_LIST_HEAD, mdi);
    } else {
        if (mdi->type != PS_DATA_METADATA_MULTI) psAbort("missing PSF_MODEL selection");
        list = psMemIncrRefCounter(mdi->data.list);
    }

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

    // try each model option listed in config
    psListIterator *iter = psListIteratorAlloc (list, PS_LIST_HEAD, FALSE);
    for (int i = 0; i < models->n; i++) {
        psMetadataItem *item = psListGetAndIncrement (iter);
        char *modelName = item->data.V;
        models->data[i] = pmPSFtryModel (stars, modelName, RADIUS, POISSON_ERRORS, psfTrendMask, PSF_PARAM_WEIGHTS, maskVal, mark);
    }

    psFree (iter);
    psFree (list);
    psFree (stars);
    psFree (psfTrendMask);

    // 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) {
            psError(PSPHOT_ERR_PSF, false, "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) {
        psError (PSPHOT_ERR_DATA, false, "Failed to fit any models when choosing PSF");
        psFree (models);
        return NULL;
    }

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

    // print/dump psf parameters
    if (psTraceGetLevel("psphot") >= 5) {
        for (int i = PM_PAR_SXX; i < try->psf->params_NEW->n; i++) {
            psPolynomial2D *poly = try->psf->params_NEW->data[i];
            for (int nx = 0; nx <= poly->nX; nx++) {
                for (int ny = 0; ny <= poly->nY; ny++) {
                    if (poly->mask[nx][ny]) continue;
                    fprintf (stderr, "%g x^%d y^%d\n", poly->coeff[nx][ny], nx, ny);
                }
            }
        }
    }

    // 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);
        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, RADIUS, "OR", PM_MASK_MARK);
            pmModelSub (source->pixels, source->maskObj, source->modelPSF, PM_MODEL_OP_FULL, maskVal);
            psImageKeepCircle (source->maskObj, x, y, RADIUS, "AND", PS_NOT_U8(PM_MASK_MARK));
        }

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

            // 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);
        psMetadata *psfData = pmPSFtoMetadata (NULL, try->psf);
        psMetadataConfigWrite (psfData, "psfmodel.dat");
        psFree (psfData);
        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(psf);
        return NULL;
    }

    char *modelName = pmModelGetType (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);

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

    pmModel *modelEXT = pmModelAlloc (psf->type);
    PS_ASSERT_PTR_NON_NULL(modelEXT, false);

    // make a model with unit central intensity at the image center
    modelEXT->params->data.F32[PM_PAR_SKY] = 0;
    modelEXT->params->data.F32[PM_PAR_I0] = 1;
    modelEXT->params->data.F32[PM_PAR_XPOS] = 0.5*image->numCols;
    modelEXT->params->data.F32[PM_PAR_YPOS] = 0.5*image->numRows;

    // construct a PSF model at this coordinate
    pmModel *modelPSF = pmModelFromPSF (modelEXT, psf);
    if (modelPSF == NULL) {
        psError(PSPHOT_ERR_PSF, false, "Failed to estimate PSF model at image centre");
        psFree(modelEXT);
        return false;
    }

    // get the correct model-radius function
    pmModelRadius modelRadiusFunc = pmModelRadius_GetFunction (psf->type);

    // get the model full-width at half-max
    psF64 FWHM_X = 2*modelRadiusFunc (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);

    psF64 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 (modelEXT);
    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;
}
