# 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 SetBorderT (psSparseBorder *border, pmReadout *readout, psArray *refSources, psArray *fitSources, psVector *index, bool constant_weights, int SKY_FIT_ORDER);

bool psphotEnsemblePSF (pmReadout *readout, psArray *refSources, psMetadata *recipe, pmPSF *psf, bool final) {

    bool isPSF;
    bool status;
    float x;
    float y;
    float f;
    // float r;

    psTimerStart ("psphot");

    // source analysis is done in spatial order
    refSources = psArraySort (refSources, pmSourceSortByY);

    // storage arrays for fitSources and sequence index
    psArray *fitSources = psArrayAllocEmpty (refSources->n);
    psVector *index = psVectorAllocEmpty (refSources->n, PS_TYPE_U32);

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

    // construct source table with stand-alone fitSources
    for (int i = 0; i < refSources->n; i++) {
        pmSource *refSource = refSources->data[i];

        // skip non-astronomical objects (very likely defects)
        if (refSource->type == PM_SOURCE_TYPE_DEFECT) continue;
        if (refSource->type == PM_SOURCE_TYPE_SATURATED) continue;
        if (final) {
            if (refSource->mode &  PM_SOURCE_MODE_SUBTRACTED) continue;
        } else {
            if (refSource->mode &  PM_SOURCE_MODE_BLEND) continue;
        }

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

        // save the original coords
        x = model->params->data.F32[PM_PAR_XPOS];
        y = model->params->data.F32[PM_PAR_YPOS];

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

        pmSource *fitSource = pmSourceAlloc ();

        // make temporary copies of the image pixels and mask
	// we need to have a copy which will not be modified by changes to its neighbor
        fitSource->mask   = psImageCopy (NULL, refSource->mask,   PS_TYPE_U8);
        fitSource->pixels = psImageCopy (NULL, refSource->pixels, PS_TYPE_F32);
        fitSource->weight = psImageCopy (NULL, refSource->weight, PS_TYPE_F32);

        // set model to unit peak, zero sky (we assume sky is subtracted)
	// XXX do this with the actual model or use keep the original guess?
        model->params->data.F32[PM_PAR_SKY] = 0.0;
        model->params->data.F32[PM_PAR_I0] = 1.0;

        // fill in the model pixel values
	// XXX this function should just work with the (normalized) modelFlux pixels
	// XXX we should not have to create a complete copy
	// XXX review the use of the object mask: can we use a copied mask?
        psImageInit (fitSource->pixels, 0.0);
        psImageKeepCircle (fitSource->mask, x, y, model->radiusFit, "OR", PM_MASK_MARK);
        pmModelAdd (fitSource->pixels, fitSource->mask, model, PM_MODEL_OP_FULL);

        // save source in list
	if (isPSF) {
	  fitSource->modelPSF = psMemIncrRefCounter (model);
	} else {
	  fitSource->modelEXT = psMemIncrRefCounter (model);
	}
        index->data.U32[fitSources->n] = i;
        psArrayAdd (fitSources, 100, fitSource);
        psFree (fitSource);
    }
    psLogMsg ("psphot.ensemble", PS_LOG_MINUTIA, "built fitSources: %f (%ld objects)\n", psTimerMark ("psphot"), refSources->n);

    // vectors to store stats for each object
    // psVector *weight = 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)
    // Ri is the current refSource of interest (sci pixels)
    // Fi is the current fitSource of interest (fit pixels)
    // note that the sci pixels are real image pixels while the fit pixels are unique for each
    // fitted source -- they hold the model flux

    for (int i = 0; i < fitSources->n; i++) {
        int N = index->data.U32[i];
        pmSource *Ri = refSources->data[N];
        pmSource *Fi = fitSources->data[i];

        // diagonal elements of the sparse matrix (auto-cross-product)
        f = pmSourceCrossProduct (Fi, Fi, CONSTANT_PHOTOMETRIC_WEIGHTS);
        psSparseMatrixElement (sparse, i, i, f);

        // XXX being used? weight->data.F32[i] = r;

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


        // find the image x model value
        f = pmSourceCrossProduct (Ri, Fi, CONSTANT_PHOTOMETRIC_WEIGHTS);
        psSparseVectorElement (sparse, i, f);

	// add the per-source weights (border region)
	switch (SKY_FIT_ORDER) {
	  case 1:
	    f = pmSourceWeight (Fi, 1, CONSTANT_PHOTOMETRIC_WEIGHTS);
	    psSparseBorderElementB (border, i, 1, f);
	    f = pmSourceWeight (Fi, 2, CONSTANT_PHOTOMETRIC_WEIGHTS);
	    psSparseBorderElementB (border, i, 2, f);

	  case 0:
	    f = pmSourceWeight (Fi, 0, CONSTANT_PHOTOMETRIC_WEIGHTS);
	    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 *Fj = fitSources->data[j];

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

            // got an overlap; calculate cross-product and add to output array
            f = pmSourceCrossProduct (Fi, Fj, CONSTANT_PHOTOMETRIC_WEIGHTS);
            psSparseMatrixElement (sparse, j, i, f);
        }
    }

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

    // set the sky, sky_x, sky_y components of border matrix
    SetBorderT (border, readout, refSources, fitSources, index, CONSTANT_PHOTOMETRIC_WEIGHTS, SKY_FIT_ORDER);

    psSparseConstraint constraint;
    constraint.paramMin   = 0.0;
    constraint.paramMax   = 1e8;
    constraint.paramDelta = 1e8;

    // solve for normalization terms (need include local sky?)
    psVector *norm = NULL;
    psVector *skyfit = NULL;
    if (SKY_FIT_LINEAR) {
	psSparseBorderSolve (&norm, &skyfit, constraint, border, 5);
	fprintf (stderr, "skyfit: %f\n", skyfit->data.F32[0]);
    } else {
	norm = psSparseSolve (NULL, constraint, sparse, 5);
	skyfit = NULL;
    }

    // adjust fitSources, set refSources and subtract
    for (int i = 0; i < fitSources->n; i++) {
        int N = index->data.U32[i];
        pmSource *Ri = refSources->data[N];
	pmModel *model = pmSourceGetModel (NULL, Ri);

        // assign linearly-fitted normalization
        if (isnan(norm->data.F32[i])) {
            psAbort("ensemble source is nan");
        }
        model->params->data.F32[PM_PAR_I0] = norm->data.F32[i];
	// XXX note: is the value of 'errors' modified by the sky fit?
        model->dparams->data.F32[PM_PAR_I0] = errors->data.F32[i];

        // subtract object
        pmModelSub (Ri->pixels, Ri->mask, model, PM_MODEL_OP_FULL);
        Ri->mode |= PM_SOURCE_MODE_SUBTRACTED;
        if (!final) Ri->mode |= PM_SOURCE_MODE_TEMPSUB;
	// XXX not sure about the use of TEMPSUB
    }

    // measure chisq for each source
    for (int i = 0; final && (i < fitSources->n); i++) {
        int N = index->data.U32[i];
        pmSource *Ri = refSources->data[N];
        pmModel *model = pmSourceGetModel (NULL, Ri);

        x = model->params->data.F32[PM_PAR_XPOS];
        y = model->params->data.F32[PM_PAR_YPOS];

        psImageKeepCircle (Ri->mask, x, y, model->radiusFit, "OR", PM_MASK_MARK);
        pmSourceChisq (model, Ri->pixels, Ri->mask, Ri->weight);
        psImageKeepCircle (Ri->mask, x, y, model->radiusFit, "AND", PS_NOT_U8(PM_MASK_MARK));
    }

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

    psLogMsg ("psphot.ensemble", PS_LOG_INFO, "measure ensemble of PSFs: %f sec\n", psTimerMark ("psphot"));
    return true;
}

// calculate the weight terms for the sky fit component of the matrix
static bool SetBorderT (psSparseBorder *border, pmReadout *readout, psArray *refSources, psArray *fitSources, psVector *index, bool constant_weights, int SKY_FIT_ORDER) {

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

    // turn off MARK for all object pixels
    // we must use the refSources here since their pixels point at the science image
    for (int i = 0; i < fitSources->n; i++) {
        int N = index->data.U32[i];
        pmSource *refSource = refSources->data[N];
	pmModel *model = pmSourceGetModel (NULL, refSource);
	if (model == NULL) continue;
        float x = model->params->data.F32[PM_PAR_XPOS];
        float y = model->params->data.F32[PM_PAR_YPOS];
	psImageMaskCircle (refSource->mask, x, y, model->radiusFit, "AND", PS_NOT_U8(PM_MASK_MARK));
    }	

    // accumulate the image statistics from the masked regions
    psF32 **image = readout->image->data.F32;
    psF32 **weight = readout->weight->data.F32;
    psU8 **mask = readout->mask->data.U8;

    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 = weight[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_U8(PM_MASK_MARK));

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