using namespace RooFit;
TRandom3 rnd;

void buildModel(RooWorkspace *wspace)
{
    TFile *fin = new TFile("example_splusb.root");
    TNtupleD* nt = (TNtupleD *)fin->Get("nt");
    
    RooRealVar mass("mass","mass obs",0.,2.);
    
    RooRealVar mu("mu","signal mean",1.0);
    RooRealVar sigma("sigma","signal width",0.05);
    RooGaussian gaus("gaus","signal PDF",mass,mu,sigma);
    
    RooRealVar slope("slope","background slope",-0.3,-5.,5.);
    RooPolynomial linear("linear","background PDF",mass,RooArgSet(slope));
    
    RooRealVar ns("ns","ns",10,0.,1000.);
    RooRealVar nb("nb","nb",90,0.,1000.);
    RooAddPdf model("model","mass PDF",RooArgList(gaus,linear),RooArgList(ns,nb));
    
    RooRealVar slope_mu("slope_mu","slope mean",-0.3);
    RooRealVar slope_sigma("slope_sigma","slope sigma",0.03);
    RooGaussian cons("cons","constrained PDF",slope,slope_mu,slope_sigma);
    
    RooDataSet data("data","data",nt,RooArgSet(mass));
    
    wspace->import(model);
    wspace->import(cons);
    wspace->import(data);
    
    delete fin;
}

void freq_toy(double target_ns = 10., int ntoys = 100)
{
    RooWorkspace *wspace = new RooWorkspace("wspace");
    buildModel(wspace);

    wspace->var("ns")->setVal(target_ns);
    wspace->var("ns")->setConstant(true);
    wspace->pdf("model")->fitTo(*wspace->data("data"),ExternalConstraints(*wspace->pdf("cons")));
    
    TH1D *pull_ns = new TH1D("pull_ns","",60.,-6.,6.);
    TH1D *pull_slope = new TH1D("pull_slope","",60.,-6.,6.);
                            
    for (int idx=0; idx<ntoys; idx++) {
        RooWorkspace* ws_toy = new RooWorkspace(*wspace);
        
        // main data
        RooDataSet *toy = ws_toy->pdf("model")->generate(*ws_toy->var("mass"),Extended(true));
            
        // complementary data
        ws_toy->var("slope_mu")->setVal(rnd.Gaus(ws_toy->var("slope")->getVal(),ws_toy->var("slope_sigma")->getVal()));
            
        ws_toy->var("ns")->setConstant(false);
        ws_toy->pdf("model")->fitTo(*toy,ExternalConstraints(*ws_toy->pdf("cons")));
        
        pull_ns->Fill((ws_toy->var("ns")->getVal()-wspace->var("ns")->getVal())/ws_toy->var("ns")->getError());
        pull_slope->Fill((ws_toy->var("slope")->getVal()-wspace->var("slope")->getVal())/ws_toy->var("slope")->getError());

        delete toy;
        delete ws_toy;
    }
    
    TCanvas *c1 = new TCanvas("c1","",400,800);
    c1->Divide(1,2);
    for (auto& hist : {pull_ns, pull_slope}) {
        hist->SetStats(false);
        hist->SetFillColor(41);
        hist->GetYaxis()->SetTitle("Entries");
        hist->GetXaxis()->SetTitleSize(0.06);
        hist->GetYaxis()->SetTitleSize(0.06);
        hist->GetYaxis()->SetTitleOffset(1.0);
    }
    c1->cd(1);
    c1->GetPad(1)->SetMargin(0.14,0.06,0.13,0.07);
    pull_ns->GetXaxis()->SetTitle("Pull(ns)");
    pull_ns->Fit("gaus","L");
    c1->cd(2);
    c1->GetPad(2)->SetMargin(0.14,0.06,0.13,0.07);
    pull_slope->GetXaxis()->SetTitle("Pull(slope)");
    pull_slope->Fit("gaus","L");
    
    delete wspace;
}

void example_03()
{
    freq_toy(10.,1000);
}