# include "psphotInternal.h"

bool psphotMakeResiduals (psArray *sources, psMetadata *recipe, pmPSF *psf, psMaskType maskVal) {

    bool status, isPSF;
    double flux, dflux;
    psU8 mflux;

    psTimerStart ("residuals");

    if (!psMetadataLookupBool(&status, recipe, "PSF.RESIDUALS")) return true;

    int SPATIAL_ORDER = psMetadataLookupS32(&status, recipe, "PSF.RESIDUALS.SPATIAL_ORDER");
    PS_ASSERT (status, false);
    if (SPATIAL_ORDER != 0 && SPATIAL_ORDER != 1) {
        psError(PSPHOT_ERR_CONFIG, true, "PSF.RESIDUALS.SPATIAL_ORDER must be 0 or 1 (not %d)",
                SPATIAL_ORDER);
        return false;
    }

    int xBin = psMetadataLookupS32(&status, recipe, "PSF.RESIDUALS.XBIN");
    PS_ASSERT (status, false);
    PS_ASSERT (xBin != 0, false);

    int yBin = psMetadataLookupS32(&status, recipe, "PSF.RESIDUALS.YBIN");
    PS_ASSERT (status, false);
    PS_ASSERT (yBin != 0, false);

    float nSigma = psMetadataLookupF32(&status, recipe, "PSF.RESIDUALS.NSIGMA");
    PS_ASSERT (status, false);

    float pixelSN = psMetadataLookupF32(&status, recipe, "PSF.RESIDUALS.PIX.SN");
    PS_ASSERT (status, false);

    char *modeString = psMetadataLookupStr(&status, recipe, "PSF.RESIDUALS.INTERPOLATION");
    PS_ASSERT (status, false);

    psImageInterpolateMode mode = psImageInterpolateModeFromString (modeString);
    if (mode == PS_INTERPOLATE_NONE) {
        psError(PSPHOT_ERR_CONFIG, false, "invalid interpolation in psphot.config");
        return false;
    }

    char *statString = psMetadataLookupStr(&status, recipe, "PSF.RESIDUALS.STATISTIC");
    PS_ASSERT (status, false);

    psStatsOptions statOption = psStatsOptionFromString (statString);
    if (!statOption) {
        psError(PSPHOT_ERR_CONFIG, false, "invalid residual statistic in psphot.config");
        return false;
    }

    // user parameters:
    // size of aperture (determine from source images?)
    // binning factor

    // select the subset of sources which are the PSFSTARs
    // for each input source:
    // - construct a residual image, renormalized
    // - construct a renormalized weight image
    // - construct a new mask image

    // construct the output residual table (Nx*DX,Ny*DY)
    // for each output pixel:
    // - construct a histogram of the values & weights (interpolate to the common pixel coordinate)
    // - measure the robust median & sigma
    // - reject (mask) input pixels which are outliers
    // - re-measure the robust median & sigma
    // - set output pixel, weight, and mask

    const int badMask = 1;              // mask bits
    const int poorMask = 2;             //       from psImageInterpolate
    const int clippedMask = 4;          // mask bit set for clipped values

    bool offImage = false;              // pixel is off the image

    // determine the maximum image size from the input sources
    int xSize = 0;
    int ySize = 0;

    psVector *xC = psVectorAllocEmpty (100, PS_TYPE_F32);
    psVector *yC = psVectorAllocEmpty (100, PS_TYPE_F32);

    // build (DATA - MODEL) [an image] for each psf star
    psArray *input = psArrayAllocEmpty (100);
    for (int i = 0; i < sources->n; i++) {

        pmSource *source = sources->data[i];

        if (!(source->mode & PM_SOURCE_MODE_PSFSTAR)) continue;

        // which model to use?
        pmModel *model = pmSourceGetModel (&isPSF, source);
        if (model == NULL) continue;  // model must be defined

        psImage *image  = psImageCopy (NULL, source->pixels,   PS_TYPE_F32);
        psImage *weight = psImageCopy (NULL, source->weight,   PS_TYPE_F32);
        psImage *mask   = psImageCopy (NULL, source->maskView, PS_TYPE_U8);
        pmModelSub (image, mask, model, PM_MODEL_OP_FUNC, maskVal);

        // re-normalize image and weight
        float Io = model->params->data.F32[PM_PAR_I0];
        psBinaryOp (image, image, "/", psScalarAlloc(Io, PS_TYPE_F32));
        psBinaryOp (weight, weight, "/", psScalarAlloc(Io*Io, PS_TYPE_F32));

        // we will interpolate the image and weight - include the mask or not?
        // XXX consider better values for the mask bits
        psImageInterpolateOptions *interp =
            psImageInterpolateOptionsAlloc(mode, image, weight, NULL, 0xff, 0.0, 0.0, badMask, poorMask, 0.0);
        psArrayAdd (input,  100, interp);

        // save the X,Y position for future reference
        xC->data.F32[xC->n] = model->params->data.F32[PM_PAR_XPOS];
        yC->data.F32[yC->n] = model->params->data.F32[PM_PAR_YPOS];
        psVectorExtend (xC, 100, 1);
        psVectorExtend (yC, 100, 1);

        xSize = PS_MAX (xSize, image->numCols);
        ySize = PS_MAX (ySize, image->numRows);

        // free up the excess references
        psFree (mask);
        psFree (image);
        psFree (weight);
        psFree (interp);
    }
    pmResiduals *resid = pmResidualsAlloc (xSize, ySize, xBin, yBin);

    // x(resid) = (x(image) - Xo)*xBin + xCenter

    psVector *fluxes  = psVectorAlloc (input->n, PS_TYPE_F32);
    psVector *dfluxes = psVectorAlloc (input->n, PS_TYPE_F32);
    psVector *fmasks  = psVectorAlloc (input->n, PS_TYPE_U8);

    // statistic to use to determine baseline for clipping
    psStats *fluxClip     = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    psStats *fluxClipDef  = psStatsAlloc (PS_STAT_ROBUST_MEDIAN | PS_STAT_ROBUST_STDEV);
    // statistic to use to determine output flux
    // XXX make API to convert statOption for MEAN/MEDIAN in to corresponding STDEV?
    psStats *fluxStats    = psStatsAlloc (statOption | PS_STAT_SAMPLE_STDEV);
    psStats *fluxStatsDef = psStatsAlloc (statOption | PS_STAT_SAMPLE_STDEV);

    // Use psF64 to minimize overflow problems?
    psImage *A = psImageAlloc(3, 3, PS_TYPE_F64); // Least-squares matrix
    psVector *B = psVectorAlloc(3, PS_TYPE_F64); // Least-squares vector

    // Solve MODEL = R + x R_x + y R_y in pixel-by-pixel a least-squares sense
    // (If SPATIAL_ORDER == 0, just solve MODEL = R)
    for (int oy = 0; resid != NULL && oy < resid->Ro->numRows; oy++) {
        for (int ox = 0; ox < resid->Ro->numCols; ox++) {

            // build the vector of data values for this output pixel
            for (int i = 0; i < input->n; i++) {

                psImageInterpolateOptions *interp = input->data[i];

                // fractional image position
                float ix = (ox + 0.5 - resid->xCenter) / (float) xBin + xC->data.F32[i] - interp->image->col0;
                float iy = (oy + 0.5 - resid->yCenter) / (float) yBin + yC->data.F32[i] - interp->image->row0;

                mflux = 0;
                offImage = false;
                if (psImageInterpolate (&flux, &dflux, &mflux, ix, iy, interp) == PS_INTERPOLATE_STATUS_OFF) {
                    // This pixel is off the image
                    offImage = true;
                }
                fluxes->data.F32[i] = flux;
                dfluxes->data.F32[i] = dflux;
                fmasks->data.U8[i] = mflux;
		if (isnan(flux)) { 
		    fmasks->data.U8[i] = 1;
		}
                // fprintf (stderr, "%f %f : %f %f (%d)\n", ix, iy, flux, dflux, fmasks->data.U8[i]);
            }

            // measure the robust median to determine a baseline reference value
            *fluxClip = *fluxClipDef;
            psVectorStats (fluxClip, fluxes, NULL, fmasks, 0xff);
            psErrorClear();             // clear (ignore) any outstanding errors

            // mark input pixels which are more than N sigma from the median
            for (int i = 0; i < fluxes->n; i++) {
                float delta = fluxes->data.F32[i] - fluxClip->robustMedian;
                float sigma = sqrt (dfluxes->data.F32[i]);
                float swing = fabs(delta) / sigma;

                // make this a user option
                if (swing > nSigma) {
                    fmasks->data.U8[i] = clippedMask;
                }
            }

            if (offImage || SPATIAL_ORDER == 0) {
                // measure the desired statistic on the unclipped pixels
                *fluxStats = *fluxStatsDef;
                psVectorStats (fluxStats, fluxes, NULL, fmasks, 0xff);
                psErrorClear();         // clear (ignore) any outstanding errors

                resid->Ro->data.F32[oy][ox] = psStatsGetValue(fluxStats, statOption);
                resid->Rx->data.F32[oy][ox] = resid->Ry->data.F32[oy][ox] = 0.0;
                //resid->weight->data.F32[oy][ox] = fluxStats->sampleStdev;

		if (resid->Ro->data.F32[oy][ox] < pixelSN*fluxStats->sampleStdev) {
		  resid->mask->data.U8[oy][ox] = 1;
		}

            } else {
                assert (SPATIAL_ORDER == 1);
                psImageInit(A, 0.0);
                psVectorInit(B, 0.0);
                for (int i = 0; i < fluxes->n; i++) {
                    if (fmasks->data.U8[i]) continue;
                    B->data.F64[0] += fluxes->data.F32[i]/dfluxes->data.F32[i];
                    B->data.F64[1] += fluxes->data.F32[i]*xC->data.F32[i]/dfluxes->data.F32[i];
                    B->data.F64[2] += fluxes->data.F32[i]*yC->data.F32[i]/dfluxes->data.F32[i];

                    A->data.F64[0][0] += 1.0/dfluxes->data.F32[i];
                    A->data.F64[1][0] += xC->data.F32[i]/dfluxes->data.F32[i];
                    A->data.F64[2][0] += yC->data.F32[i]/dfluxes->data.F32[i];

                    A->data.F64[1][1] += PS_SQR(xC->data.F32[i])/dfluxes->data.F32[i];
                    A->data.F64[2][2] += PS_SQR(yC->data.F32[i])/dfluxes->data.F32[i];
                    A->data.F64[1][2] += xC->data.F32[i]*yC->data.F32[i]/dfluxes->data.F32[i];
                }

                A->data.F64[0][1] = A->data.F64[1][0];
                A->data.F64[0][2] = A->data.F64[2][0];
                A->data.F64[2][1] = A->data.F64[1][2];

                if (!psMatrixGJSolve(A, B)) {
                    psError(PSPHOT_ERR_PSF, false, "Singular matrix solving for (y,x) = (%d,%d)'s residuals",
                            oy, ox);
                    psFree(resid); resid = NULL;
                    break;
                }

                resid->Ro->data.F32[oy][ox] = B->data.F64[0];
                resid->Rx->data.F32[oy][ox] = B->data.F64[1];
                resid->Ry->data.F32[oy][ox] = B->data.F64[2];

		float dRo = sqrt(A->data.F32[0][0]);
		if (resid->Ro->data.F32[oy][ox] < pixelSN*dRo) {
		  resid->mask->data.U8[oy][ox] = 1;
		}
                //resid->weight->data.F32[oy][ox] = XXX;
            }
        }
    }

    psFree (A);
    psFree (B);

    psLogMsg ("psphot.pspsf", PS_LOG_INFO, "generate residuals for %ld objects: %f sec\n", input->n, psTimerMark ("residuals"));

    psFree (xC);
    psFree (yC);
    psFree (input);

    psFree (fluxes);
    psFree (dfluxes);
    psFree (fmasks);

    psFree (fluxStats);
    psFree (fluxStatsDef);
    psFree (fluxClip);
    psFree (fluxClipDef);

    if (resid != NULL && psTraceGetLevel("psphot") > 5) {
      psphotSaveImage (NULL, resid->Ro,     "resid.ro.fits");
      psphotSaveImage (NULL, resid->Rx,     "resid.rx.fits");
      psphotSaveImage (NULL, resid->Ry,     "resid.ry.fits");
      psphotSaveImage (NULL, resid->weight, "resid.wt.fits");
      psphotSaveImage (NULL, resid->mask,   "resid.mk.fits");
    }

    psf->residuals = resid;
    return (resid != NULL) ? true : false;
}
