# include "psphotInternal.h"
# include <gsl/gsl_sf_gamma.h>

float psphotModelContour (psImage *image, psImage *weight, psImage *mask, psMaskType maskVal, pmModel *model, float Ro);

// we need to call this function after sources have been fitted to the PSF model and
// subtracted.  To determine the CR-nature, this function examines the 9 pixels in the 3x3
// square containing the peak and compares the observed flux to the model.  To determine
// the EXT-nature, this function measures the amount of positive or negative total
// deviation from the psf model at the r = FWHM/2 position

bool psphotSourceSize (pmConfig *config, pmReadout *readout, psArray *sources, psMetadata *recipe, long first)
{

    bool status;

    psTimerStart ("psphot");

    // user-defined masks to test for good/bad pixels (build from recipe list if not yet set)
    psMaskType maskVal = psMetadataLookupU8(&status, recipe, "MASK.PSPHOT"); // Mask value for bad pixels
    assert (maskVal);

    // bit to mask the cosmic-ray pixels
    psMaskType crMask  = pmConfigMaskGet("CR", config); // Mask value for cosmic rays

    float CR_NSIGMA_LIMIT = psMetadataLookupF32 (&status, recipe, "PSPHOT.CR.NSIGMA.LIMIT");
    assert (status);

    float EXT_NSIGMA_LIMIT = psMetadataLookupF32 (&status, recipe, "PSPHOT.EXT.NSIGMA.LIMIT");
    assert (status);

    int grow = psMetadataLookupS32(&status, recipe, "PSPHOT.CR.GROW"); // Growth size for CRs
    if (!status || grow < 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "PSPHOT.CR.GROW is not positive.");
        return false;
    }

    // loop over all source
    for (int i = first; i < sources->n; i++) {
        pmSource *source = sources->data[i];

        // skip source if it was already measured
        if (isfinite(source->crNsigma)) continue;

        // source must have been subtracted
        if (!(source->mode & PM_SOURCE_MODE_SUBTRACTED)) continue;

        psF32 **resid  = source->pixels->data.F32;
        psF32 **weight = source->weight->data.F32;
        psU8 **mask    = source->maskObj->data.U8;

        // check for extendedness: measure the delta flux significance at the 1 sigma contour
        source->extNsigma = psphotModelContour (source->pixels, source->weight, source->maskObj, maskVal, source->modelPSF, 1.0);

        // XXX prevent a source from being both CR and EXT?
        if (source->extNsigma > EXT_NSIGMA_LIMIT) {
          source->mode |= PM_SOURCE_MODE_EXT_LIMIT;
        }

        int xPeak = source->peak->xf - source->pixels->col0 + 0.5;
        int yPeak = source->peak->yf - source->pixels->row0 + 0.5;

        // XXX for now, skip sources which are too close to a boundary
        // XXX raise a flag?
        if (xPeak < 1) continue;
        if (xPeak > source->pixels->numCols - 2) continue;
        if (yPeak < 1) continue;
        if (yPeak > source->pixels->numRows - 2) continue;

        // XXX for now, just skip any sources with masked pixels
        // XXX raise a flag?
        bool keep = true;
        for (int iy = -1; (iy <= +1) && keep; iy++) {
            for (int ix = -1; (ix <= +1) && keep; ix++) {
                if (mask[yPeak+iy][xPeak+ix]) { keep &= false; }
            }
        }
        if (!keep) continue;

	if (mask[yPeak][xPeak] & maskVal) continue;

        // XXX need to deal with edge peaks... and mask
	float nX = 0;
	float nY = 0;
	float nL = 0;
	float nR = 0;

	if (!(mask[yPeak+0][xPeak-1] & maskVal) && !(mask[yPeak+0][xPeak+1] & maskVal)) {
	    float cX = 2*resid[yPeak][xPeak] - resid[yPeak+0][xPeak-1] - resid[yPeak+0][xPeak+1];
	    float dcX = 4*weight[yPeak][xPeak] + weight[yPeak+0][xPeak-1] + weight[yPeak+0][xPeak+1];
	    nX = cX / sqrt(dcX);
	}

	if (!(mask[yPeak-1][xPeak+0] & maskVal) && !(mask[yPeak+1][xPeak+0] & maskVal)) {
	    float cY = 2*resid[yPeak][xPeak] - resid[yPeak-1][xPeak+0] - resid[yPeak+1][xPeak+0];
	    float dcY = 4*weight[yPeak][xPeak] + weight[yPeak-1][xPeak+0] + weight[yPeak+1][xPeak+0];
	    nY = cY / sqrt(dcY);
	}

	if (!(mask[yPeak-1][xPeak-1] & maskVal) && !(mask[yPeak+1][xPeak+1] & maskVal)) {
	    float dcL = 4*weight[yPeak][xPeak] + weight[yPeak-1][xPeak-1] + weight[yPeak+1][xPeak+1];
	    float cL = 2*resid[yPeak][xPeak] - resid[yPeak-1][xPeak-1] - resid[yPeak+1][xPeak+1];
	    nL = cL / sqrt(dcL);
	}

	if (!(mask[yPeak+1][xPeak-1] & maskVal) && !(mask[yPeak-1][xPeak+1] & maskVal)) {
	    float dcR = 4*weight[yPeak][xPeak] + weight[yPeak+1][xPeak-1] + weight[yPeak-1][xPeak+1];
	    float cR = 2*resid[yPeak][xPeak] - resid[yPeak+1][xPeak-1] - resid[yPeak-1][xPeak+1];
	    nR = cR / sqrt(dcR);
	}

        // P(chisq > chisq_obs; Ndof) = gamma_Q (Ndof/2, chisq/2)
        // Ndof = 4 ? (four measurements, no free parameters)
        // XXX this value is going to be biased low because of systematic errors.
        // we need to calibrate it somehow
        // source->psfProb = gsl_sf_gamma_inc_Q (2, 0.5*chisq);

        // not strictly accurate: overcounts the chisq contribution from the center pixel (by
        // factor of 4); also biases a bit low if any pixels are masked
	// XXX I am not sure I want to keep this value...
        source->psfChisq = PS_SQR (nX) + PS_SQR (nY) + PS_SQR (nX) + PS_SQR (nR);

        float fCR = 0.0;
        int nCR = 0;
        if (nX > 0.0) {
            fCR += nX;
            nCR ++;
        }
        if (nY > 0.0) {
            fCR += nY;
            nCR ++;
        }
        if (nL > 0.0) {
            fCR += nL;
            nCR ++;
        }
        if (nR > 0.0) {
            fCR += nR;
            nCR ++;
        }
        source->crNsigma  = (nCR > 0)  ? fCR / nCR : 0.0;
        assert (isfinite(source->crNsigma));

        // this source is thought to be a cosmic ray.  flag the detection and mask the pixels
        if (source->crNsigma > CR_NSIGMA_LIMIT) {
            source->mode |= PM_SOURCE_MODE_CR_LIMIT;
            pmPeak *peak = source->peak;
            psAssert (peak, "NULL peak");

            // replace the source flux
            pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
            source->mode &= ~PM_SOURCE_MODE_SUBTRACTED;

            psImage *mask   = source->maskView;
            psImage *pixels = source->pixels;
            psImage *weight = source->weight;

            # define SN_LIMIT 5.0

            int xo = peak->x - pixels->col0;
            int yo = peak->y - pixels->row0;

            // mark the pixels in this row to the left, then the right
            for (int ix = xo; ix >= 0; ix--) {
                float SN = pixels->data.F32[yo][ix] / sqrt(weight->data.F32[yo][ix]);
                if (SN > SN_LIMIT) {
                    mask->data.U8[yo][ix] |= crMask;
                }
            }
            for (int ix = xo + 1; ix < pixels->numCols; ix++) {
                float SN = pixels->data.F32[yo][ix] / sqrt(weight->data.F32[yo][ix]);
                if (SN > SN_LIMIT) {
                    mask->data.U8[yo][ix] |= crMask;
                }
            }

            // for each of the neighboring rows, mark the high pixels if they have a marked neighbor
            // first go up:
            for (int iy = PS_MIN(yo, mask->numRows-2); iy >= 0; iy--) {
                // mark the pixels in this row to the left, then the right
                for (int ix = 0; ix < pixels->numCols; ix++) {
                    float SN = pixels->data.F32[iy][ix] / sqrt(weight->data.F32[iy][ix]);
                    if (SN < SN_LIMIT) continue;

                    bool valid = false;
                    valid |= (mask->data.U8[iy+1][ix] & crMask);
                    valid |= (ix > 0) ? (mask->data.U8[iy+1][ix-1] & crMask) : 0;
                    valid |= (ix <= mask->numCols) ? (mask->data.U8[iy+1][ix+1] & crMask) : 0;

                    if (!valid) continue;
                    mask->data.U8[iy][ix] |= crMask;
                }
            }
            // next go down:
            for (int iy = PS_MIN(yo+1, mask->numRows-1); iy < pixels->numRows; iy++) {
                // mark the pixels in this row to the left, then the right
                for (int ix = 0; ix < pixels->numCols; ix++) {
                    float SN = pixels->data.F32[iy][ix] / sqrt(weight->data.F32[iy][ix]);
                    if (SN < SN_LIMIT) continue;

                    bool valid = false;
                    valid |= (mask->data.U8[iy-1][ix] & crMask);
                    valid |= (ix > 0) ? (mask->data.U8[iy-1][ix-1] & crMask) : 0;
                    valid |= (ix <= mask->numCols) ? (mask->data.U8[iy-1][ix+1] & crMask) : 0;

                    if (!valid) continue;
                    mask->data.U8[iy][ix] |= crMask;
                }
            }
        }
    }

    // now that we have masked pixels associated with CRs, we can grow the mask
    if (grow > 0) {
        psImage *newMask = psImageConvolveMask(NULL, readout->mask, crMask, crMask, -grow, grow, -grow, grow);
        if (!newMask) {
            psError(PS_ERR_UNKNOWN, false, "Unable to grow CR mask");
            return false;
        }
        psFree(readout->mask);
        readout->mask = newMask;
    }

    psLogMsg ("psphot.size", PS_LOG_INFO, "measure source sizes for %ld sources: %f sec\n",
              sources->n - first, psTimerMark ("psphot"));

    psphotVisualPlotSourceSize (sources);
    psphotVisualShowSourceSize (sources);

    return true;
}

