using namespace RooFit;

void buildModel(RooWorkspace *wspace)
{
    TFile *fin = new TFile("example_06.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","full PDF",RooArgList(gaus,linear),RooArgList(ns,nb));
    
    RooDataSet data("data","data",nt,RooArgSet(mass));
    
    wspace->import(model);
    wspace->import(data);

    delete fin;
}

double FC_scan(double target_ns = 10., int ntoys = 100, bool display = false)
{
    RooWorkspace *wspace = new RooWorkspace("wspace");
    buildModel(wspace);

    RooFitResult *res0 = wspace->pdf("model")->fitTo(*wspace->data("data"),Save(true));
    wspace->var("ns")->setVal(target_ns);
    wspace->var("ns")->setConstant(true);
    RooFitResult *res1 = wspace->pdf("model")->fitTo(*wspace->data("data"),Save(true));
    double logR_data = res0->minNll()-res1->minNll();
    
    TH1D *h_teststat = new TH1D("h_teststat","",50,-5.,0.);
    double beta = 0.;
    for (int idx=0; idx<ntoys; idx++) {
        RooWorkspace* wspace_toy = new RooWorkspace(*wspace);
        RooDataSet *toy = wspace_toy->pdf("model")->generate(*wspace_toy->var("mass"));
        
        RooFitResult *res1 = wspace_toy->pdf("model")->fitTo(*toy,Save(true));
        wspace_toy->var("ns")->setConstant(false);
        RooFitResult *res0 = wspace_toy->pdf("model")->fitTo(*toy,Save(true));
        
        double logR = res0->minNll()-res1->minNll();
        if (logR>logR_data) beta += 1.;
        h_teststat->Fill(logR);
        
        delete res0;
        delete res1;
        delete toy;
        delete wspace_toy;
    }
    beta /= (double)ntoys;
    
    if (display) {
        TCanvas *c1 = new TCanvas("c1","",600,600);
        c1->SetMargin(0.14,0.06,0.13,0.07);
        c1->SetLogy();
        h_teststat->SetStats(false);
        h_teststat->SetFillColor(41);
        h_teststat->GetXaxis()->SetTitle("log(R)");
        h_teststat->GetYaxis()->SetTitle("# of toys");
        h_teststat->GetXaxis()->SetTitleSize(0.06);
        h_teststat->GetYaxis()->SetTitleSize(0.06);
        h_teststat->GetYaxis()->SetTitleOffset(1.0);
        h_teststat->Draw();
        TLine l1;
        l1.SetLineWidth(3);
        l1.DrawLine(logR_data,0.,logR_data,h_teststat->GetMaximum());
    } else delete h_teststat;
    
    delete res0;
    delete res1;
    delete wspace;
    
    return beta;
}

void example_08()
{
    RooWorkspace *wspace = new RooWorkspace("wspace");
    buildModel(wspace);
    
    RooFitResult *res0 = wspace->pdf("model")->fitTo(*wspace->data("data"),Save(true),Minos(true));
    
    TH1D *scan_beta = new TH1D("scan_beta","",101,-0.1,20.1);
    TH1D *scan_fc = new TH1D("scan_fc","",101,-0.1,20.1);
    
    for (int i=1; i<=scan_beta->GetNbinsX(); i++) {
        wspace->var("ns")->setVal(scan_beta->GetBinCenter(i));
        wspace->var("ns")->setConstant(true);
        RooFitResult *res1 = wspace->pdf("model")->fitTo(*wspace->data("data"),Save(true));
        double d2NLL = (res1->minNll()-res0->minNll())*2.;
        scan_beta->SetBinContent(i,1.-TMath::Prob(d2NLL,1));
        
        if ((i%5)==1) {
            double beta = FC_scan(scan_beta->GetBinCenter(i),100);
            double beta_err = sqrt(beta*(1.-beta)/100);
            scan_fc->SetBinContent(i,beta);
            scan_fc->SetBinError(i,beta_err);
        }
    }
    
    TCanvas *c1 = new TCanvas("c1","",600,600);
    c1->SetMargin(0.14,0.06,0.13,0.07);
    
    RooRealVar* best_ns = (RooRealVar*)res0->floatParsFinal().find("ns");
    
    TLine lin;
    lin.SetLineColor(kGray+2);
    lin.SetLineStyle(kDashed);
    
    TBox box;
    box.SetFillColor(kRed-10);
    
    scan_beta->SetStats(false);
    scan_beta->SetLineWidth(2);
    scan_beta->GetXaxis()->SetTitle("Signal Yield");
    scan_beta->GetYaxis()->SetTitle("Confidence Level (#beta)");
    scan_beta->GetXaxis()->SetTitleSize(0.06);
    scan_beta->GetYaxis()->SetTitleSize(0.06);
    scan_beta->GetYaxis()->SetTitleOffset(1.0);
    scan_beta->Draw("axis");
    
    box.DrawBox(best_ns->getVal()+best_ns->getErrorLo(),0.,
                best_ns->getVal()+best_ns->getErrorHi(),scan_beta->GetMaximum());
    
    for(auto prob: {0.,0.6827,0.9545})
        lin.DrawLine(0.,prob,20.,prob);
    
    scan_beta->Draw("csame");
    scan_fc->SetMarkerStyle(20);
    scan_fc->SetMarkerColor(50);
    scan_fc->SetLineWidth(2);
    scan_fc->SetLineColor(50);
    scan_fc->Draw("esame");

    TF1 fl("fl", "pol3", 0.,best_ns->getVal());
    TF1 fr("fr", "pol3",best_ns->getVal(),20.);
    fl.SetLineStyle(2);
    fr.SetLineStyle(2);
    scan_fc->Fit("fl","+R");
    scan_fc->Fit("fr","+R");
    
    printf("MINOS 68.3%% interval: [%.2f, %.2f]\n",
           best_ns->getVal()+best_ns->getErrorLo(),
           best_ns->getVal()+best_ns->getErrorHi());
    printf("F&C   68.3%% interval: [%.2f, %.2f]\n",fl.GetX(0.6827),fr.GetX(0.6827));
    
}