# include "psphotInternal.h"

// measure the aperture residual statistics and 2D variations
bool psphotApResid (pmReadout *readout, psArray *sources, psMetadata *recipe, pmPSF *psf, psMaskType maskVal, psMaskType mark)
{
    int Nfail = 0;
    int Nskip = 0;
    int Npsf;
    bool status;
    pmModel *model;
    pmSource *source;

    PS_ASSERT_PTR_NON_NULL(psf, false);
    // XXX drop this? PS_ASSERT_PTR_NON_NULL(psf->ApTrend, false);
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(sources, false);
    PS_ASSERT_PTR_NON_NULL(recipe, false);

    psTimerStart ("psphot");

    // S/N limit to perform full non-linear fits
    float MAX_AP_OFFSET = psMetadataLookupF32 (&status, recipe, "MAX_AP_OFFSET");

    // this is the smallest radius allowed: need to at least extend growth curve down to this...
    float PSF_FIT_PAD   = psMetadataLookupF32 (&status, recipe, "PSF_FIT_PADDING");
    bool IGNORE_GROWTH = psMetadataLookupBool (&status, recipe, "IGNORE_GROWTH");
    bool INTERPOLATE_AP = psMetadataLookupBool (&status, recipe, "INTERPOLATE_AP");

    // XXX is this still needed?  the pmTrend2D stuff should be auto-adjusting...
    int NSTAR_APERTURE_CORRECTION_MIN = psMetadataLookupS32(&status, recipe, "NSTAR_APERTURE_CORRECTION_MIN");
    if (!status) {
        NSTAR_APERTURE_CORRECTION_MIN = 5;
    }

    pmSourcePhotometryMode photMode = 0;
    if (!IGNORE_GROWTH) photMode |= PM_SOURCE_PHOT_GROWTH;
    if (INTERPOLATE_AP) photMode |= PM_SOURCE_PHOT_INTERP;

    // set limits on the aperture magnitudes
    pmSourceMagnitudesInit (recipe);

    // measure the aperture loss as a function of radius for PSF
    float REF_RADIUS = psMetadataLookupF32 (&status, recipe, "PSF_REF_RADIUS");
    psf->growth = pmGrowthCurveAlloc (PSF_FIT_PAD, 100.0, REF_RADIUS);

    if (!pmGrowthCurveGenerate (readout, psf, IGNORE_GROWTH, maskVal, mark)) {
        psError(PSPHOT_ERR_APERTURE, false, "Fitting aperture corrections");
        psFree(psf->growth); psf->growth = NULL;
        return false;
    }

    psVector *mask    = psVectorAllocEmpty (300, PS_TYPE_U8);
    psVector *mag     = psVectorAllocEmpty (300, PS_TYPE_F32);
    psVector *xPos    = psVectorAllocEmpty (300, PS_TYPE_F32);
    psVector *yPos    = psVectorAllocEmpty (300, PS_TYPE_F32);
    psVector *apResid = psVectorAllocEmpty (300, PS_TYPE_F32);
    psVector *dMag    = psVectorAllocEmpty (300, PS_TYPE_F32);
    Npsf = 0;

    // select all good PM_SOURCE_TYPE_STAR entries
    for (int i = 0; i < sources->n; i++) {
        source = sources->data[i];
        model = source->modelPSF;

        if (source->type != PM_SOURCE_TYPE_STAR) continue;
        if (source->mode &  PM_SOURCE_MODE_SATSTAR) continue;
        if (source->mode &  PM_SOURCE_MODE_BLEND) continue;
        if (source->mode &  PM_SOURCE_MODE_FAIL) continue;
        if (source->mode &  PM_SOURCE_MODE_POOR) continue;

        // get growth-corrected, apTrend-uncorrected magnitudes in scaled apertures
        // will fail if below S/N threshold or model is missing
        if (!pmSourceMagnitudes (source, psf, photMode, maskVal, mark)) {
            Nskip ++;
            continue;
        }

        if (!isfinite(source->apMag) || !isfinite(source->apMag)) {
            Nfail ++;
            continue;
        }

	// aperture residual for this source
	float dap = source->apMag - source->psfMag;

        // sanity clipping : if the model is very discrepant, put your expectation in the recipe
        if ((MAX_AP_OFFSET > 0) && (fabs(dap) > MAX_AP_OFFSET)) {
            Nfail ++;
            continue;
        }

	mag->data.F32[Npsf]     = source->psfMag;
        apResid->data.F32[Npsf] = dap;
        xPos->data.F32[Npsf]    = model->params->data.F32[PM_PAR_XPOS];
        yPos->data.F32[Npsf]    = model->params->data.F32[PM_PAR_YPOS];

        mask->data.U8[Npsf] = 0;

        dMag->data.F32[Npsf] = model->dparams->data.F32[PM_PAR_I0] / model->params->data.F32[PM_PAR_I0];

        psVectorExtend (mask,    100, 1);
        psVectorExtend (xPos,    100, 1);
        psVectorExtend (yPos,    100, 1);
        psVectorExtend (dMag,    100, 1);
        psVectorExtend (apResid, 100, 1);
        Npsf ++;
    }
    
    psLogMsg ("psphot.apresid", PS_LOG_DETAIL, "measure aperture residuals for %d objects (%d skipped, %d failed, %ld invalid)\n",
              Npsf, Nskip, Nfail, sources->n - Npsf - Nskip - Nfail);

    // XXX choose a better value here?
    if (Npsf < NSTAR_APERTURE_CORRECTION_MIN) {
        psError(PSPHOT_ERR_APERTURE, true, "Only %d valid aperture residual sources (need %d), giving up",
                Npsf, NSTAR_APERTURE_CORRECTION_MIN);
        return false;
    }

    // XXX deprecating the old code which allowed the ApResid to be fitted as a function of flux and r^2/flux
    // XXX is this asymmetric clipping still needed?  this analysis should come after neighbors are subtracted...
    // 3hi/1lo sigma clipping on the rflux vs metric fit
    // systematic error information
    float errorScale = 0.0;
    float errorFloor = 0.0;

    float errorFloorMin = FLT_MAX;
    int entryMin = -1;

    // *** iterate over spatial scale until error Floor increases
    // *** fit out the dap vs mag trend?
    // *** stop if Npsf / (Nx * Ny) < 3
    for (int i = 1; i < 10; i++) {

	if (!psphotApResidTrend (readout, psf, Npsf, i, &errorScale, &errorFloor, mask, xPos, yPos, apResid, dMag)) {
	    break;
	}

	// store the resulting errorFloor values and the scales, redo the best
	if (errorFloor < errorFloorMin) {
	    errorFloorMin = errorFloor;
	    entryMin = i;
	}
    }
    if (entryMin == -1) {
	psAbort ("failed on ApResid Trend");
    }

    psphotApResidTrend (readout, psf, Npsf, entryMin, &errorScale, &errorFloor, mask, xPos, yPos, apResid, dMag);

    // construct the fitted values and the residuals
    psVector *apResidFit = pmTrend2DEvalVector (psf->ApTrend, xPos, yPos);
    psVector *apResidRes = (psVector *) psBinaryOp (NULL, (void *) apResid, "-", (void *) apResidFit);
    psVector *dMagSys = (psVector *) psBinaryOp (NULL, (void *) dMag, "*", (void *) psScalarAlloc(errorScale, PS_TYPE_F32));

    psphotSaveImage (NULL, psf->ApTrend->map->map, "apresid.map.fits");

    if (psTraceGetLevel("psphot") >= 5) {
	FILE *dumpFile = fopen ("apresid.dat", "w");
	for (int i = 0; i < xPos->n; i++) {
	    fprintf (dumpFile, "%f %f  %f %f %f  %f %f  %x\n", 
		     xPos->data.F32[i], yPos->data.F32[i], 
		     mag->data.F32[i], dMag->data.F32[i], dMagSys->data.F32[i],
		     apResid->data.F32[i], apResidRes->data.F32[i],
		     mask->data.U8[i]);
	}
	fclose (dumpFile);
    }

    // apply ApTrend results
    float xc = 0.5*readout->image->numCols + readout->image->col0 + 0.5;
    float yc = 0.5*readout->image->numRows + readout->image->row0 + 0.5;
    psf->ApResid  = pmTrend2DEval (psf->ApTrend, xc, yc); // ap-fit at chip center
    psf->dApResid = errorFloor;
    psf->nApResid = Npsf;

    // save results for later output
    psMetadataAdd (recipe, PS_LIST_TAIL, "SKYBIAS",  PS_DATA_F32 | PS_META_REPLACE, "aperture sky bias",   0.0);
    psMetadataAdd (recipe, PS_LIST_TAIL, "SKYSAT",   PS_DATA_F32 | PS_META_REPLACE, "aperture sky bias",   0.0);
    psMetadataAdd (recipe, PS_LIST_TAIL, "APMIFIT",  PS_DATA_F32 | PS_META_REPLACE, "aperture residual",   psf->ApResid);
    psMetadataAdd (recipe, PS_LIST_TAIL, "DAPMIFIT", PS_DATA_F32 | PS_META_REPLACE, "ap residual scatter", psf->dApResid);
    psMetadataAdd (recipe, PS_LIST_TAIL, "APLOSS",   PS_DATA_F32 | PS_META_REPLACE, "ap residual scatter", psf->growth->apLoss);
    psMetadataAdd (recipe, PS_LIST_TAIL, "NAPMIFIT", PS_DATA_S32 | PS_META_REPLACE, "ap residual scatter", psf->nApResid);

    // psLogMsg ("psphot.apresid", PS_LOG_INFO, "measure full-frame aperture residuals for %d of %d objects: %f sec\n", Nkeep, Npsf, psTimerMark ("psphot"));
    psLogMsg ("psphot.apresid", PS_LOG_DETAIL, "aperture residual: %f +/- %f\n", psf->ApResid, psf->dApResid);

    psFree (mag);
    psFree (mask);
    psFree (xPos);
    psFree (yPos);

    psFree (apResid);
    psFree (apResidFit);
    psFree (apResidRes);

    psFree (dMagSys);
    psFree (dMag);

    return true;
}

