using namespace std;
using namespace RooFit;

void rfWoes()
{
    /////////Play with these parameters////////////////
    double rangeMin = 0.1;   //Bottom of fit range
    double rangeMax = 0.9;   //Top of fit range
    double sigInj = 0;       //Injected signal strength
    double sigWidth = 0.001; //Frozen signal width
    ///////////////////////////////////////////////////

    RooWorkspace* wksp = new RooWorkspace("wksp");

    wksp->factory("x[0.,1.]");
    ((RooRealVar*)wksp->var("x"))->setRange("fullRange",0.,1.);

    //The problem: define a subrange
    ((RooRealVar*)wksp->var("x"))->setRange("fitRange",rangeMin,rangeMax);

    //Create the background model
    wksp->factory("Exponential::bkgModel(x,k[-15.,-100.,0.])");

    //Create the signal model with a peak in the centre of the range
    wksp->factory("BreitWigner::sigModel(x,m["+TString::Format("%.4f",(rangeMin+rangeMax)/2.)+",0.,1.],w["+TString::Format("%.5f",sigWidth)+",0.,10.])");

    //Fix the signal width and pole mass
    ((RooRealVar*)wksp->var("m"))->setConstant(kTRUE);
    ((RooRealVar*)wksp->var("w"))->setConstant(kTRUE);

    //Create a bkgN*background + sigN*signal model
    wksp->factory("SUM::bkgModelPlusSigModel(bkgN[1000000.,0.,100000000.]*bkgModel, sigN[0.,-10000000.,10000000.]*sigModel)");

    wksp->pdf("bkgModelPlusSigModel")->fixAddCoefNormalization(RooArgSet(*((RooRealVar*)wksp->var("x")))); //This sometimes matters
    ((RooRealVar*)wksp->var("sigN"))->setVal(sigInj);

    //How many true bkg event in fitRange? bkgN*integral(bkgModel)
    double trueBkgEventsInRange = (((RooRealVar*)wksp->var("bkgN"))->getVal())*(((RooAbsReal*)wksp->pdf("bkgModel")->createIntegral(*((RooRealVar*)wksp->var("x")),NormSet(*((RooRealVar*)wksp->var("x"))),Range("fitRange")))->getVal());

    //How many true sig events in fitRange? sigN*integral(sigModel)
    double trueSigEventsInRange =(((RooRealVar*)wksp->var("sigN"))->getVal())*(((RooRealVar*)wksp->var("sigN"))->getVal())*(((RooAbsReal*)wksp->pdf("sigModel")->createIntegral(*((RooRealVar*)wksp->var("x")),NormSet(*((RooRealVar*)wksp->var("x"))),Range("fitRange")))->getVal());

    //Generate some asimov data
    RooDataHist *toyData = (RooDataHist*)wksp->pdf("bkgModelPlusSigModel")->generateBinned(*((RooRealVar*)wksp->var("x")),((RooRealVar*)wksp->var("bkgN"))->getVal(),ExpectedData(kTRUE),Extended(kTRUE));

    //How many total events in fitRange?
    double integral = toyData->sumEntries("x","fitRange");

    RooPlot* bkgPlusSigFitPlot = ((RooRealVar*)wksp->var("x"))->frame();

    //Plot the data
    toyData->plotOn(bkgPlusSigFitPlot,Name("data"), MarkerSize(1));

    //Configure normalization for fit range (problematic?)
    wksp->pdf("bkgModel")->setNormRange("fitRange");
    wksp->pdf("sigModel")->setNormRange("fitRange");
    wksp->pdf("bkgModelPlusSigModel")->setNormRange("fitRange");
    wksp->pdf("bkgModelPlusSigModel")->fixAddCoefRange("fitRange");

    //An extended likelihood fit using only bkg
    ((RooRealVar*)wksp->var("sigN"))->setVal(0.);
    ((RooRealVar*)wksp->var("sigN"))->setConstant(kTRUE);
    wksp->pdf("bkgModelPlusSigModel")->fitTo(*toyData,Range("fitRange"),Extended(kTRUE));

    double expectedEventsBkgOnly = wksp->pdf("bkgModelPlusSigModel")->expectedEvents(RooArgSet(*((RooRealVar*)wksp->var("x"))));
    double expectedBkgEventsBkgOnly = ((RooRealVar*)wksp->var("bkgN"))->getVal();

     //Plot the bkg only fit
    *wksp->pdf("bkgModelPlusSigModel")->plotOn(bkgPlusSigFitPlot,Name("bkgOnlyFit"),LineColor(kBlue),Range("fitRange"),Normalization(integral,RooAbsReal::NumEvent),NormRange("fitRange"));

    //An extended likelihood fit using bkg plus sig
    ((RooRealVar*)wksp->var("sigN"))->setConstant(kFALSE);
    wksp->pdf("bkgModelPlusSigModel")->fitTo(*toyData,Range("fitRange"),Extended(kTRUE));

    double expectedEventsBkgPlusSig = wksp->pdf("bkgModelPlusSigModel")->expectedEvents(RooArgSet(*((RooRealVar*)wksp->var("x"))));

    //Plot the bkg component
    *wksp->pdf("bkgModelPlusSigModel")->plotOn(bkgPlusSigFitPlot, Components("bkgModel"), Name("bkgComponent"),LineColor(kBlue),LineStyle(kDashed),Range("fitRange"),Normalization(integral,RooAbsReal::NumEvent),NormRange("fitRange"));

    //Plot the (positive) sig component
    if(((RooRealVar*)wksp->var("sigN"))->getVal() > 0.)
    {
        *wksp->pdf("bkgModelPlusSigModel")->plotOn(bkgPlusSigFitPlot, Components("sigModel"), Name("sigComponent"),LineColor(kRed),LineStyle(kDashed),Range("fitRange"),Normalization(integral,RooAbsReal::NumEvent),NormRange("fitRange"),Precision(1e-6));
    }

    //Plot the bkg plus sig model
    *wksp->pdf("bkgModelPlusSigModel")->plotOn(bkgPlusSigFitPlot, Name("bkgPlusSigFit"),LineColor(kGreen),Range("fitRange"),Normalization(integral,RooAbsReal::NumEvent),NormRange("fitRange"),Precision(1e-6));

    bkgPlusSigFitPlot->SetMinimum(1e-1);
    bkgPlusSigFitPlot->SetTitle("range = ["+TString::Format("%.3f",rangeMin)+","+TString::Format("%.3f",rangeMax)+"], sigWidth = "+TString::Format("%.5f",sigWidth)+", sigN = "+TString::Format("%.1f",((RooRealVar*)wksp->var("sigN"))->getVal()));

    //Draw on a canvas
    TCanvas* canvas = new TCanvas("plot","plot",800,600);
    canvas->Divide(1,2);
    canvas->cd(1);
    bkgPlusSigFitPlot->Draw();
    gPad->SetLogy();

    //Build a legend
    TLegend *legend = new TLegend(0.65,0.65,0.9,0.9);
    legend->SetFillColor(kWhite);
    legend->SetLineColor(kWhite);
    legend->AddEntry("data","Toy data", "P");
    legend->AddEntry("bkgOnlyFit","Bkg-only fit","L");
    legend->AddEntry("bkgPlusSigFit","Bkg+Sig fit", "L");
    legend->AddEntry("bkgComponent","Bkg component", "L");    
    if(((RooRealVar*)wksp->var("sigN"))->getVal() > 0.)
    {
        legend->AddEntry("sigComponent","Sig component", "L");
    }
    legend->Draw();

    //Draw pull plot
    RooPlot* pullPlot = ((RooRealVar*)wksp->var("x"))->frame(Title("pull plot"));
    pullPlot->addPlotable(bkgPlusSigFitPlot->residHist("data","bkgPlusSigFit",true),"p");
    canvas->cd(2);
    pullPlot->SetYTitle("(data-pdf)/error");
    pullPlot->Draw();
    TLine* xAxis = new TLine(0.,0.,1.,0.);
    xAxis->SetLineWidth(3);
    xAxis->Draw("same");
    gPad->SetGridy();

    canvas->Print("rfWoes_sigWidth_"+TString::Format("%.5f",sigWidth)+".pdf","pdf");

    //Results
    cout << "\n\n\n" << endl;
    cout << "-------------------------------Results in fit range-------------------------------------" << endl;
    cout << "                  |\tdata \t\t|\tbkgOnlyFit \t|\tbkgPlusSigFit" << endl;
    cout << "total events      |\t"<< integral << "\t\t|\t" << expectedEventsBkgOnly << "\t\t|\t" << expectedEventsBkgPlusSig << endl;
    cout << "background events |\t"<< trueBkgEventsInRange << "\t\t|\t" << expectedBkgEventsBkgOnly << "\t\t|\t" << ((RooRealVar*)wksp->var("bkgN"))->getVal() << endl;
    cout << "signal events     |\t"<< trueSigEventsInRange << "\t\t|\t" << 0.0000 << "\t\t|\t" << ((RooRealVar*)wksp->var("sigN"))->getVal() << endl;
    cout << "chi2              |\t" <<  0.0000 << "\t\t|\t" <<  bkgPlusSigFitPlot->chiSquare("bkgOnlyFit","data",1) << "\t|\t" << bkgPlusSigFitPlot->chiSquare("bkgPlusSigFit","data",2) << endl;
    cout << "----------------------------------------------------------------------------------------\n\n\n";
}