//Last edited: May 20, 2004
//David Faden
//dfaden@cs.iastate.edu

//This code depends on the RooFit libraries. It assumes that they
//have been loaded prior to its running. You can do this with the
//following:
//  gSystem->Load("~dfaden/roofit/lib/libRooFitCore.so");
//  gSystem->Load("~dfaden/roofit/lib/libRooFitModels.so");

/**
 * Fit the mass using only the signal + background data set.
 * First fit an exponential background to the a cut of the data
 * set computed using the given cut. Then fit the sum of the
 * background and a Gaussian to the full data set.
 *
 * @param minMass The minimum mass value to read and fit to.
 * @param maxMass The maximum mass value to read and fit to.
 * @param bgCut The cut of the data set read to fit the background to.
 */
void displayFitOnMassWithSplitFitBackground(const double minMass = 2,
					    const double maxMass = 6,
					    const char* bgCut = 
					    "mass < 2.6 || mass > 3.4")
{
  RooRealVar mass("mass", "mass", minMass, maxMass);
  RooDataSet* signalPlusBackgroundData =
    RooDataSet::read("data/mass_bg_eq_0.txt", RooArgSet(mass), "Q");
  
  const char* name = getUniqueName("backgroundData", 0);
  RooDataSet* backgroundData =
    new RooDataSet(name, "background", signalPlusBackgroundData,
		   RooArgSet(mass), bgCut);
  delete[] name;

  RooRealVar c("c", "c", -0.8, -4, -0.01);
  RooExponential expPdf("expPdf", "Exponential Background PDF", mass, c);

  fitToBackground(expPdf, mass, backgroundData);
  delete backgroundData;

  fixVarOtherThanMassOfPdf(expPdf);
  
  RooRealVar mean("mean", "mean", 3.1, 0, 6);
  RooRealVar sigma("sigma", "sigma", 0.1, 0.0001, 2);

  RooGaussian gaussian("gaussian", "J/psi", mass, mean, sigma);

  fitSumToSignalPlusBackground(gaussian, expPdf, mass,
			       signalPlusBackgroundData);
  delete signalPlusBackgroundData;
}


/**
 * Fit the sum of a Gaussian and a/(b + mass**n) to a
 * set of mass data points and display the result.
 * 
 * @param minMass The minimum mass value to read and fit to.
 * @param maxMass The maximum mass value to read and fit to.
 * @param readOptions Options used for reading data. "Q" for quiet,
 * "D" for extra debugging info.
 * @param fixBackgroundPdf If true, make all the parameters of the
 * background PDF constant after fitting to the background data and
 * before fitting to the signal + background data.
 */
void displayFitOnMassWithInverseMassBackground(const double minMass = 2,
					       const double maxMass = 6,
					       const char* readOptions = "Q",
					       Bool_t fixBackgroundPdf = kTRUE)
{
  RooRealVar mass("mass", "mass", minMass, maxMass);
  RooRealVar a("a", "a", 0.0000001, 10);
  RooRealVar b("b", "b", 0.0000001, 10);
  RooRealVar n("n", "n", 0.0000001, 10);
  RooGenericPdf background("background", "(a/b + mass**n) Background PDF",
  		      "a/(b + mass**n)", RooArgList(a,b,n,mass));
  displayFitOnMass(background, readOptions, fixBackgroundPdf);
}

/**
 * Fit the sum of a Gaussian and an exponential to a
 * set of mass data points and display the result.
 * 
 * @param minMass The minimum mass value to read and fit to.
 * @param maxMass The maximum mass value to read and fit to.
 * @param readOptions Options used for reading data. "Q" for quiet,
 * "D" for extra debugging info.
 * @param fixBackgroundPdf If true, make all the parameters of the
 * background PDF constant after fitting to the background data and
 * before fitting to the signal + background data.
 */
void displayFitOnMassWithExpBackground(const double minMass = 2,
				       const double maxMass = 6,
				       const char* readOptions = "Q",
				       Bool_t fixBackgroundPdf = kTRUE)
{
  RooRealVar mass("mass", "mass", minMass, maxMass);
  RooRealVar c("c", "c", -0.8, -4, -0.01);

  RooExponential expPdf("expPdf","Exponential Background PDF", mass, c);
  displayFitOnMass(expPdf, readOptions, fixBackgroundPdf);
}

/**
 * Fit the sum of a Gaussian and the given background pdf to a
 * set of mass data points and display the result.
 * 
 * @param minMass The minimum mass value to read and fit to.
 * @param maxMass The maximum mass value to read and fit to.
 * @param readOptions Options used for reading data. "Q" for quiet,
 * "D" for extra debugging info.
 * @param fixBackgroundPdf If true, make all the parameters of the
 * background PDF constant after fitting to the background data and
 * before fitting to the signal + background data.
 */
