# include "psphotInternal.h"

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

    fprintf (stdout, "\n");
    psLogMsg ("psphot", PS_LOG_INFO, "--- psphot Extended Source Fits ---");

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

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

    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
	if (!psphotExtendedSourceFitsReadout (config, view, filerule, i, recipe)) {
            psError (PSPHOT_ERR_CONFIG, false, "failed on to fit extended sources for %s entry %d", filerule, i);
	    return false;
	}
    }
    return true;
}

// non-linear model fitting for extended sources
bool psphotExtendedSourceFitsReadout (pmConfig *config, const pmFPAview *view, const char *filerule, int index, psMetadata *recipe) {

    bool status;
    int Next = 0;
    int Nconvolve = 0;
    int NconvolvePass = 0;
    int Nplain = 0;
    int NplainPass = 0;
    int Nfaint = 0;
    int Nfail = 0;

    psphotSersicModelClassInit();

    psTimerStart ("psphot.extended");

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

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

    psLogMsg("psphot", PS_LOG_INFO, "extended source fits for image %d", index);
    psphotVisualShowImage(readout);

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

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

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

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

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

    psImageMaskType markVal = psMetadataLookupImageMask(&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;

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

    psphotInitRadiusEXT (recipe, readout);

    // 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) {
        // XXX we could cull the bad entries or build a validated model folder
        psAbort ("Invalid type for EXTENDED_SOURCE_MODEL entry %s, not a metadata folder", item->name);
      }

      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 = psRegionAlloc(0,0,0,0);
    *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, pmSourceSortByFlux);

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

            psArrayAdd(job->args, 1, readout);
            psArrayAdd(job->args, 1, cells->data[j]); // sources
            psArrayAdd(job->args, 1, models);
            psArrayAdd(job->args, 1, AnalysisRegion); // XXX make a pointer

            PS_ARRAY_ADD_SCALAR(job->args, psfSize, PS_TYPE_S32);
            PS_ARRAY_ADD_SCALAR(job->args, maskVal, PS_TYPE_IMAGE_MASK);
            PS_ARRAY_ADD_SCALAR(job->args, markVal, PS_TYPE_IMAGE_MASK);

            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Next
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nconvolve
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for NconvolvePass
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nplain
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for NplainPass
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfaint
            PS_ARRAY_ADD_SCALAR(job->args, 0, PS_TYPE_S32); // this is used as a return value for Nfail

// set this to 0 to run without threading
# if (1)	    
            if (!psThreadJobAddPending(job)) {
                psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		psFree(AnalysisRegion);
                return false;
            } 
# else
	    if (!psphotExtendedSourceFits_Threaded(job)) {
		psError(PS_ERR_UNKNOWN, false, "Unable to guess model.");
		psFree(AnalysisRegion);
		return false;
	    }
	    psScalar *scalar = NULL;
	    scalar = job->args->data[7];
	    Next += scalar->data.S32;
	    scalar = job->args->data[8];
	    Nconvolve += scalar->data.S32;
	    scalar = job->args->data[9];
	    NconvolvePass += scalar->data.S32;
	    scalar = job->args->data[10];
	    Nplain += scalar->data.S32;
	    scalar = job->args->data[11];
	    NplainPass += scalar->data.S32;
	    scalar = job->args->data[12];
	    Nfaint += scalar->data.S32;
	    scalar = job->args->data[13];
	    Nfail += scalar->data.S32;
	    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.");
	    psFree(AnalysisRegion);
            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");
            } else {
                psScalar *scalar = NULL;
                scalar = job->args->data[7];
                Next += scalar->data.S32;
                scalar = job->args->data[8];
                Nconvolve += scalar->data.S32;
                scalar = job->args->data[9];
                NconvolvePass += scalar->data.S32;
                scalar = job->args->data[10];
                Nplain += scalar->data.S32;
                scalar = job->args->data[11];
                NplainPass += scalar->data.S32;
                scalar = job->args->data[12];
                Nfaint += scalar->data.S32;
                scalar = job->args->data[13];
                Nfail += scalar->data.S32;
            }
            psFree(job);
	}
    }
    psFree (cellGroups);
    psFree(AnalysisRegion);

    psphotSersicModelClassCleanup();

    psphotVisualShowResidualImage (readout, false);

    psLogMsg ("psphot", PS_LOG_WARN, "extended source model fits: %f sec for %d objects\n", psTimerMark ("psphot.extended"), 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);
    psLogMsg ("psphot", PS_LOG_INFO, "  %d too faint to fit, %d failed\n", Nfaint, Nfail);
    return true;
}

