//=============================================================================
/** 
 * @file RooFitToyGen.C 
 * @brief Get familar with RooFit toy generation
 *
 * @author Tim-Philip Huecking, <huecking@hep.physik.uni-siegen.de>
 * 
 * Created:
 * @date 2019-10-17
 * Changed:
 *
 */
//=============================================================================

using namespace std; 
using namespace RooFit; 

void RooFitIncludeWeightBug()
{   
    //observable
    RooRealVar mass( "mass" , "mass" , 0 , 150 , "a.u." );
    
    //event numbers
    RooRealVar nSigExp("nSigExp", "Number of expected signal events", 20000
            , 0 , 10000000 ); 
    RooRealVar nBkgExp("nBkgExp", "Number of expected background events", 100000
             , 0 , 1000000);
    
    //setup sig and bkg pdfs
    RooRealVar sigMean("sigMean", "Mean of signal Gaussian", 60 , -100 , 200 ); 
    RooRealVar sigSigma("sigSigma", "Sigma of signal Gaussian",4 , 0, 100   ); 
    RooRealVar bkgTau("bkgTau","Tau of background exponential", -1/40. , -100 , -1E-10 ); 

    RooGaussian sigPDF("sigPDF", "PDF of the signal", mass ,sigMean , sigSigma ); 
    RooExponential bkgPDF("bkgPDF" , "PDF ot the background" , mass, bkgTau ); 
    
    // make single PDFs extended 
    RooExtendPdf eSigPDF("eSigPDF", "Extended signal PDF", sigPDF , nSigExp ); 
    RooExtendPdf eBkgPDF("ebkgPDF", "Extended bkg PDF", bkgPDF , nBkgExp ); 

    // create full model 
    RooAddPdf massPDF("massPDF", "Total PDF mass", RooArgList(sigPDF,bkgPDF),
            RooArgList(nSigExp,nBkgExp) );
    RooAddPdf* pmassPDF = &massPDF;
    
    // generate data
    RooDataSet* toyData1 = massPDF.generate(mass, Extended(kTRUE));
    RooDataSet* toyData2 = new RooDataSet(*toyData1, "toyData2");
    RooDataSet* toyData3 = new RooDataSet(*toyData1, "toyData3");
    RooDataSet* toyData4 = new RooDataSet(*toyData1, "toyData4");

    // perform a fit 
    auto fitRes = massPDF.fitTo(*toyData1, RooFit::Extended(kTRUE), Save(kTRUE)); 
    fitRes->Print(); 

    // 
    // do sPlots    
    // 
    
    cout << "Data set contents before doing sPlots: " << endl;
    toyData1->Print();

    // use AddSWeight method
    RooStats::SPlot splot("splot1","splot1");
    splot.SetSData(toyData1);
    splot.AddSWeight(pmassPDF, 
            RooArgList(nSigExp, nBkgExp), RooArgSet(), kTRUE);
    RooStats::SPlot splot2("splot2","splot2");
    splot2.SetSData(toyData2);
    splot2.AddSWeight(pmassPDF, 
            RooArgList(nSigExp, nBkgExp), RooArgSet(), kFALSE);

    // do everything in constructor
    RooStats::SPlot sploti3("splot3","splot3", *toyData3, pmassPDF, 
            RooArgList(nSigExp, nBkgExp), RooArgSet(), kTRUE);
    RooStats::SPlot sploti4("splot4","splot4", *toyData4, pmassPDF, 
            RooArgList(nSigExp, nBkgExp), RooArgSet(), kFALSE);
    
    // print results
    cout << endl << "==========" << endl;
    cout << "The weights are indicated by *_sw" << endl; 
    cout << "toyData1 (Expect weights to be added):" << endl;
    toyData1->Print();
    cout << "toyData2 (Expect weights not to be added)" << endl;
    toyData2->Print();
    cout <<"toydata 3 using constructor (Expect weights to be added)" << endl; 
    toyData3->Print();
    cout <<"toydata 4 using constructor (Expect weights not to be added)" << endl; 
    toyData4->Print();
        
}