/*
  (aprMag' - fitMag) = rflux*skyBias + ApTrend(x,y)
  (aprMag - rflux*skyBias) - fitMag = ApTrend(x,y)
  (aprMag - rflux*skyBias) = fitMag + ApTrend(x,y)
*/

bool psphotMagErrorScale (float *errorScale, float *errorFloor, psVector *dMag, psVector *dap, int nGroup) {

    psStats *statsS = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    psStats *statsM = psStatsAlloc (PS_STAT_SAMPLE_MEAN);

    // measure the trend in bins with 10 values each; if < 10 total, use them all
    int nBin = PS_MAX (dMag->n / nGroup, 1);

    // output vectors for ApResid trend 
    psVector *dSo = psVectorAlloc (nBin, PS_TYPE_F32);
    psVector *dMo = psVectorAlloc (nBin, PS_TYPE_F32);
    psVector *dRo = psVectorAlloc (nBin, PS_TYPE_F32);

    // use dMag to group the dMag and dap vectors
    psVector *index = psVectorSortIndex (NULL, dMag);

    // subset vectors for dMag and dap values within the given range
    psVector *dMSubset = psVectorAllocEmpty (nGroup, PS_TYPE_F32);
    psVector *dASubset = psVectorAllocEmpty (nGroup, PS_TYPE_F32);

    int n = 0;
    for (int i = 0; i < dMo->n; i++) {
	int j;
	for (j = 0; (j < nGroup) && (n < dMag->n); j++, n++) {
	    int N = index->data.U32[n];
	    dMSubset->data.F32[j] = dMag->data.F32[N];
	    dASubset->data.F32[j] = dap->data.F32[N];
	}
	dMSubset->n = j;
	dASubset->n = j;

	psStatsInit (statsS);
	psStatsInit (statsM);

	psVectorStats (statsS, dASubset, NULL, NULL, 0xff);
	psVectorStats (statsM, dMSubset, NULL, NULL, 0xff);

	dSo->data.F32[i] = statsS->robustStdev;
	dMo->data.F32[i] = statsM->sampleMean;

	dRo->data.F32[i] = statsS->robustStdev / statsM->sampleMean;
    }
    psFree (dMSubset);
    psFree (dASubset);
    
    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN);
    psVectorStats (stats, dRo, NULL, NULL, 0);

    *errorScale = stats->sampleMedian;
    *errorFloor = dSo->data.F32[0];

    psFree (stats);
    psFree (index);

    psFree (dRo);
    psFree (dMo);
    psFree (dSo);
    
    psFree (statsS);
    psFree (statsM);

    return true;
}

