# include "psphotInternal.h"
# ifndef ROUND
# define ROUND(X) ((int) ((X) + 0.5*SIGN(X)))
# endif

bool psphotKronWindowSetSource(pmSource *source, psImage *kronWindow, bool useKronRadius, bool insert);
bool psphotKronWindowMag (pmSource *source, psImage *kronWindow, float radius, float minKronRadius, psImageMaskType maskVal);

bool psphotKronIterate (pmConfig *config, const pmFPAview *view, const char *filerule)
{
    bool status = true;

    // return true;

    fprintf (stdout, "\n");
    psLogMsg ("psphot", PS_LOG_INFO, "--- psphot Kron Iterate ---");

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

    int num = psphotFileruleCount(config, filerule);

    // skip the chisq image (optionally?)
    int chisqNum = psMetadataLookupS32 (&status, config->arguments, "PSPHOT.CHISQ.NUM");
    if (!status) chisqNum = -1;

    // loop over the available readouts
    for (int i = 0; i < num; i++) {
        if (i == chisqNum) continue; // skip chisq image

        // find the currently selected readout
        pmFPAfile *file = pmFPAfileSelectSingle(config->files, filerule, i); // File of interest
        psAssert (file, "missing file?");

        pmReadout *readout = pmFPAviewThisReadout(view, file->fpa);
        psAssert (readout, "missing readout?");

        pmDetections *detections = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.DETECTIONS");
        psAssert (detections, "missing detections?");

        psArray *sources = detections->newSources ? detections->newSources : detections->allSources;
        psAssert (sources, "missing sources?");

        pmPSF *psf = psMetadataLookupPtr (&status, readout->analysis, "PSPHOT.PSF");
        // psAssert (psf, "missing psf?");

        if (!psphotKronIterateReadout (config, recipe, view, readout, sources, psf)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed to measure magnitudes for %s entry %d", filerule, i);
            return false;
        }
    }
    return true;
}

int psphotKapaChannel (int channel);
bool psphotVisualShowMask (int kapaFD, psImage *inImage, const char *name, int channel);
bool psphotVisualRangeImage (int kapaFD, psImage *inImage, const char *name, int channel, float min, float max);

bool psphotKronIterateReadout(pmConfig *config, psMetadata *recipe, const pmFPAview *view, pmReadout *readout, psArray *sources, pmPSF *psf) {

    bool status = false;

    if (!sources->n) {
        psLogMsg ("psphot", PS_LOG_INFO, "no sources, skipping masked kron");
        return true;
    }

    psTimerStart ("psphot.kron");

    // determine the number of allowed threads
    int nThreads = psMetadataLookupS32(&status, config->arguments, "NTHREADS"); // Number of threads
    if (!status) {
        nThreads = 0;
    }

    float RADIUS = psMetadataLookupF32 (&status, readout->analysis, "PSF_MOMENTS_RADIUS");
    if (!status) {
        RADIUS = psMetadataLookupF32 (&status, recipe, "PSF_MOMENTS_RADIUS");
    }

    float MIN_KRON_RADIUS = psMetadataLookupF32 (&status, readout->analysis, "MOMENTS_MIN_KRON");
    if (!status) {
        MIN_KRON_RADIUS = 0.25*RADIUS;
    }

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

    // source analysis is done in S/N order (brightest first)
    sources = psArraySort (sources, pmSourceSortByFlux);
    if (!sources->n) {
        psLogMsg ("psphot", PS_LOG_INFO, "no sources, skipping blend");
        return true;
    }

    // generate the window image: multiply the flux by this to downweight neighbors
    psImage *kronWindow = psImageAlloc (readout->image->numCols, readout->image->numRows, PS_TYPE_F32);
    psImageInit (kronWindow, 1.0);

    // start with the currently known moments (Mxx, Myy, Mxy) and generate a window image
    for (int i = 0; i < sources->n; i++) {

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

	// set a window function for each source based on the moments
	// (this skips really bad sources (no peak, no moments, DEFECT)
	psphotKronWindowSetSource (source, kronWindow, false, true);
    }

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

            psArrayAdd(job->args, 1, readout);
            psArrayAdd(job->args, 1, kronWindow);
            psArrayAdd(job->args, 1, cells->data[j]); // sources
            PS_ARRAY_ADD_SCALAR(job->args, markVal,            PS_TYPE_IMAGE_MASK);
            PS_ARRAY_ADD_SCALAR(job->args, maskVal,            PS_TYPE_IMAGE_MASK);
            PS_ARRAY_ADD_SCALAR(job->args, RADIUS,             PS_TYPE_F32);
            PS_ARRAY_ADD_SCALAR(job->args, MIN_KRON_RADIUS,    PS_TYPE_F32);

// set this to 0 to run without threading
# if (1)
            if (!psThreadJobAddPending(job)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
                return false;
            }
# else
	    if (!psphotKronIterate_Threaded(job)) {
		psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		psFree(AnalysisRegion);
		return false;
	    }
	    psFree(job);
# endif
        }

        // wait for the threads to finish and manage results
        if (!psThreadPoolWait (false, true)) {
            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");
            }
            psFree(job);
        }
    }
    psFree (cellGroups);
    psFree (kronWindow);

    psLogMsg ("psphot.kron", PS_LOG_WARN, "measure masked kron magnitudes : %f sec for %ld objects\n", psTimerMark ("psphot.kron"), sources->n);
    return true;
}

