# include "psphotInternal.h"

// fit flux (and optionally sky model) to all reasonable sources
// with the linear fitting process.  sources must have an associated
// model with selected pixels, and the fit radius must be defined

// given the set of sources, each of which points to the pixels in the
// science image, we construct a set of simulated sources with their own pixels.
// these are used to determine the simultaneous linear fit of fluxes.
// the analysis is performed wrt the simulated pixel values

static bool SetBorderMatrixElements (psSparseBorder *border, pmReadout *readout, psArray *sources, bool constant_weights, int SKY_FIT_ORDER, psImageMaskType markVal);

// for now, let's store the detections on the readout->analysis for each readout
bool psphotFitSourcesLinear (pmConfig *config, const pmFPAview *view, const char *filerule, bool final)
{
    bool status = true;

    fprintf (stdout, "\n");
    psLogMsg ("psphot", PS_LOG_INFO, "--- psphot Fit Source (Linear) ---");

    // select the appropriate recipe information
    psMetadata *recipe  = psMetadataLookupPtr (&status, config->recipes, PSPHOT_RECIPE);
    assert (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->allSources;
        psAssert (sources, "missing sources?");

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

        if (!psphotFitSourcesLinearReadout (recipe, readout, sources, psf, final)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed to fit sources (linear) for %s entry %d", filerule, i);
            return false;
        }

	psphotVisualShowResidualImage (readout, (num > 0)); 
	psphotVisualShowPeaks (detections);
	psphotVisualShowObjectRegions (readout, recipe, sources);
    }
    return true;
}

