# include "psphotInternal.h"

# define SKIPSTAR(MSG) { psTrace ("psphot", 3, "invalid : %s", MSG); continue; }
// measure the aperture residual statistics and 2D variations

bool psphotApResid (pmConfig *config, pmReadout *readout, psArray *sources, pmPSF *psf)
{
    int Nfail = 0;
    int Nskip = 0;
    int Npsf;
    bool status;
    pmModel *model;
    pmSource *source;

    PS_ASSERT_PTR_NON_NULL(config, false);
    PS_ASSERT_PTR_NON_NULL(readout, false);
    PS_ASSERT_PTR_NON_NULL(sources, false);
    PS_ASSERT_PTR_NON_NULL(psf, false);

    psTimerStart ("psphot.apresid");

    // select the appropriate recipe information
    psMetadata *recipe  = psMetadataLookupPtr (&status, config->recipes, PSPHOT_RECIPE);
    assert (recipe);

    // determine the number of allowed threads
    int nThreads = psMetadataLookupS32(&status, config->arguments, "NTHREADS"); // Number of threads
    if (!status) {
	nThreads = 0;
    }
    // nThreads = 0; // XXX until testing is complete, do not thread this function

    bool measureAptrend = psMetadataLookupBool (&status, recipe, "MEASURE.APTREND");
    if (!measureAptrend) {
        // save nan values since these were not calculated
        psMetadataAdd (recipe, PS_LIST_TAIL, "SKYBIAS",  PS_DATA_F32 | PS_META_REPLACE, "aperture sky bias",   NAN);
        psMetadataAdd (recipe, PS_LIST_TAIL, "SKYSAT",   PS_DATA_F32 | PS_META_REPLACE, "aperture-determined saturation",   NAN);
        psMetadataAdd (recipe, PS_LIST_TAIL, "APMIFIT",  PS_DATA_F32 | PS_META_REPLACE, "aperture residual",   NAN);
        psMetadataAdd (recipe, PS_LIST_TAIL, "DAPMIFIT", PS_DATA_F32 | PS_META_REPLACE, "ap residual scatter", NAN);
        psMetadataAdd (recipe, PS_LIST_TAIL, "APLOSS",   PS_DATA_F32 | PS_META_REPLACE, "aperture loss (mag)", NAN);
        psMetadataAdd (recipe, PS_LIST_TAIL, "NAPMIFIT", PS_DATA_S32 | PS_META_REPLACE, "number of apresid stars", 0);
        return true;
    }

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

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

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

    // 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...
    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 APTREND_NSTAR_MIN = psMetadataLookupS32(&status, recipe, "APTREND.NSTAR.MIN");
    assert (status);

    // maximum order for aperture correction
    int APTREND_ORDER_MAX = psMetadataLookupS32(&status, recipe, "APTREND.ORDER.MAX");
    assert (status);

    if (APTREND_ORDER_MAX < 1) {
        psError(PSPHOT_ERR_CONFIG, true, "APTREND.ORDER.MAX must be 1 or more");
        return false;
    }

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

    // threaded measurement of the source magnitudes
    // choose Cx, Cy (see psphotThreadTools.c for overview of the concepts)
    int Cx = 1, Cy = 1;
    psphotChooseCellSizes (&Cx, &Cy, readout, nThreads);

    psArray *cellGroups = psphotAssignSources (Cx, Cy, sources);

    for (int i = 0; i < cellGroups->n; i++) {

	psArray *cells = cellGroups->data[i];

	for (int j = 0; j < cells->n; j++) {

	    // allocate a job -- if threads are not defined, this just runs the job
	    psThreadJob *job = psThreadJobAlloc ("PSPHOT_APRESID_MAGS");

	    psArrayAdd(job->args, 1, cells->data[j]); // sources
	    psArrayAdd(job->args, 1, psf);
	    PS_ARRAY_ADD_SCALAR(job->args, photMode, PS_TYPE_S32);
	    PS_ARRAY_ADD_SCALAR(job->args, maskVal,  PS_TYPE_IMAGE_MASK);

	    PS_ARRAY_ADD_SCALAR(job->args, 0,        PS_TYPE_S32); // this is used as a return value for Nskip
	    PS_ARRAY_ADD_SCALAR(job->args, 0,        PS_TYPE_S32); // this is used as a return value for Nfail

	    if (!psThreadJobAddPending(job)) {
		psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		psFree (job);
		return false;
	    }
	    psFree(job);

	}

	// wait for the threads to finish and manage results
	if (!psThreadPoolWait (false)) {
	    psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
	    return false;
	}

	// we have only supplied one type of job, so we can assume the types here
	psThreadJob *job = NULL;
	while ((job = psThreadJobGetDone()) != NULL) {
	    if (job->args->n < 1) {
		fprintf (stderr, "error with job\n");
	    } else {
		psScalar *scalar = NULL;
		scalar = job->args->data[4];
		Nskip += scalar->data.S32;
		scalar = job->args->data[5];
		Nfail += scalar->data.S32;
	    }
	    psFree(job);
	}
    }

    psFree (cellGroups);

    // gather the stats to assess the aperture residuals
    psVector *mask    = psVectorAllocEmpty (300, PS_TYPE_VECTOR_MASK);
    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;

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

        if (source->type != PM_SOURCE_TYPE_STAR) SKIPSTAR ("NOT STAR");
        if (source->mode &  PM_SOURCE_MODE_SATSTAR) SKIPSTAR ("SATSTAR");
        if (source->mode &  PM_SOURCE_MODE_BLEND) SKIPSTAR ("BLEND");
        if (source->mode &  PM_SOURCE_MODE_FAIL) SKIPSTAR ("FAIL STAR");
        if (source->mode &  PM_SOURCE_MODE_POOR) SKIPSTAR ("POOR STAR");

        if (!isfinite(source->apMag) || !isfinite(source->psfMag)) {
            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 ++;
            psTrace ("psphot", 3, "fail : bad dap %f %f", dap, MAX_AP_OFFSET);
            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.PS_TYPE_VECTOR_MASK_DATA[Npsf] = 0;

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

        psVectorExtend (mag,     100, 1);
        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 < APTREND_NSTAR_MIN) {
        psWarning("Only %d valid aperture residual sources (need %d), giving up", Npsf, APTREND_NSTAR_MIN);
	goto escape;
    }

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

    // Fit out the dap vs mag trend, iterate over spatial scale until error Floor increases.
    // Stop if Npsf / (Nx * Ny) < 3
    for (int i = 1; i <= APTREND_ORDER_MAX; 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) {
        psWarning("Failed to find a valid aperture residual value");
	goto escape;
    }

    // XXX catch error condition
    psphotApResidTrend (readout, psf, Npsf, entryMin, &errorScale, &errorFloor, mask, xPos, yPos, apResid, dMag);

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

    if (psTraceGetLevel("psphot") >= 2) {
        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.PS_TYPE_VECTOR_MASK_DATA[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-determined saturation",   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, "aperture loss (mag)", psf->growth->apLoss);
    psMetadataAdd (recipe, PS_LIST_TAIL, "NAPMIFIT", PS_DATA_S32 | PS_META_REPLACE, "number of apresid stars", psf->nApResid);

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

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

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

    psFree (dMagSys);
    psFree (dMag);

    psphotVisualPlotApResid (sources);

    return true;

escape:
    // save nan values since these were not calculated
    psMetadataAdd (recipe, PS_LIST_TAIL, "SKYBIAS",  PS_DATA_F32 | PS_META_REPLACE, "aperture sky bias",   NAN);
    psMetadataAdd (recipe, PS_LIST_TAIL, "SKYSAT",   PS_DATA_F32 | PS_META_REPLACE, "aperture-determined saturation",   NAN);
    psMetadataAdd (recipe, PS_LIST_TAIL, "APMIFIT",  PS_DATA_F32 | PS_META_REPLACE, "aperture residual",   NAN);
    psMetadataAdd (recipe, PS_LIST_TAIL, "DAPMIFIT", PS_DATA_F32 | PS_META_REPLACE, "ap residual scatter", NAN);
    psMetadataAdd (recipe, PS_LIST_TAIL, "APLOSS",   PS_DATA_F32 | PS_META_REPLACE, "aperture loss (mag)", NAN);
    psMetadataAdd (recipe, PS_LIST_TAIL, "NAPMIFIT", PS_DATA_S32 | PS_META_REPLACE, "number of apresid stars", 0);

    psFree (mag);
    psFree (mask);
    psFree (xPos);
    psFree (yPos);
    psFree (apResid);
    psFree (dMag);
    return false;
}

/*
  (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, psVector *mask, 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);
    psVector *mkSubset = psVectorAllocEmpty (nGroup, PS_TYPE_VECTOR_MASK);

    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];
            mkSubset->data.PS_TYPE_VECTOR_MASK_DATA[j] = mask->data.PS_TYPE_VECTOR_MASK_DATA[N];
        }
        dMSubset->n = j;
        dASubset->n = j;
        mkSubset->n = j;

        psStatsInit (statsS);
        psStatsInit (statsM);

        if (j > 2) {
            bool status = true;
            status &= psVectorStats (statsS, dASubset, NULL, mkSubset, 0xff);
            status &= psVectorStats (statsM, dMSubset, NULL, mkSubset, 0xff);
            if (!status) { psErrorClear (); }
            dSo->data.F32[i] = statsS->robustStdev;
            dMo->data.F32[i] = statsM->sampleMean;
            dRo->data.F32[i] = statsS->robustStdev / statsM->sampleMean;
            //      fprintf (stderr, "%d (%d) : sys: %f, phot: %f, rat: %f\n", i, j, dSo->data.F32[i], dMo->data.F32[i], dRo->data.F32[i]);
        } else {
            dSo->data.F32[i] = NAN;
            dMo->data.F32[i] = NAN;
            dRo->data.F32[i] = NAN;
        }
    }
    psFree (dMSubset);
    psFree (dASubset);
    psFree (mkSubset);

    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN);
    if (!psVectorStats (stats, dRo, NULL, NULL, 0)) {
        // XXX better testing of raised error
        psErrorClear();
    }

    *errorScale = stats->sampleMedian;
    for (int i = 0; i < dSo->n; i++) {
        *errorFloor = dSo->data.F32[i];
        if (isfinite(*errorFloor)) break;
    }

    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;
        float AR = readout->image->numRows / (float) readout->image->numCols;
        Ny = (int) (Nx * AR + 0.5);
        Ny = PS_MAX (1, Ny);
    } else {
        Ny = scale;
        float AR = readout->image->numRows / (float) readout->image->numCols;
        Nx = (int) (Ny * AR + 0.5);
        Nx = PS_MAX (1, Nx);
    }

    // require at least 10 stars per spatial bin
    if (Npsf < 10*Nx*Ny) {
        return false;
    }

    // the mask marks the values not used to calculate the ApTrend
    psVectorInit (mask, 0);

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

    // XXX somewhat arbitrary: soften the errors so the few bright stars do not totally dominate:
    psVector *dMagSoft = psVectorAlloc (dMag->n, PS_TYPE_F32);
    for (int i = 0; i < dMag->n; i++) {
        dMagSoft->data.F32[i] = hypot(dMag->data.F32[i], 0.01);
    }

    // XXX test for errors here
    pmTrend2DFit (psf->ApTrend, mask, 0xff, xPos, yPos, apResid, dMagSoft);

    // construct the fitted values and the residuals
    psVector *apResidFit = pmTrend2DEvalVector (psf->ApTrend, mask, 0xff, 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, mask, 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 (dMagSoft);
    psFree (apResidFit);
    psFree (apResidRes);

    return true;
}

bool psphotApResidMags_Threaded (psThreadJob *job) {

    int Nskip = 0;
    int Nfail = 0;

    psScalar *scalar = NULL;

    psArray *sources                = job->args->data[0];
    pmPSF *psf                      = job->args->data[1];
    pmSourcePhotometryMode photMode = PS_SCALAR_VALUE(job->args->data[2],S32);
    psImageMaskType maskVal         = PS_SCALAR_VALUE(job->args->data[3],PS_TYPE_IMAGE_MASK_DATA);

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

        if (source->type != PM_SOURCE_TYPE_STAR) SKIPSTAR ("NOT STAR");
        if (source->mode &  PM_SOURCE_MODE_SATSTAR) SKIPSTAR ("SATSTAR");
        if (source->mode &  PM_SOURCE_MODE_BLEND) SKIPSTAR ("BLEND");
        if (source->mode &  PM_SOURCE_MODE_FAIL) SKIPSTAR ("FAIL STAR");
        if (source->mode &  PM_SOURCE_MODE_POOR) SKIPSTAR ("POOR STAR");

        if (!pmSourceMagnitudes (source, psf, photMode, maskVal)) {
            Nskip ++;
            psTrace ("psphot", 3, "skip : bad source mag");
            continue;
        }

        if (!isfinite(source->apMag) || !isfinite(source->psfMag)) {
            Nfail ++;
            psTrace ("psphot", 3, "fail : nan mags : %f %f", source->apMag, source->psfMag);
            continue;
        }
    }

    // change the value of a scalar on the array (wrap this and put it in psArray.h)
    scalar = job->args->data[4];
    scalar->data.S32 = Nskip;

    scalar = job->args->data[5];
    scalar->data.S32 = Nfail;

    return true;
}