bool psphotKronIterate_Threaded (psThreadJob *job) {

    pmReadout *readout              = job->args->data[0];
    psImage *kronWindow             = job->args->data[1];
    psArray *sources                = job->args->data[2];
    psImageMaskType markVal         = PS_SCALAR_VALUE(job->args->data[3],PS_TYPE_IMAGE_MASK_DATA);
    psImageMaskType maskVal         = PS_SCALAR_VALUE(job->args->data[4],PS_TYPE_IMAGE_MASK_DATA);
    float RADIUS                    = PS_SCALAR_VALUE(job->args->data[5],F32);
    float MIN_KRON_RADIUS           = PS_SCALAR_VALUE(job->args->data[6],F32);

    // XXX TEST : set iteration to 1
    for (int j = 0; j < 1; j++) {
	for (int i = 0; i < sources->n; i++) {

	    pmSource *source = sources->data[i];
	    if (!source->peak) continue; // XXX how can we have a peak-less source?

	    // allocate space for moments
	    if (!source->moments) continue;

	    // replace object in image
	    bool reSubtract = false;
	    if (source->tmpFlags & PM_SOURCE_TMPF_SUBTRACTED) {
		pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
		reSubtract = true;
	    }

	    // use S/N to control max window size
	    // float kronSN = source->moments->KronFlux / source->moments->KronFluxErr;

	    // maxWindow -> 1.5*RADIUS for kronSN = 5.0, keeping S.B. constant (kronSN ~ flux)
	    // (kronSN / maxWindow^2) = (5.0 / (1.5 RADIUS)^2)
	    // maxWindow = 1.5 * RADIUS * sqrt(kronSN / 5.0)
	    // XXX float maxWindow = (isfinite(kronSN) && (kronSN > 5.0)) ? 1.5 * RADIUS * sqrt(kronSN / 5.0) : 1.5*RADIUS;

	    // iterate to the window radius
	    // XXX float windowRadius = PS_MIN(PS_MAX(RADIUS, 4.0*source->moments->Mrf), maxWindow);

	    // XXX TEST : use a window based on the radial profile numbers: max is skyRadius, min is RADIUS
	    // if we lack the skyRadius (eg MATCHED sources), go to the default value
	    float maxWindow = isfinite(source->skyRadius) ? source->skyRadius : RADIUS;
	    float windowRadius = PS_MAX(RADIUS, maxWindow);

	    // re-allocate image, weight, mask arrays for each peak with box big enough to fit BIG_RADIUS
	    pmSourceRedefinePixels (source, readout, source->peak->x, source->peak->y, windowRadius + 2);
	    psAssert (source->pixels, "WTF?");

	    // clear the window function for this source based on the moments
	    psphotKronWindowSetSource (source, kronWindow, (j > 0), false);

	    // this function populates moments->Mrf,KronFlux,KronFluxErr
	    psphotKronWindowMag (source, kronWindow, windowRadius, MIN_KRON_RADIUS, maskVal);
	    psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal));

	    // set a window function for each source based on the moments
	    psphotKronWindowSetSource (source, kronWindow, true, true);

	    // if we subtracted it above, re-subtract the object, leave local sky
	    if (reSubtract) {
		pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
	    }
	}
    }
    return true;
}

