#include "RooWorkspace.h"
#include "RooAbsPdf.h"
#include "RooDataSet.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "RooRealVar.h"
#include "RooRandom.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooChebychev.h"
#include "RooAddPdf.h"
#include "TFile.h"
#include "TCanvas.h"

#include "RooStats/ModelConfig.h"
#include "RooStats/HybridCalculator.h"
#include "RooStats/FrequentistCalculator.h"
#include "RooStats/ToyMCSampler.h"
#include "RooStats/HypoTestPlot.h"

#include "RooStats/NumEventsTestStat.h"
#include "RooStats/ProfileLikelihoodTestStat.h"
#include "RooStats/SimpleLikelihoodRatioTestStat.h"
#include "RooStats/RatioOfProfiledLikelihoodsTestStat.h"
#include "RooStats/MaxLikelihoodEstimateTestStat.h"

#include "RooStats/HypoTestInverter.h"
#include "RooStats/HypoTestInverterResult.h"
#include "RooStats/HypoTestInverterPlot.h"


using namespace RooFit;
using namespace RooStats;


void GausCheb_BRConv_Model(int nsig = 0,    // number of signal events 
                  int nbkg = 150 ,  // number of background events
                  double fErrPer = 0.05 )  // Percentage error of the const factor f
{ 

  TString saveName = Form("BRConv_NP_a20a21_nSig%d_nBkg%d",nsig,nbkg);

  RooRealVar Mass("Mass","Mass(#chi_{b}) (GeV/c^{2})", 7.11, 12.61,"");
  
  RooRealVar a20("a20", "Par a20", -7.8e-01, -50., 50.);
  RooRealVar a21("a21", "Par a21", 4.0e-01, -50., 50.);
  
  RooChebychev *bkg_pdf = new RooChebychev("bkg_pdf","bkg_pdf", Mass, RooArgSet(a20,a21));
  
  RooRealVar mean("mean","mean", 9.859); //mean - variable
  RooRealVar sigma("sigma","sigma", 0.07); //sigma - variable
  RooGaussian* sig_pdf = new RooGaussian("sig_pdf", "sig_pdf", Mass, mean, sigma);

  
  double nbkgErr = 0.8*sqrt(nbkg);
  RooRealVar Nbkg("Nbkg", "Nbkg", nbkg, nbkg-3.*nbkgErr, nbkg+3.*nbkgErr);
  
  double f =  1.5581;
  double sf =  f*(1.+fErrPer);

  cout << "\n\tConstant: " << f << " +- " << sf << "\n" << endl;
  
  RooRealVar SES("SES","SES", f, f-5*sf, f+5*sf);

  RooRealVar SESMean("SESMean","SESMean",f);
  RooRealVar SESErr("SESErr","SESErr",sf);
  RooAbsPdf* gErrSES = new RooGaussian("gErrSES","gErrSES", SES, SESMean, SESErr);
  
  cout << "\n\tSES: " << SES.getVal() << " +- " << sf << "\n" << endl;
  
  RooRealVar BR("BR","BR [x10^{-4}]", 0, 0, 20, "BR [x10^{-4}]");
  RooFormulaVar BRconv("BRconv","@0*@1",RooArgList(BR,SES));
  
  /// PDF : BRconv * sig_pdf + Nbkg * bkg_pdf
  ///      sig_pdf = Gaus(M|m,s)
  ///      bkg_pdf = Cheb(M|a20,a21)
  ///      BRconv =  BR * SES
  ///      
  RooAddPdf* model = new RooAddPdf("model","model", RooArgList(*sig_pdf,*bkg_pdf),RooArgList(BRconv,Nbkg));

  /// Import staff to a RooWorkspace
  RooWorkspace* w = new RooWorkspace("w");
  w->import(BR);
  w->import(*model);
  
  /// Define our observable and parameters
  w->defineSet("obs",RooArgSet(Mass));
  w->defineSet("poi",RooArgSet(BR));

  w->var("BR")->setVal(nsig/SES.getVal());
  w->var("Nbkg")->setVal(nbkg);
  
  
  /// generate the data
  /// use fixed random numbers for reproducibility (use 0 for changing every time)
  RooRandom::randomGenerator()->SetSeed(111);

  Mass.setBins(55);

  RooDataSet * data = model->generate( Mass);  // will generate accordint to total S+B events
  data->SetName("obsData");
  w->import(*data);

  data->Print();
  TCanvas *c_fit = new TCanvas(Form("c_GausChebModel_%s.root",saveName.Data()), Form("c_GausChebModel_%s.root",saveName.Data()), 800,600);

  RooPlot * plot = Mass.frame(Title("Gaussian Signal over Polynomial Background"));
  data->plotOn(plot);
  plot->Draw();

  RooFitResult * r = model->fitTo(*data, RooFit::Save(true), RooFit::Minimizer("Minuit2","Migrad"));
  r->Print();

  model->plotOn(plot);
  /// draw the two separate pdf's
  model->plotOn(plot, RooFit::Components("bkg_pdf"), RooFit::LineStyle(kDashed) );
  model->plotOn(plot, RooFit::Components("sig_pdf"), RooFit::LineColor(kRed), RooFit::LineStyle(kDashed) );

  model->paramOn(plot,Layout(0.5,0.9,0.85));

  plot->Draw();

  RooStats::ModelConfig mc("ModelConfig",w);
  mc.SetPdf(*model);
  mc.SetParametersOfInterest(*w->set("poi"));
  mc.SetObservables(*w->var("Mass"));
  
  /// define set of nuisance parameters
  w->defineSet("nuisParams","a20,a21");

  mc.SetNuisanceParameters(*w->set("nuisParams"));

  /// import model in the workspace 
  w->import(mc);

  /// write the workspace in the file
  TString fileName = Form("GausChebModel_%s.root",saveName.Data());
  w->writeToFile(fileName,true);
  cout << "model written to file " << fileName << endl;
   
   
}