bool psphotFitSourcesLinearReadout (psMetadata *recipe, pmReadout *readout, psArray *sources, pmPSF *psf, bool final) {

    bool status;
    float x;
    float y;
    float f;

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

    psTimerStart ("psphot.linear");

    // 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 spatial order
    sources = psArraySort (sources, pmSourceSortByY);

    // storage array for fitSources
    psArray *fitSources = psArrayAllocEmpty (sources->n);

    // option to limit analysis to a specific region
    char *region = psMetadataLookupStr (&status, recipe, "ANALYSIS_REGION");
    psRegion AnalysisRegion = psRegionFromString (region);
    AnalysisRegion = psRegionForImage (readout->image, AnalysisRegion);
    if (psRegionIsNaN (AnalysisRegion)) psAbort("analysis region mis-defined");

    bool CONSTANT_PHOTOMETRIC_WEIGHTS =
        psMetadataLookupBool(&status, recipe, "CONSTANT_PHOTOMETRIC_WEIGHTS");
    if (!status) {
        psAbort("You must provide a value for the BOOL recipe CONSTANT_PHOTOMETRIC_WEIGHTS");
    }
    int SKY_FIT_ORDER = psMetadataLookupS32(&status, recipe, "SKY_FIT_ORDER");
    if (!status) {
        SKY_FIT_ORDER = 0;
    }
    bool SKY_FIT_LINEAR = psMetadataLookupBool(&status, recipe, "SKY_FIT_LINEAR");
    if (!status) {
        SKY_FIT_LINEAR = false;
    }
    
    float MIN_VALID_FLUX = psMetadataLookupF32(&status, recipe, "PSF_FIT_MIN_VALID_FLUX");
    if (!status) {
        MIN_VALID_FLUX = 0.0;
    }
    float MAX_VALID_FLUX = psMetadataLookupF32(&status, recipe, "PSF_FIT_MAX_VALID_FLUX");
    if (!status) {
        MAX_VALID_FLUX = 1e+8;
    }

    // XXX test: choose a larger-than expected radius:
    float covarFactor = psImageCovarianceFactorForAperture(readout->covariance, 10.0); // Covariance matrix
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "covariance factor: %f\n", covarFactor);

    // XXX do not apply covarFactor for the moment...
    // covarFactor = 1.0;

    int Nsat = 0;

    // select the sources which will be used for the fitting analysis
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];

	// psAssert (source->peak, "source without peak??");
	// psAssert (source->peak->footprint, "peak without footprint??");

        // turn this bit off and turn it on again if we pass this test
        source->mode &= ~PM_SOURCE_MODE_LINEAR_FIT;

        // skip non-astronomical objects (very likely defects)
        if (source->type == PM_SOURCE_TYPE_DEFECT) continue;
        if (source->type == PM_SOURCE_TYPE_SATURATED) continue;

        // do not include CRs in the full ensemble fit
        if (source->mode & PM_SOURCE_MODE_CR_LIMIT) continue;

	// do not include MOMENTS_FAILURES in the fit
        if (source->mode & PM_SOURCE_MODE_MOMENTS_FAILURE) continue;

	// XXX count saturated stars
        if (source->mode & PM_SOURCE_MODE_SATSTAR) {
	    Nsat ++;
	}

        if (final) {
            if (source->tmpFlags & PM_SOURCE_TMPF_SUBTRACTED) continue;
        } else {
            // if (source->mode & PM_SOURCE_MODE_BLEND) continue;
        }

        // generate model for sources without, or skip if we can't
        if (!source->modelFlux) {
            if (!pmSourceCacheModel (source, maskVal)) continue;
        }

        // save the original coords
        x = source->peak->xf;
        y = source->peak->yf;

        // is the source in the region of interest?
        if (x < AnalysisRegion.x0) continue;
        if (y < AnalysisRegion.y0) continue;
        if (x > AnalysisRegion.x1) continue;
        if (y > AnalysisRegion.y1) continue;

	// check the integral of the model : is it large enough?
	// apply mask?
	float modelSum = 0.0;
	float maskedSum = 0.0;
	for (int iy = 0; iy < source->modelFlux->numRows; iy++) {
	    for (int ix = 0; ix < source->modelFlux->numCols; ix++) {
		modelSum += source->modelFlux->data.F32[iy][ix];
		if (source->maskObj->data.PS_TYPE_IMAGE_MASK_DATA[iy][ix] & maskVal) continue;
		maskedSum += source->modelFlux->data.F32[iy][ix];
	    }
	}
	if (modelSum < 0.5) continue; // skip sources with no model constraint (somewhat arbitrary limit)
	if (modelSum < 0.8) {
	    fprintf (stderr, "low-sig model @ %f, %f (%f masked sum, %f sum, %f peak)\n",
		     source->peak->xf, source->peak->yf, maskedSum, modelSum, source->peak->rawFlux);
	}
	if (maskedSum < 0.5) {
	    fprintf (stderr, "worrying model @ %f, %f (%f masked sum, %f sum, %f peak)\n",
		     source->peak->xf, source->peak->yf, maskedSum, modelSum, source->peak->rawFlux);
	}

	bool isPSF = false;
        pmModel *model = pmSourceGetModel (&isPSF, source);

	// clear the 'mark' pixels and remask on the fit aperture 
        psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal));
        psImageKeepCircle (source->maskObj, source->peak->x, source->peak->y, model->fitRadius, "OR", markVal);

	// we call this function multiple times. for the first time, we have only PSF models for all objects
	// the second time has extended sources.  If we ever fit the PSF model, we should raise this bit
	source->mode |= PM_SOURCE_MODE_LINEAR_FIT;
	if (isPSF) {
	    source->mode |= PM_SOURCE_MODE_PSFMODEL;
	}	    
	
        psArrayAdd (fitSources, 100, source);
    }
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "built fitSources: %f sec (%ld objects)\n", psTimerMark ("psphot.linear"), sources->n);

    if (fitSources->n == 0) {
        psFree(fitSources);
        return true;
    }

    // vectors to store stats for each object
    // psVector *variance = psVectorAlloc (fitSources->n, PS_TYPE_F32);
    psVector *errors = psVectorAlloc (fitSources->n, PS_TYPE_F32);

    // create the border matrix (includes the sparse matrix)
    // for just sky: 1 row; for x,y terms: 3 rows
    psSparse *sparse = psSparseAlloc (fitSources->n, 100);
    int nBorder = (SKY_FIT_ORDER == 0) ? 1 : 3;
    psSparseBorder *border = psSparseBorderAlloc (sparse, nBorder);

    // fill out the sparse matrix elements and border elements (B)
    // SRCi is the current source of interest
    // SRCj is a possibly overlapping source
    for (int i = 0; i < fitSources->n; i++) {
        pmSource *SRCi = fitSources->data[i];

        // diagonal elements of the sparse matrix (auto-cross-product)
        f = pmSourceModelDotModel (SRCi, SRCi, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
        psSparseMatrixElement (sparse, i, i, f);

        // the formal error depends on the weighting scheme
        if (CONSTANT_PHOTOMETRIC_WEIGHTS) {
            float var = pmSourceModelDotModel (SRCi, SRCi, false, covarFactor, maskVal);
            errors->data.F32[i] = 1.0 / sqrt(var);
        } else {
            errors->data.F32[i] = 1.0 / sqrt(f);
        }


        // find the image x model value
        f = pmSourceDataDotModel (SRCi, SRCi, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
        psSparseVectorElement (sparse, i, f);

        // add the per-source variances (border region)
        switch (SKY_FIT_ORDER) {
          case 1:
            f = pmSourceModelWeight (SRCi, 1, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
            psSparseBorderElementB (border, i, 1, f);
            f = pmSourceModelWeight (SRCi, 2, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
            psSparseBorderElementB (border, i, 2, f);

          case 0:
            f = pmSourceModelWeight (SRCi, 0, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
            psSparseBorderElementB (border, i, 0, f);
            break;

          default:
            psAbort("Invalid SKY_FIT_ORDER %d\n", SKY_FIT_ORDER);
            break;
        }

        // loop over all other stars following this one
        for (int j = i + 1; j < fitSources->n; j++) {
            pmSource *SRCj = fitSources->data[j];

            // skip over disjoint source images, break after last possible overlap
            if (SRCi->pixels->row0 + SRCi->pixels->numRows < SRCj->pixels->row0) break;
            if (SRCj->pixels->row0 + SRCj->pixels->numRows < SRCi->pixels->row0) continue;
            if (SRCi->pixels->col0 + SRCi->pixels->numCols < SRCj->pixels->col0) continue;
            if (SRCj->pixels->col0 + SRCj->pixels->numCols < SRCi->pixels->col0) continue;

            // got an overlap; calculate cross-product and add to output array
            f = pmSourceModelDotModel (SRCi, SRCj, CONSTANT_PHOTOMETRIC_WEIGHTS, covarFactor, maskVal);
            psSparseMatrixElement (sparse, j, i, f);
        }
    }

# if (0)
    static int npass = 0;
    char name[128];
    FILE *f1 = NULL;
    int fd = -1;

    snprintf (name, 128, "sparse.Aij.%02d.dat", npass);
    f1 = fopen (name, "w");
    psAssert (f1, "failed to open file\n");
    fd = fileno (f1);
    p_psVectorPrint (fd, sparse->Aij, "Aij");
    fclose (f1);

    snprintf (name, 128, "sparse.Bfj.%02d.dat", npass);
    f1 = fopen (name, "w");
    psAssert (f1, "failed to open file\n");
    fd = fileno (f1);
    p_psVectorPrint (fd, sparse->Bfj, "Bfj");
    fclose (f1);

    snprintf (name, 128, "sparse.Qii.%02d.dat", npass);
    f1 = fopen (name, "w");
    psAssert (f1, "failed to open file\n");
    fd = fileno (f1);
    p_psVectorPrint (fd, sparse->Qii, "Qii");
    fclose (f1);
    npass ++;
# endif

    psSparseResort (sparse);
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "built matrix: %f sec (%d elements)\n", psTimerMark ("psphot.linear"), sparse->Nelem);

    // set the sky, sky_x, sky_y components of border matrix
    SetBorderMatrixElements (border, readout, fitSources, CONSTANT_PHOTOMETRIC_WEIGHTS, SKY_FIT_ORDER, markVal);

    psSparseConstraint constraint;
    constraint.paramMin   = MIN_VALID_FLUX;
    constraint.paramMax   = MAX_VALID_FLUX;
    constraint.paramDelta = 1e7;

    // solve for normalization terms (need include local sky?)
    psVector *norm = NULL;
    psVector *skyfit = NULL;
    if (SKY_FIT_LINEAR) {
        psSparseBorderSolve (&norm, &skyfit, constraint, border, 5);
        psLogMsg ("psphot", PS_LOG_MINUTIA, "skyfit: %f\n", skyfit->data.F32[0]);
    } else {
        norm = psSparseSolve (NULL, constraint, sparse, 5);
        skyfit = NULL;
    }
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "solve matrix: %f sec (%d elements)\n", psTimerMark ("psphot.linear"), sparse->Nelem);

    // Philosophical question: we measure bright objects in three passes: 1) linear fit; 2)
    // non-linear fit; 3) linear fit: should we retain the chisq and errors from the
    // intermediate non-linear fit?  the non-linear fit provides better values for the position
    // errors, and for extended sources, the shape errors

    // adjust I0 for fitSources and subtract
    for (int i = 0; i < fitSources->n; i++) {
        pmSource *source = fitSources->data[i];
        pmModel *model = pmSourceGetModel (NULL, source);

        // assign linearly-fitted normalization
        if (isnan(norm->data.F32[i])) {
            psAbort("linear fitted source is nan");
        }

        model->params->data.F32[PM_PAR_I0] = norm->data.F32[i];
        model->dparams->data.F32[PM_PAR_I0] = errors->data.F32[i];

	// clear the 'mark' pixels so the subtraction covers the full window
        psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal));

        // subtract object & add noise
        pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
    }
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "sub models: %f sec (%d elements)\n", psTimerMark ("psphot.linear"), sparse->Nelem);

    // measure chisq for each source
    // for (int i = 0; final && (i < fitSources->n); i++) {
    for (int i = 0; i < fitSources->n; i++) {
        pmSource *source = fitSources->data[i];
        pmModel *model = pmSourceGetModel (NULL, source);
        if (!(source->mode & PM_SOURCE_MODE_NONLINEAR_FIT)) {
	    model->nPar = 1; // LINEAR-only sources have 1 parameter; NONLINEAR sources have their original value
	}
        pmSourceChisq (model, source->pixels, source->maskObj, source->variance, maskVal);
    }
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "get chisqs: %f sec (%d elements)\n", psTimerMark ("psphot.linear"), sparse->Nelem);

    // psFree (index);
    psFree (sparse);
    psFree (fitSources);
    psFree (norm);
    psFree (skyfit);
    psFree (errors);
    psFree (border);

    psLogMsg ("psphot.ensemble", PS_LOG_WARN, "measure ensemble of PSFs: %f sec\n", psTimerMark ("psphot.linear"));

    psphotVisualPlotChisq (sources);
    // psphotVisualShowFlags (sources);

    // We have to place this visualization here because the models are not realized until
    // psphotGuessModels or fitted until psphotFitSourcesLinear.
    // psphotVisualShowPSFStars (recipe, psf, sources);

    return true;
}

