#include<iostream>
#include<vector>

#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooExponential.h"
#include "RooAddPdf.h"
#include "RooExtendPdf.h"
#include "RooPlot.h"

#include "TFile.h"

using namespace RooFit;

//#define GenerateDataSet
#define Blinding

void GenerateRootFile(RooExponential *Exponential, RooRealVar *mass)
{
  TFile *outFile      = new TFile("ToyExponential.root","RECREATE");
  std::cout << "Generate toy set" << std::endl;
  RooDataSet *ExpoSet = Exponential->generate(*mass,10000);
  ExpoSet->Write();
  outFile->Close();
}

void FitData(RooExponential *Exponential, RooRealVar *mass)
{
  TFile *inFile         = new TFile("ToyExponential.root","OPEN");
  RooDataSet *ExpoSet   = (RooDataSet*)inFile->Get("ExpoData");
  std::cout << "Extend PDF" << std::endl;
  RooRealVar *nExpo     = new RooRealVar("nExpo","nExpo",10000,0,100000);
  RooExtendPdf *ExtExpo = new RooExtendPdf("ExtExpo","ExtExpo",*Exponential,*nExpo);
  
  TString fitRange = "allRange";
#ifdef Blinding
  fitRange = "lowerSB,upperSB";
#endif
  
  std::cout << "Going for the fit" << std::endl;
  ExtExpo->fitTo(*ExpoSet,Extended(kTRUE),Save(),Range(fitRange),PrintLevel(3));

  std::cout << "Plotting Result" << std::endl;
  RooPlot *frame = mass->frame("mass","mass");
  ExpoSet->plotOn(frame,Binning(30));
  ExtExpo->plotOn(frame,LineColor(kBlue),Range(fitRange),NormRange(fitRange));
  frame->Draw();
}

int RooFitFit()
{
  RooRealVar *mass = new RooRealVar("mass","mass",0,10);
  mass->setRange("allRange",0,10);
#ifdef Blinding
  mass->setRange("lowerSB",0,3);
  mass->setRange("upperSB",6,10);
#endif
  Double_t aExpo = -.5;
  RooRealVar *aExp = new RooRealVar("aExp","aExp",aExpo,-100,100);
  
  RooExponential *Expo = new RooExponential("Expo","Expo",*mass,*aExp);

#ifdef GenerateDataSet
  GenerateRootFile(Expo,mass);
#else
  FitData(Expo,mass);
#endif

  return 0;
}