bool psphotApResidTrend (pmReadout *readout, pmPSF *psf, int Npsf, int scale, float *errorScale, float *errorFloor, psVector *mask, psVector *xPos, psVector *yPos, psVector *apResid, psVector *dMag) {

    int Nx, Ny;
	
    if (readout->image->numCols > readout->image->numRows) {
	Nx = scale;
	Ny = PS_MAX (1, Nx * (readout->image->numRows / readout->image->numCols));
    } else {
	Ny = scale;
	Nx = PS_MAX (1, Ny * (readout->image->numCols / readout->image->numRows));
    }

    if (Npsf < 3*Nx*Ny) {
	return false;
    }

    // XXX stats structure for use by ApTrend : make parameters user setable
    psStats *stats = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    stats->min = 2.0;
    stats->max = 3.0;

    // measure Trend2D for the current spatial scale
    psFree (psf->ApTrend);
    psf->ApTrend = pmTrend2DAlloc (PM_TREND_MAP, readout->image, Nx, Ny, stats);
    pmTrend2DFit (psf->ApTrend, mask, 0xff, xPos, yPos, apResid, dMag);
    
    // construct the fitted values and the residuals
    psVector *apResidFit = pmTrend2DEvalVector (psf->ApTrend, xPos, yPos);
    psVector *apResidRes = (psVector *) psBinaryOp (NULL, (void *) apResid, "-", (void *) apResidFit);

    // measure systematic errorFloor & systematic / photon scale factor
    // XXX this is a bit arbitrary, but it forces ~3 stars from the bright bin per spatial bin
    int nGroup = PS_MAX (3*Nx*Ny, 10);
    psphotMagErrorScale (errorScale, errorFloor, dMag, apResidRes, nGroup);

    psLogMsg ("psphot.apresid", PS_LOG_INFO, "result of %d x %d grid (%d stars per bin)\n", Nx, Ny, nGroup);
    psLogMsg ("psphot.apresid", PS_LOG_INFO, "systematic error / photon error: %f\n", *errorScale);
    psLogMsg ("psphot.apresid", PS_LOG_INFO, "systematic scatter floor: %f\n", *errorFloor);

    psFree (stats);
    psFree (apResidFit);
    psFree (apResidRes);

    return true;
}
    