// XXX do we need this?
// XXX disallow the simultaneous sky fit and remove this code...

// Calculate the weight terms for the sky fit component of the matrix.  This function operates
// on the pixels which correspond to all of the sources of interest.  These elements fill in
// the border matrix components in the sparse matrix equation.
static bool SetBorderMatrixElements (psSparseBorder *border, pmReadout *readout, psArray *sources, bool constant_weights, int SKY_FIT_ORDER, psImageMaskType markVal) {

    // generate the image-wide weight terms
    // turn on MARK for all image pixels
    psRegion fullArray = psRegionSet (0, 0, 0, 0);
    fullArray = psRegionForImage (readout->mask, fullArray);
    psImageMaskRegion (readout->mask, fullArray, "OR", markVal);

    // turn off MARK for all object pixels
    for (int i = 0; i < sources->n; i++) {
        pmSource *source = sources->data[i];
        pmModel *model = pmSourceGetModel (NULL, source);
        if (model == NULL) continue;
        float x = model->params->data.F32[PM_PAR_XPOS];
        float y = model->params->data.F32[PM_PAR_YPOS];
        psImageMaskCircle (source->maskView, x, y, model->fitRadius, "AND", PS_NOT_IMAGE_MASK(markVal));
    }

    // accumulate the image statistics from the masked regions
    psF32 **image  = readout->image->data.F32;
    psF32 **variance = readout->variance->data.F32;
    psImageMaskType  **mask   = readout->mask->data.PS_TYPE_IMAGE_MASK_DATA;

    double w, x, y, x2, xy, y2, xc, yc, wt, f, fo, fx, fy;
    w = x = y = x2 = xy = y2 = fo = fx = fy = 0;

    int col0 = readout->image->col0;
    int row0 = readout->image->row0;

    for (int j = 0; j < readout->image->numRows; j++) {
        for (int i = 0; i < readout->image->numCols; i++) {
            if (mask[j][i]) continue;
            if (constant_weights) {
                wt = 1.0;
            } else {
                wt = variance[j][i];
            }
            f = image[j][i];
            w   += 1/wt;
            fo  += f/wt;
            if (SKY_FIT_ORDER == 0) continue;

            xc  = i + col0;
            yc  = j + row0;
            x  +=    xc/wt;
            y  +=    yc/wt;
            x2 += xc*xc/wt;
            xy += xc*yc/wt;
            y2 += yc*yc/wt;
            fx +=  f*xc/wt;
            fy +=  f*yc/wt;
        }
    }

    // turn off MARK for all image pixels
    psImageMaskRegion (readout->mask, fullArray, "AND", PS_NOT_IMAGE_MASK(markVal));

    // set the Border T elements
    psSparseBorderElementG (border, 0, fo);
    psSparseBorderElementT (border, 0, 0, w);
    if (SKY_FIT_ORDER > 0) {
        psSparseBorderElementG (border, 0, fx);
        psSparseBorderElementG (border, 0, fy);
        psSparseBorderElementT (border, 1, 0, x);
        psSparseBorderElementT (border, 2, 0, y);
        psSparseBorderElementT (border, 0, 1, x);
        psSparseBorderElementT (border, 1, 1, x2);
        psSparseBorderElementT (border, 2, 1, xy);
        psSparseBorderElementT (border, 0, 2, y);
        psSparseBorderElementT (border, 1, 2, xy);
        psSparseBorderElementT (border, 2, 2, y2);
    }

    return true;
}