// given the PSF ellipse parameters, navigate around the 1sigma contour, return the total
// deviation in sigmas.  This is measure on the residual image - should we ignore negative
// deviations?
    float psphotModelContour (psImage *image, psImage *weight, psImage *mask, psMaskType maskVal, pmModel *model, float Ro) {

    psF32 *PAR = model->params->data.F32;

    // Ro = (x / SXX)^2 + (y / SYY)^2 + x y SXY;
    // y^2 (1/SYY^2) + y (x SXY) + (x / SXX)^2 - Ro = 0;
    // y = [-(x SXY) +/- sqrt ((x SXY)^2 - 4 (1/SYY^2) ((x/SXX)^2 - Ro))] * [SYY^2 / 2];
    // y = [-B +/- sqrt (B^2 - 4 A C)] / [2 A];

    // min/max value of x is where T -> 0
    // solve this for x2:
    float Q = Ro * PS_SQR(PAR[PM_PAR_SXX]) / (1.0 - PS_SQR(PAR[PM_PAR_SXX]*PAR[PM_PAR_SYY]*PAR[PM_PAR_SXY]) / 4.0);
    if (Q < 0.0) return NAN; // ellipse is imaginary

    int xMax = sqrt(Q);
    int xMin = -1.0*xMax;

    int nPts = 0;
    float nSigma = 0.0;

    for (int x = xMin; x <= xMax; x++) {
        float A = PS_SQR (1.0 / PAR[PM_PAR_SYY]);
        float B = x * PAR[PM_PAR_SXY];
        float C = PS_SQR (x / PAR[PM_PAR_SXX]) - Ro;

        float T = PS_SQR(B) - 4*A*C;
        if (T < 0.0) continue;

        float yP = (-B + sqrt (T)) / (2.0 * A);
        float yM = (-B - sqrt (T)) / (2.0 * A);

        int xPix  = x  + PAR[PM_PAR_XPOS] - image->col0 + 0.5;
        int yPixM = yM + PAR[PM_PAR_YPOS] - image->row0 + 0.5;
        int yPixP = yP + PAR[PM_PAR_YPOS] - image->row0 + 0.5;

        if (xPix < 0) continue;
        if (xPix >= image->numCols) continue;

        if ((yPixM >= 0) && (yPixM < image->numRows)) {
            if (!mask || !(mask->data.U8[yPixM][xPix] & maskVal)) {
                float dSigma = image->data.F32[yPixM][xPix] / sqrt (weight->data.F32[yPixM][xPix]);
                nSigma += dSigma;
                nPts ++;
            }
        }

        if (yPixM == yPixP) continue;

        if ((yPixP >= 0) && (yPixP < image->numRows)) {
            if (!mask || !(mask->data.U8[yPixP][xPix] & maskVal)) {
                float dSigma = image->data.F32[yPixP][xPix] / sqrt (weight->data.F32[yPixP][xPix]);
                nSigma += dSigma;
                nPts ++;
            }
        }
    }
    nSigma /= nPts;
    return nSigma;
}
