# include "psphotInternal.h"

// non-linear model fitting for extended sources
bool psphotExtendedSourceFits (pmReadout *readout, psArray *sources, psMetadata *recipe) {

    bool status;
    int Next = 0;
    int Nconvolve = 0;
    int NconvolvePass = 0;
    int Nplain = 0;
    int NplainPass = 0;
    bool savePics = false;

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

    psMaskType markVal = psMetadataLookupU8(&status, recipe, "MARK.PSPHOT"); // Mask value for bad pixels
    assert (markVal);

    // maskVal is used to test for rejected pixels, and must include markVal
    maskVal |= markVal;

    // perform full extended source non-linear fits?
    if (!psMetadataLookupBool (&status, recipe, "EXTENDED_SOURCE_FITS")) {
	psLogMsg ("psphot", PS_LOG_INFO, "skipping extended source measurements\n");
	return true;
    }

    // select the collection of desired models
    psMetadata *models = psMetadataLookupMetadata (&status, recipe, "EXTENDED_SOURCE_MODELS");
    if (!status) {
	psWarning ("extended source model fits requested but model model is missing (EXTENDED_SOURCE_MODELS)\n");
	return true;
    }
    if (models->list->n == 0) {
	psWarning ("extended source model fits requested but no models are specified\n");
	return true;
    }

    // validate the model entries
    psMetadataIterator *iter = psMetadataIteratorAlloc (models, PS_LIST_HEAD, NULL);
    psMetadataItem *item = NULL;
    while ((item = psMetadataGetAndIncrement (iter)) != NULL) {

      if (item->type != PS_DATA_METADATA) {
	psAbort ("Invalid type for EXTENDED_SOURCE_MODEL entry %s, not a metadata folder", item->name);
	// XXX we could cull the bad entries or build a validated model folder
      }

      psMetadata *model = (psMetadata *) item->data.md;

      // check on the model type
      char *modelName = psMetadataLookupStr (&status, model, "MODEL");
      int modelType = pmModelClassGetType (modelName);
      if (modelType < 0) {
	psAbort ("Unknown model class for EXTENDED_SOURCE_MODEL entry %s: %s", item->name, modelName);
      }	
      psMetadataAddS32 (model, PS_LIST_TAIL, "MODEL_TYPE", PS_META_REPLACE, "", modelType);

      // check on the SNLIM, set a float value
      char *SNword = psMetadataLookupStr (&status, model, "SNLIM");
      if (!status) {
	psAbort("SNLIM not defined for extended source model %s\n", item->name);
      }
      float SNlim = atof (SNword);
      psMetadataAddF32 (model, PS_LIST_TAIL, "SNLIM_VALUE", PS_META_REPLACE, "", SNlim);
	
      // check on the PSF-Convolution status
      char *convolvedWord = psMetadataLookupStr (&status, model, "PSF_CONVOLVED");
      if (!status || (strcasecmp (convolvedWord, "true") && strcasecmp (convolvedWord, "false"))) {
	psAbort ("PSF_CONVOLVED entry invalid or missing for EXTENDED_SOURCE_MODEL entry %s", item->name);
      }	
      bool convolved = !strcasecmp (convolvedWord, "true");
      psMetadataAddBool (model, PS_LIST_TAIL, "PSF_CONVOLVED_VALUE", PS_META_REPLACE, "", convolved);
    }
    psFree (iter);

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

    // what fraction of the PSF is used? (radius in pixels : 2 -> 5x5 box)
    int psfSize = psMetadataLookupS32 (&status, recipe, "PCM_BOX_SIZE");
    assert (status);

    // source analysis is done in S/N order (brightest first)
    sources = psArraySort (sources, pmSourceSortBySN);

    // choose the sources of interest
    for (int i = 0; i < sources->n; i++) {

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

        // skip PSF-like and non-astronomical objects
        if (source->type == PM_SOURCE_TYPE_STAR) continue;
        if (source->type == PM_SOURCE_TYPE_DEFECT) continue;
        if (source->type == PM_SOURCE_TYPE_SATURATED) continue;

        // XXX this should use peak?
        if (source->peak->x < AnalysisRegion.x0) continue;
        if (source->peak->y < AnalysisRegion.y0) continue;
        if (source->peak->x > AnalysisRegion.x1) continue;
        if (source->peak->y > AnalysisRegion.y1) continue;

        // if model is NULL, we don't have a starting guess
	// XXX use the parameters guessed from moments
        // if (source->modelEXT == NULL) continue;

        // replace object in image
        if (source->mode & PM_SOURCE_MODE_SUBTRACTED) {
            pmSourceAdd (source, PM_MODEL_OP_FULL, maskVal);
        }
        Next ++;

	// save the modelFlux here in case we need to subtract it (for failure)
	psImage *modelFluxStart = psMemIncrRefCounter (source->modelFlux);

	if (savePics) {
	  psphotSaveImage (NULL, readout->image, "image.xp.fits");
	}

	// array to store the pointers to the model flux images while the models are being fitted
	psArray *modelFluxes = psArrayAllocEmpty (4);
    
	// allocate the array to store the model fits
	if (source->modelFits == NULL) {
	    source->modelFits = psArrayAllocEmpty (4);
	}

	// loop here over the models chosen for each source (exclude by S/N)
	psMetadataIterator *iter = psMetadataIteratorAlloc (models, PS_LIST_HEAD, NULL);
	psMetadataItem *item = NULL;
	while ((item = psMetadataGetAndIncrement (iter)) != NULL) {

	  // XXX this should have been forced above
	  assert (item->type == PS_DATA_METADATA);
	  psMetadata *model = (psMetadata *) item->data.md;

	  // check the SNlim and skip model if source is too faint
	  float SNlim = psMetadataLookupF32 (&status, model, "SNLIM_VALUE");
	  assert (status);

	  // limit selection to some SN limit
	  assert (source->peak); // how can a source not have a peak?
	  if (source->peak->SN < SNlim) continue;

	  // check on the model type
	  pmModelType modelType = psMetadataLookupS32 (&status, model, "MODEL_TYPE");
	  assert (status);

	  // check on the PSF-Convolution status
	  bool convolved = psMetadataLookupBool (&status, model, "PSF_CONVOLVED_VALUE");
	  assert (status);

	  // XXX psTraceSetLevel ("psLib.math.psMinimizeLMChi2", 6);
	  // XXX psTraceSetLevel ("psphot.psphotModelWithPSF_LMM", 6);

	  // fit the model as convolved or not
	  pmModel *modelFit = NULL;
	  if (convolved) {
	      modelFit = psphotPSFConvModel (readout, source, modelType, maskVal, markVal, psfSize);
	      if (!modelFit) {
		  psTrace ("psphot", 5, "failed to fit psf-conv model for object at %f, %f", source->moments->x, source->moments->y);
		  continue;
	      } 
	      psTrace ("psphot", 4, "fit psf-conv model for %f, %f : %s chisq = %f\n", source->moments->x, source->moments->y, pmModelClassGetName (modelFit->type), modelFit->chisq);
	      Nconvolve ++;
	      if (!(modelFit->flags & (PM_MODEL_STATUS_BADARGS | PM_MODEL_STATUS_NONCONVERGE | PM_MODEL_STATUS_OFFIMAGE))) {
		  NconvolvePass ++;
	      }
	  } else {
	      psFree (source->modelFlux);
	      source->modelFlux = NULL;
	      modelFit = psphotFitEXT (readout, source, modelType, maskVal, markVal);
	      if (!modelFit) {
		  psTrace ("psphot", 5, "failed to fit plain model for object at %f, %f", source->moments->x, source->moments->y);
		  continue;
	      } 
	      pmSourceCacheModel (source, maskVal);
	      psTrace ("psphot", 4, "fit plain model for %f, %f : %s chisq = %f\n", source->moments->x, source->moments->y, pmModelClassGetName (modelFit->type), modelFit->chisq);
	      Nplain ++;
	      if (!(modelFit->flags & (PM_MODEL_STATUS_BADARGS | PM_MODEL_STATUS_NONCONVERGE | PM_MODEL_STATUS_OFFIMAGE))) {
		  NplainPass ++;
	      }
	  }

	  // save each of the model flux images and store the best
	  psArrayAdd (modelFluxes, 4, source->modelFlux);

	  // test for fit quality / result
	  psArrayAdd (source->modelFits, 4, modelFit);

	  psFree (modelFit);
	}
	psFree (iter);

	// evaluate the relative quality of the models, choose one
	float minChisq = NAN;
	int minModel = -1;
	for (int i = 0; i < source->modelFits->n; i++) {
	    pmModel *model = source->modelFits->data[i];
	    
	    if (!(model->flags & PM_MODEL_STATUS_FITTED)) continue;
	    
	    if (model->flags & (PM_MODEL_STATUS_BADARGS)) continue;
	    if (model->flags & (PM_MODEL_STATUS_NONCONVERGE)) continue;
	    if (model->flags & (PM_MODEL_STATUS_OFFIMAGE)) continue;

	    if ((minModel < 0) || (model->chisq < minChisq)) {
		minChisq = model->chisq;
		minModel = i;
	    }
	}	    

	if (minModel == -1) {
	  // re-subtract the object, leave local sky
	  psTrace ("psphot", 5, "failed to fit extended source model to object at %f, %f", source->moments->x, source->moments->y);

	  // replace original model, subtract it
	  psFree (source->modelFlux);
	  source->modelFlux = modelFluxStart;

	  pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
	  source->mode |= PM_SOURCE_MODE_SUBTRACTED;

	  psFree (modelFluxes);

	  if (savePics) {
	    psphotSaveImage (NULL, readout->image, "image.xp.fits");
	    char key[10];
	    fprintf (stdout, "continue? ");
	    fgets (key, 8, stdin);
	    if (key[0] == 'n') {
	      savePics = false;
	    }
	  }
	  continue;
	}	

	// save the best extended model in modelEXT
	psFree (source->modelEXT);
	source->modelEXT = psMemIncrRefCounter (source->modelFits->data[minModel]);

	// save the modelFlux for the best model
	psFree (source->modelFlux);
	source->modelFlux = psMemIncrRefCounter (modelFluxes->data[minModel]);

        // subtract the best fit from the object, leave local sky
        pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);
        source->mode |= PM_SOURCE_MODE_SUBTRACTED;

	// the initial model flux is no longer needed
	psFree (modelFluxStart);
	psFree (modelFluxes);

	psTrace ("psphot", 4, "best ext model for %f, %f : %s chisq = %f\n", source->moments->x, source->moments->y, pmModelClassGetName (source->modelEXT->type), source->modelEXT->chisq);
	psTrace ("psphot", 5, "extended source model for source at %7.1f, %7.1f", source->moments->x, source->moments->y);

	if (savePics) {
	  psphotSaveImage (NULL, readout->image, "image.xm.fits");
	  char key[10];
	  fprintf (stdout, "continue? ");
	  fgets (key, 8, stdin);
	  if (key[0] == 'n') {
	    savePics = false;
	  }
	}
    }

    psLogMsg ("psphot", PS_LOG_INFO, "extended source analysis: %f sec for %d objects\n", psTimerMark ("psphot"), Next);
    psLogMsg ("psphot", PS_LOG_INFO, "  %d convolved models (%d passed)\n", Nconvolve, NconvolvePass);
    psLogMsg ("psphot", PS_LOG_INFO, "  %d plain models (%d passed)\n", Nplain, NplainPass);
    return true;
}