bool psphotKronWindowMag (pmSource *source, psImage *kronWindow, float radius, float minKronRadius, psImageMaskType maskVal) {

    PS_ASSERT_PTR_NON_NULL(source, false);
    PS_ASSERT_PTR_NON_NULL(source->peak, false);
    PS_ASSERT_PTR_NON_NULL(source->pixels, false);
    PS_ASSERT_FLOAT_LARGER_THAN(radius, 0.0, false);

    psF32 R2 = PS_SQR(radius);
    float rsigma2 = 0.5 / PS_SQR(radius/2.0);

    // a note about coordinates: coordinates of objects throughout psphot refer to the primary
    // image coordinates.  the source->pixels image has an offset relative to its parent of
    // col0,row0: a pixel (x,y) in the primary image has coordinates of (x-col0, y-row0) in
    // this subimage.  we subtract off the peak coordinates, adjusted to this subimage, to have
    // minimal round-off error in the sums.  since these values are subtracted just to minimize
    // the dynamic range and are added back below, the exact value does not matter. these are
    // (int) so they can be used in the image index below.

    // Now calculate higher-order moments, using the above-calculated first moments to adjust coordinates
    // Xn  = SUM (x - xc)^n * (z - sky)

    psF32 RF = 0.0;
    psF32 RS = 0.0;

    // the peak position is less accurate but less subject to extreme deviations
    float dX = source->moments->Mx - source->peak->xf;
    float dY = source->moments->My - source->peak->yf;
    float dR = hypot(dX, dY);
    float Xo = (dR < 2.0) ? source->moments->Mx : source->peak->xf;
    float Yo = (dR < 2.0) ? source->moments->My : source->peak->yf;

    // center of mass in subimage.  Note: the calculation below uses pixel index, so we correct
    // xCM, yCM from pixel coords to pixel index here.
    psF32 xCM = Xo - 0.5 - source->pixels->col0; // coord of peak in subimage
    psF32 yCM = Yo - 0.5 - source->pixels->row0; // coord of peak in subimage

    int Xwo = source->pixels->col0;
    int Ywo = source->pixels->row0;

    psF32 **vPix = source->pixels->data.F32;
    psF32 **vWin = kronWindow->data.F32;
    psF32 **vWgt = source->variance->data.F32;
    
    psImageMaskType **vMsk = (source->maskObj == NULL) ? NULL : source->maskObj->data.PS_TYPE_IMAGE_MASK_DATA;

    for (psS32 row = 0; row < source->pixels->numRows ; row++) {

	psF32 yDiff = row - yCM;
	if (fabs(yDiff) > radius) continue;

	// coordinate of mirror pixel
	int yFlip = yCM - yDiff;
	if (yFlip < 0) continue;
	if (yFlip >= source->pixels->numRows) continue;

	for (psS32 col = 0; col < source->pixels->numCols ; col++) {
	    // check mask and value for this pixel
	    if (vMsk && (vMsk[row][col] & maskVal)) continue;
	    if (isnan(vPix[row][col])) continue;

	    psF32 xDiff = col - xCM;
	    if (fabs(xDiff) > radius) continue;

	    // coordinate of mirror pixel
	    int xFlip = xCM - xDiff;
	    if (xFlip < 0) continue;
	    if (xFlip >= source->pixels->numCols) continue;

	    // check mask and value for mirror pixel
	    if (vMsk && (vMsk[yFlip][xFlip] & maskVal)) continue;
	    if (isnan(vPix[yFlip][xFlip])) continue;

	    // radius is just a function of (xDiff, yDiff)
	    psF32 r2  = PS_SQR(xDiff) + PS_SQR(yDiff);
	    if (r2 > R2) continue;

	    // flux * window
	    float z = r2 * rsigma2;
	    assert (z >= 0.0);

	    // weight by window image and wide Gaussian
	    float weight1  = vWin[row+Ywo][col+Xwo]*exp(-z);
	    float weight2  = vWin[yFlip+Ywo][xFlip+Xwo]*exp(-z);

	    float fDiff1 = vPix[row][col]*weight1;
	    float fDiff2 = vPix[yFlip][xFlip]*weight2;

	    float pDiff = (fDiff1 > 0.0) ? sqrt(fabs(fDiff1*fDiff2)) : -sqrt(fabs(fDiff1*fDiff2));

	    // Kron Flux uses the 1st radial moment (maybe Gaussian windowed?)
	    psF32 rf = pDiff * sqrt(r2);
	    psF32 rs = 0.5 * (fDiff1 + fDiff2);

	    RF  += rf;
	    RS  += rs;
	}
    }

    // Saturate the 1st radial moment
    float Mrf = MAX(minKronRadius, RF/RS);
    if (sqrt(source->peak->detValue) < 10.0) {
	Mrf = MIN (radius, Mrf);
    }

    // Calculate the Kron magnitude (make this block optional?)
    float radKron  = 2.5*Mrf;
    float radKron2 = radKron*radKron;

    int nKronPix = 0;
    float Sum = 0.0;
    float Var = 0.0;
    float Win = 0.0;

    for (psS32 row = 0; row < source->pixels->numRows ; row++) {

	psF32 yDiff = row - yCM;
	if (fabs(yDiff) > radKron) continue;

	for (psS32 col = 0; col < source->pixels->numCols ; col++) {
	    // check mask and value for this pixel
	    if (vMsk && (vMsk[row][col] & maskVal)) continue;
	    if (isnan(vPix[row][col])) continue;

	    psF32 xDiff = col - xCM;
	    if (fabs(xDiff) > radKron) continue;

	    // radKron is just a function of (xDiff, yDiff)
	    psF32 r2  = PS_SQR(xDiff) + PS_SQR(yDiff);
	    if (r2 > radKron2) continue;

	    float weight1  = vWin[row+Ywo][col+Xwo];
	    float fDiff1 = vPix[row][col]*weight1;

	    float pDiff = fDiff1;
	    psF32 wDiff = vWgt[row][col] * weight1;

	    Sum += pDiff;
	    Var += wDiff;
	    Win += weight1;
	    nKronPix ++;
	}
    }

    source->moments->Mrf = Mrf;
    source->moments->KronFlux    = Sum;
    source->moments->KronFluxErr = sqrt(Var);

    return true;
}

