#include "RooAbsPdf.h"
#include "RooRealVar.h"
#include "RooDataHist.h"
#include "RooDataSet.h"
#include "RooPlot.h"
#include "RooWorkspace.h"
#include "TFile.h"

using namespace RooFit;

void test(bool load = true){
  TString fileName("mass_workspace.root");
  RooWorkspace *myWS;
  if(load){
    TFile *in_file = TFile::Open(fileName);
    myWS  = (RooWorkspace*)in_file->Get("myWS");
  }else{
    myWS = new RooWorkspace("myWS");  
    myWS->factory("mass[1764.840,1964.840]");

    // Pdf in observable 
    myWS->factory("Gaussian::sigPdf(mass,1864,sigSigma[0,200])") ;
    myWS->var("sigSigma")->setVal(10);
    myWS->factory("Uniform::bkgPdf(mass)") ;
    myWS->factory("SUM::model(S[1000,0,100000]*sigPdf,B[1000,0,100000]*bkgPdf)") ;

    // Background only pdf
    // myWS->factory("ExtendPdf::modelBkg(bkgPdf,B)") ;
    // // Priors
    // myWS->factory("Gaussian::prior_sigSigma(sigSigma, 5,5)") ;
    // myWS->factory("Gaussian::prior_B(B,1000,200)") ;
    // myWS->factory("PROD::priorNuisance(prior_sigSigma,prior_B)") ;
    // myWS->factory("Uniform::priorPOI(S)") ;

    // Definition of observables and parameters of interest
    myWS->defineSet("observables","mass");
    //myWS->defineSet("parameters","B,sigSigma");
    //myWS->defineSet("POI","S");
    myWS->writeToFile(fileName);
    
  }
  myWS->Print();
  // Var values 
  cout <<"S = " <<  myWS->var("S")->getVal() << endl;
  cout <<"B = " <<  myWS->var("B")->getVal() << endl;
  cout <<"sigSigma = " <<  myWS->var("sigSigma")->getVal() << endl;
  
  
  RooAbsData* data = 0;
  TString data_name = "toy_data";
  //  myWS->var("sigSigma")->setVal(10);
  myWS->var("S")->setVal(100);
  
  data = myWS->pdf("model")->generateBinned(*myWS->set("observables"), ExpectedData(1), Name(data_name), Verbose(1)); // add Extended() for poisson fluct
  //data = myWS->pdf("model")->generate(*myWS->set("observables"), 10000,Name(data_name)); // add Extended() for poisson fluct
  
  RooPlot* plot = myWS->var("mass")->frame();
  data->plotOn(plot);
  myWS->pdf("model")->plotOn(plot);
  myWS->pdf("sigPdf")->plotOn(plot);
  myWS->pdf("bkgPdf")->plotOn(plot);

  plot->DrawClone();

}