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 CL_scan(double *CLs, double *CLs_err, double target_ns = 10., int ntoys = 100, bool display = false)
{
    RooWorkspace *wspace = new RooWorkspace("wspace");
    buildModel(wspace);
    double Q_thres[4] = {0.,0.,0.,0.};

    wspace->var("ns")->setMax(target_ns);
    wspace->var("ns")->setVal(target_ns);
    wspace->var("ns")->setConstant(false);
    RooFitResult *res0 = wspace->pdf("model")->fitTo(*wspace->data("data"),ExternalConstraints(*wspace->pdf("cons")),Save(true));
    wspace->var("ns")->setVal(target_ns);
    wspace->var("ns")->setConstant(true);
    RooFitResult *res1 = wspace->pdf("model")->fitTo(*wspace->data("data"),ExternalConstraints(*wspace->pdf("cons")),Save(true));
    Q_thres[0] = max(0.,res1->minNll()-res0->minNll())*2.; // observed test stat
    
    vector<double> stat[2];
    for (int hypo=0; hypo<=1; hypo++) {
        for (int idx=0; idx<ntoys; idx++) {
            RooWorkspace* ws_toy = new RooWorkspace(*wspace);
            if (hypo==0) ws_toy->var("ns")->setVal(0.);
            if (hypo==1) ws_toy->var("ns")->setVal(target_ns);
            
            // main data
            RooDataSet *toy = ws_toy->pdf("model")->generate(*ws_toy->var("mass"));
            
            // 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")->setVal(target_ns);
            ws_toy->var("ns")->setConstant(false);
            RooFitResult *res0 = ws_toy->pdf("model")->fitTo(*toy,ExternalConstraints(*ws_toy->pdf("cons")),Save(true));
            ws_toy->var("ns")->setVal(target_ns);
            ws_toy->var("ns")->setConstant(true);
            RooFitResult *res1 = ws_toy->pdf("model")->fitTo(*toy,ExternalConstraints(*ws_toy->pdf("cons")),Save(true));
            double Q = max(0.,res1->minNll()-res0->minNll())*2.; // sampled test stat
            stat[hypo].push_back(Q);
            
            delete res0;
            delete res1;
            delete toy;
            delete ws_toy;
        }
    }
    
    sort(stat[0].begin(),stat[0].end());
    Q_thres[1] = stat[0][(int)(0.159*ntoys)]; // null expected 16% (-1s)
    Q_thres[2] = stat[0][(int)(0.500*ntoys)]; // null expected 50% (median)
    Q_thres[3] = stat[0][(int)(0.841*ntoys)]; // null expected 84% (+1s)
    
    for (int type=0; type<4; type++) {
        double CL[2] = {0.,0.}, CL_err[2] = {0.,0.};
        for (int hypo=0; hypo<=1; hypo++) {
            for (int idx=0; idx<ntoys; idx++)
                if (stat[hypo][idx]>=Q_thres[type]) CL[hypo] += 1.;
            CL[hypo] /= (double)ntoys;
            CL_err[hypo] = sqrt(CL[hypo]*(1.-CL[hypo])/ntoys);
        }
        CLs[type] = CL[1]/CL[0];
        if (CLs[type]>0.) CLs_err[type] = sqrt(pow(CL_err[0]/CL[0],2)+pow(CL_err[1]/CL[1],2))*CLs[type];
    }
    
    if (display) {
        TH1D* h_stat0 = new TH1D("h_stat0","",30,0.,30.);
        TH1D* h_stat1 = new TH1D("h_stat1","",30,0.,30.);
        for (int idx=0; idx<ntoys; idx++) {
            h_stat0->Fill(stat[0][idx]);
            h_stat1->Fill(stat[1][idx]);
        }
        
        TCanvas *c1 = new TCanvas("c1","",600,600);
        c1->SetMargin(0.14,0.06,0.13,0.07);
        c1->SetLogy();
        h_stat1->SetStats(false);
        h_stat1->SetFillColor(kBlue-10);
        h_stat1->GetXaxis()->SetTitle("-2log(L(#theta,#hat{#lambda}_{#theta})/L(#hat{#theta},#hat{#lambda}))");
        h_stat1->GetYaxis()->SetTitle("# of toys");
        h_stat1->GetXaxis()->SetTitleSize(0.06);
        h_stat1->GetYaxis()->SetTitleSize(0.06);
        h_stat1->GetXaxis()->SetTitleOffset(0.8);
        h_stat1->GetYaxis()->SetTitleOffset(1.0);
        h_stat1->Draw();
        h_stat0->SetFillColor(kRed);
        h_stat0->SetFillStyle(3365);
        h_stat0->Draw("same");
        TLine l1;
        l1.SetLineWidth(3);
        l1.DrawLine(Q_thres[0],0.,Q_thres[0],h_stat1->GetMaximum());
    }
    
    printf("Target ns = %.1f, observed CLs = %.3f +- %.3f\n",target_ns,CLs[0],CLs_err[0]);
    printf("16%% null expected: %.3f +- %.3f\n",CLs[1],CLs_err[1]);
    printf("50%% null expected: %.3f +- %.3f\n",CLs[2],CLs_err[2]);
    printf("84%% null expected: %.3f +- %.3f\n",CLs[3],CLs_err[3]);
    
    delete res0;
    delete res1;
    delete wspace;
}

void example_04()
{
    double CLs[4], CLs_err[4];
    CL_scan(CLs,CLs_err,18.,1000,true);
}