bool psphotKronWindowSetSource(pmSource *source, psImage *kronWindow, bool useKronRadius, bool insert) {

    if (!source) return false;
    if (!source->peak) return false; // XXX how can we have a peak-less source?
    if (!source->moments) return false;
    if (source->type == PM_SOURCE_TYPE_DEFECT) return false;
    if (source->type == PM_SOURCE_TYPE_SATURATED) return false;
    psAssert(kronWindow, "need a window");

    // we have a source with moments Mx, My, Mxx, Mxy, Myy.  we just need to define a Gaussian that has 
    // these values (and peak of 1.0)

    int Nx = kronWindow->numCols;
    int Ny = kronWindow->numRows;

    float Xo = source->moments->Mx;
    float Yo = source->moments->My;

    float Mxx = source->moments->Mxx;
    float Mxy = source->moments->Mxy;
    float Myy = source->moments->Myy;

    float Mmajor = 0.5*(Mxx + Myy) + 0.5*sqrt(PS_SQR(Mxx - Myy) + 4.0*PS_SQR(Mxy));
    float Mminor = 0.5*(Mxx + Myy) - 0.5*sqrt(PS_SQR(Mxx - Myy) + 4.0*PS_SQR(Mxy));

    // Mxx, Mxy, Myy define the elliptical shape, but Mrf defines the width 
    float scale = PS_SQR(0.5 * source->moments->Mrf) / Mmajor;

    float Sxx = scale * Mmajor * Mminor / Myy; // sigma_x^2
    float Sxy = Mxy / (scale * Mmajor * Mminor);
    float Syy = scale * Mmajor * Mminor / Mxx; // sigma_y^2

    float Smajor = sqrt(Mmajor);

    int minX = PS_MIN(PS_MAX(Xo - 5*Smajor, 0), Nx - 1);
    int maxX = PS_MIN(PS_MAX(Xo + 5*Smajor, 0), Nx - 1);
    int minY = PS_MIN(PS_MAX(Yo - 5*Smajor, 0), Ny - 1);
    int maxY = PS_MIN(PS_MAX(Yo + 5*Smajor, 0), Ny - 1);

    float rMxx = 0.5 / Sxx;
    float rMyy = 0.5 / Syy;
    
    for (int iy = minY; iy < maxY; iy++) {
	for (int ix = minX; ix < maxX; ix++) {

	    float dX = (ix + 0.5 - Xo);
	    float dY = (iy + 0.5 - Yo);

	    float z = rMxx * PS_SQR(dX) + rMyy * PS_SQR(dY) + Sxy*dX*dY;

	    float f = insert ? 1.001 - exp(-z) : 1.0 / (1.001 - exp(-z));

	    kronWindow->data.F32[iy][ix] *= f;
	}
    }
    return true;
}