// non-linear model fitting for extended sources
bool psphotExtendedSourceFits_Threaded (psThreadJob *job) {

    bool status;
    int Next = 0;
    int Nfaint = 0;
    int Nfail = 0;
    int Nconvolve = 0;
    int NconvolvePass = 0;
    int Nplain = 0;
    int NplainPass = 0;
    float fitRadius, windowRadius;
    psScalar *scalar = NULL;

    // arguments: readout, sources, models, region, psfSize, maskVal, markVal
    pmReadout *readout      = job->args->data[0];
    psArray *sources        = job->args->data[1];
    psMetadata *models      = job->args->data[2];
    psRegion *region        = job->args->data[3];
    int psfSize             = PS_SCALAR_VALUE(job->args->data[4],S32);
    psImageMaskType maskVal = PS_SCALAR_VALUE(job->args->data[5],PS_TYPE_IMAGE_MASK_DATA);
    psImageMaskType markVal = PS_SCALAR_VALUE(job->args->data[6],PS_TYPE_IMAGE_MASK_DATA);

    // Define source fitting parameters for extended source fits
    pmSourceFitOptions *fitOptions = pmSourceFitOptionsAlloc();
    fitOptions->mode           = PM_SOURCE_FIT_EXT;
    fitOptions->saveCovariance = true;  // XXX make this a user option?
    fitOptions->covarFactor    = psImageCovarianceFactorForAperture(readout->covariance, 10.0); // Covariance matrix

    // XXX for now, use the defaults for the rest:
    // fitOptions->nIter         = fitIter;
    // fitOptions->tol           = fitTol;
    // fitOptions->poissonErrors = poisson;
    // fitOptions->weight        = PS_SQR(skySig);

    // 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->mode & PM_SOURCE_MODE_EXT_LIMIT)) 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 < region->x0) continue;
        if (source->peak->y < region->y0) continue;
        if (source->peak->x > region->x1) continue;
        if (source->peak->y > region->y1) continue;

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

	// set the fit radius based on the first radial moment (also sets the mask pixels)
	psphotSetRadiusMomentsExact(&fitRadius, &windowRadius, readout, source, markVal); // NOTE : 6 allocs

	// UPDATE : we have changed the moments calculation.  There is now an iteration within 
	// psphotKronMasked to determine moments appropriate for a larger object.  The values
	// Mrf, KronFlux, and KronFluxErr are calculated for the iterated radius.  The other
	// values are left at the psf-based values.

        // allocate the array to store the model fits
        if (source->modelFits == NULL) {
            source->modelFits = psArrayAllocEmpty (models->list->n);
        }

        // 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 FIT_SN_LIM = psMetadataLookupF32 (&status, model, "SNLIM_VALUE");
          assert (status);

          // limit selection to some SN limit
          // assert (source->peak); // how can a source not have a peak?
	  // limit selection to some SN limit
	  bool skipSource = false;
	  if (source->mode & PM_SOURCE_MODE_EXT_LIMIT) {
	      skipSource = (source->moments->KronFlux < FIT_SN_LIM * source->moments->KronFluxErr);
	  } else {
	      skipSource = (sqrt(source->peak->detValue) < FIT_SN_LIM);
	  }
          if (skipSource) {
	      Nfaint ++;
	      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);

          // fit the model as convolved or not
          pmModel *modelFit = NULL;
          if (convolved) {
	      // NOTE : 4 more allocs to here
              modelFit = psphotFitPCM (readout, source, fitOptions, modelType, maskVal, markVal, psfSize); // NOTE : 2313 allocs in here
              if (!modelFit) {
                  psTrace ("psphot", 5, "failed to fit psf-conv model for object at %f, %f", source->moments->Mx, source->moments->My);
		  Nfail ++;
                  continue;
              }
              psTrace ("psphot", 4, "fit psf-conv model for %f, %f : %s chisq = %f (npix: %d, niter: %d)\n", 
		       source->moments->Mx, source->moments->My, pmModelClassGetName (modelFit->type), modelFit->chisq, modelFit->nPix, modelFit->nIter);
              Nconvolve ++;
              if (!(modelFit->flags & (PM_MODEL_STATUS_BADARGS | PM_MODEL_STATUS_NONCONVERGE | PM_MODEL_STATUS_OFFIMAGE))) {
                  NconvolvePass ++;
		  source->mode |= PM_SOURCE_MODE_EXTENDED_FIT;
              }
          } else {
              psFree (source->modelFlux);
              source->modelFlux = NULL;
              modelFit = psphotFitEXT (readout, source, fitOptions, modelType, maskVal, markVal);
              if (!modelFit) {
                  psTrace ("psphot", 5, "failed to fit plain model for object at %f, %f", source->moments->Mx, source->moments->My);
		  Nfail ++;
                  continue;
              }
              psTrace ("psphot", 4, "fit plain model for %f, %f : %s chisq = %f (npix: %d, niter: %d)\n", 
		       source->moments->Mx, source->moments->My, pmModelClassGetName (modelFit->type), modelFit->chisq, modelFit->nPix, modelFit->nIter);
              Nplain ++;
              if (!(modelFit->flags & (PM_MODEL_STATUS_BADARGS | PM_MODEL_STATUS_NONCONVERGE | PM_MODEL_STATUS_OFFIMAGE))) {
                  NplainPass ++;
		  source->mode |= PM_SOURCE_MODE_EXTENDED_FIT;
              }
          }

	  // XXX deprecate?
	  // XXX pmSourceExtFitPars *extFitPars = pmSourceExtFitParsAlloc();
	  // XXX extFitPars->Mxx = source->moments->Mxx;
	  // XXX extFitPars->Mxy = source->moments->Mxy;
	  // XXX extFitPars->Myy = source->moments->Myy;
	  // XXX extFitPars->Mrf = source->moments->Mrf;
	  // XXX extFitPars->Mrh = source->moments->Mrh;
	  // XXX extFitPars->peakMag = (source->peak->rawFlux > 0) ? -2.5*log10(source->peak->rawFlux) : NAN;
	  
	  // save kron mag, but assign apMag & psfMag on output (not yet calculated)
	  // XXX extFitPars->krMag = -2.5*log10(source->moments->KronFlux);

          // psArrayAdd (source->extFitPars, 4, extFitPars);
	  // psFree (extFitPars);

          // test for fit quality / result
	  modelFit->fitRadius = fitRadius;
          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;
            }
        }

	// clear the circular mask
	psImageMaskPixels (source->maskObj, "AND", PS_NOT_IMAGE_MASK(markVal)); 

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

	  // ensure the modelEXT is cached
	  pmSourceCacheModel (source, maskVal);

          pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);

          continue;
        }

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

	// adjust the window so the subtraction covers the faint wings
	psphotSetRadiusMoments(&fitRadius, &windowRadius, readout, source, markVal);

	// cache the model flux
	if (source->modelEXT->isPCM) {
	    pmPCMCacheModel (source, maskVal, psfSize);
	} else {
	    pmSourceCacheModel (source, maskVal);
	}

        // subtract the best fit from the object, leave local sky
        pmSourceSub (source, PM_MODEL_OP_FULL, maskVal);

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

    // change the value of a scalar on the array (wrap this and put it in psArray.h)
    scalar = job->args->data[7];
    scalar->data.S32 = Next;

    scalar = job->args->data[8];
    scalar->data.S32 = Nconvolve;

    scalar = job->args->data[9];
    scalar->data.S32 = NconvolvePass;

    scalar = job->args->data[10];
    scalar->data.S32 = Nplain;

    scalar = job->args->data[11];
    scalar->data.S32 = NplainPass;

    scalar = job->args->data[12];
    scalar->data.S32 = Nfaint;

    scalar = job->args->data[13];
    scalar->data.S32 = Nfail;

    return true;
}