void displayFitOnMass(RooAbsPdf& background, const char* readOptions, Bool_t fixBackgroundPdf)
{
  RooRealVar mass(getMassFromPdf(background));
  mass.setConstant(kFALSE);

  RooDataSet* backgroundData;
  RooDataSet* signalPlusBackgroundData;
  readData(&backgroundData, &signalPlusBackgroundData, mass, readOptions);

  fitToBackground(background, mass, backgroundData);
  delete backgroundData;

  if (fixBackgroundPdf) {
    fixVarOtherThanMassOfPdf(background);
  }

  RooRealVar mean("mean", "mean", 3.1, 0, 6);
  RooRealVar sigma("sigma", "sigma", 0.1, 0.0001, 2);

  RooGaussian gaussian("gaussian", "J/psi", mass, mean, sigma);

  fitSumToSignalPlusBackground(gaussian, background,
			       mass, signalPlusBackgroundData);
  delete signalPlusBackgroundData;
}

//BUG: fitSumToSignalPlusBackground and fitToBackground
//have very similar code. The common parts should be
//factored out into a new function.

/**
 * Fit the sum of the given pdfs to the given data set. Open
 * a new canvas displaying the result.
 */
void fitSumToSignalPlusBackground(RooAbsPdf& signal, RooAbsPdf& background,
				  const RooRealVar& mass,
				  const RooDataSet* signalPlusBackgroundData)
{
  RooRealVar numberOfJPsi("numberOfJ", "Number of J/psi",
			  5,
			  0,
			  100000);
  RooExtendPdf extendedSignal("extendedSignal", "Extended Signal PDF",
			      signal, numberOfJPsi);

  RooRealVar numberOfBackground("numberOfBackground",
				"Number of background",
				0,
				100000);
  RooExtendPdf extendedBackground("extendedBackgroundPdf",
				  "Extended Background PDF",
				  background,
				  numberOfBackground);
  

  //intro7.cc line 45 tells us to note that we don't include
  //coefficients when adding extended PDFs.
  RooAddPdf sum("sum", "sum", RooArgList(extendedSignal,
					 extendedBackground));

  sum.fitTo(*signalPlusBackgroundData, "mhe");

  //BUG: We are currently leaking the memory for frame.
  //It may not be worth the trouble of fixing this leak though.
  //The best way to do so would probably be to get notice of the
  //closing of the frame via ROOT's signal/slot mechanism.
  //(Deleting frame at the end of this function deletes the
  //contents of the display, leaving a blank canvas.)
  RooPlot* frame = mass.frame();
  signalPlusBackgroundData->plotOn(frame);

  extendedSignal.plotOn(frame,
			Normalization(1.0,RooAbsReal::RelativeExpected),
			LineColor(kRed));
  extendedBackground.plotOn(frame,
			    Normalization(1.0,RooAbsReal::RelativeExpected),
			    LineColor(kGreen));
  sum.plotOn(frame,
	     Normalization(1.0,RooAbsReal::RelativeExpected),
	     LineColor(kBlue));

  //We create our own canvas so that the display function
  //may be called several times in one ROOT session and have
  //a separate canvas for each call. Before
  //the following code was added, each call to this function
  //would draw over the canvas put up by the previous call.
  makeNewDefaultCanvas(mass, background.ClassName());
  frame->Draw();
}

/**
 * Fit the given pdf to the given data set. Open a new canvas
 * displaying the result.
 */
void fitToBackground(RooAbsPdf& pdf, const RooRealVar& mass,
		     const RooDataSet* backgroundData)
{

  //I am removing the "e" option as I was getting lots of warnings
  //about the pdf not supporting "extended maximum likelihood".
  //Without the "e" option, we do not get these warnings.
  pdf.fitTo(*backgroundData, "mh");

  //BUG: We're currently leaking the memory of the new canvas
  //and the frame.
  makeNewDefaultCanvas(mass, "background only data", pdf.ClassName());
  RooPlot* frame = mass.frame();
  backgroundData->plotOn(frame);
  pdf.plotOn(frame, LineColor(kGreen));
  frame->Draw();
}

/**
 * Read the background and signal+background data sets. Trim the larger
 * of the two (by dropping data points that happen to occur later in the
 * data) so that they are equal in size.
 *
 * On a successful return from this function, the pointers pointed to
 * by backgroundData and signalPlusBackgroundData will have
 * been set to new data sets. It's the caller's responsibility to
 * delete the memory of the new data set objects.
 */
void readData(RooDataSet** backgroundData,
	      RooDataSet** signalPlusBackgroundData,
	      const RooRealVar& mass, const char* readOptions)
{
  *backgroundData = 0;
  *signalPlusBackgroundData = 0;

  RooDataSet* bgData =
    RooDataSet::read("data/mass_bg_gt_0.txt", RooArgSet(mass), readOptions);

  RooDataSet* signalPlusBgData =
    RooDataSet::read("data/mass_bg_eq_0.txt", RooArgSet(mass), readOptions);

  Int_t numberOfSignalPlusBg = signalPlusBgData->numEntries();
  Int_t numberOfBg = bgData->numEntries();

  if (numberOfSignalPlusBg > numberOfBg) {
    RooDataSet* newData
      = firstNMassDataPoints(numberOfBg, signalPlusBgData);
    delete signalPlusBgData;
    signalPlusBgData = newData;
  }
  else if (numberOfBg > numberOfSignalPlusBg) {
    RooDataSet* newData
      = firstNMassDataPoints(numberOfSignalPlusBg, bgData);
    delete bgData;
    bgData = newData;
  }

  *backgroundData = bgData;
  *signalPlusBackgroundData = signalPlusBgData;
}

/**
 * Get the RooRealVar named mass from the given PDF.
 */
RooRealVar& getMassFromPdf(const RooAbsPdf& pdf)
{
  RooRealVar varNamedMass("mass", "mass", -99999, -99999);
  RooArgSet* deps = pdf.getDependents(RooArgSet(varNamedMass));
  TIterator* iterator = deps.createIterator();
  RooRealVar* mass = dynamic_cast<RooRealVar*>(iterator->Next());
  delete iterator;
  return *mass;
}

/**
 * Make all of the vars of the given pdf other than
 * mass constant.
 */
void fixVarOtherThanMassOfPdf(RooAbsPdf& pdf)
{
  RooRealVar varNamedMass("mass", "mass", -99999, -99999);
  RooArgSet* param = pdf.getParameters(RooArgSet(varNamedMass));
  TIterator* iterator = param.createIterator();
  RooRealVar* var;
  while (var = dynamic_cast<RooRealVar*>(iterator->Next())) {
    var->setConstant(kTRUE);
  }
  delete iterator;
}

/**
 * Return a new data set consisting of the first n
 * data points of the given data set. The caller is
 * responsible for deleting the returned object.
 */
RooDataSet* firstNMassDataPoints(Int_t n, const RooDataSet* dataSet)
{
  const char* name = getUniqueName(dataSet->GetName(), 0);
  RooDataSet* firstN =
    new RooDataSet(name, dataSet->GetTitle(), *(dataSet->get()));
  delete[] name;

  for (Int_t i = 0; i < n; i++) {
    firstN->add(*(dataSet->get(i)));
  }
  return firstN;
}

/**
 * Create a new canvas guaranteed to have a unique name in ROOT (and so not
 * overwrite a previous object), open it, and make it the default
 * canvas for drawing operations.
 * 
 * @param mass RooRealVar that specifies the min and max values to display
 * in the title.
 * @param suffix1 Suffix to append to title.
 * @param suffix2 Suffix to append to title, after suffix1.
 */
void makeNewDefaultCanvas(const RooRealVar& mass, const char* suffix1 = "", const char* suffix2 = "")
{
  const char* canvasName = getUniqueName("mass", 0);
  char title[100];
  sprintf(title, "Mass (%f, %f) %s %s", mass.getFitMin(), mass.getFitMax(), suffix1, suffix2);

  //BUG: We are currently leaking the memory for canvas.
  TCanvas* canvas = new TCanvas(canvasName, title, 640, 480);

  canvas->Show();
  
  //cd() makes canvas the current canvas for drawing.
  canvas->cd();
  delete[] canvasName;
}

//The following code comes from offline/packages/tec

/**
 * getUniqueName returns a name that is unique with regards to
 * ROOT objects. The returned name will be of the form
 * (baseName) (positive integer).
 * <p>
 * The caller is responsible for deleting the memory of the
 * returned C string.
 * 
 * @param baseName Valid prefix for a variable name, such as "hist".
 * @param baseNumber Starting point for suffix; the more unique, the better.
 * @return A name unique in ROOT.
 */
char* getUniqueName(const char* baseName, int baseNumber)
{
  //XXX! Would it better to return a TString stack object?

  Int_t number = abs(baseNumber);
  Bool_t notUnique = kTRUE;
  char* name = 0;

  while (notUnique) {
    const char* cname;
    TString nameString(baseName);
    nameString += number;

    cname = (const char*) nameString;
    if (!gROOT->FindObject(cname)) {
      notUnique = kFALSE;

      //All of the work with cname must be
      //done here, before nameString goes
      //out of scope.
      name = new char[strlen(cname) + 1];
      strcpy(name, cname);
    }
    else {
      number++;
    }
  }

  return name;
}