using namespace RooFit;
TRandom3 rnd;

void CL_scan(double *CLs, double target_ns = 5., int ntoys = 1000, bool display = false)
{
    const int nobs = 8;
    const double nb = 4.5;
    double Q_thres[4] = {0.,0.,0.,0.};
    
    double L0 = TMath::Poisson(nobs,nb);
    double L1 = TMath::Poisson(nobs,nb+target_ns);
    Q_thres[0] = -2.*(log(L1)-log(L0)); // observed test statistic
    
    vector<double> stat[2];
    for (int hypo=0; hypo<=1; hypo++) {
        for (int idx=0; idx<ntoys; idx++) {
            int n = 0;
            if (hypo==0) n = rnd.Poisson(nb); // toy sampling b-only
            if (hypo==1) n = rnd.Poisson(nb+target_ns); // toy sampling s+b
            
            double L0 = TMath::Poisson(n,nb);
            double L1 = TMath::Poisson(n,nb+target_ns);
            double Q = -2.*(log(L1)-log(L0)); // sampled test statistic
            stat[hypo].push_back(Q);
        }
    }
    
    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.};
        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;
        }
        CLs[type] = CL[1]/CL[0];
    }

    if (display) {
        TH1D* h_stat0 = new TH1D("h_stat0","",50,-50.,50.);
        TH1D* h_stat1 = new TH1D("h_stat1","",50,-50.,50.);
        
        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_{s+b}/L_{b})");
        h_stat1->GetYaxis()->SetTitle("# of toys");
        h_stat1->GetXaxis()->SetTitleSize(0.06);
        h_stat1->GetYaxis()->SetTitleSize(0.06);
        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[2],0.,Q_thres[2],h_stat1->GetMaximum());
        l1.SetLineStyle(2);
        l1.DrawLine(Q_thres[0],0.,Q_thres[0],h_stat1->GetMaximum());
    }
    
    printf("Target ns = %.1f, observed CLs = %.3f\n",target_ns,CLs[0]);
    printf("16%% null expected: %.3f\n",CLs[1]);
    printf("50%% null expected: %.3f\n",CLs[2]);
    printf("84%% null expected: %.3f\n",CLs[3]);
    
    return CLs;
}

void example_01b()
{
    vector<double> vec_x, vec_obs, vec_exp[3];
    double target_ns = 0.;
    while (target_ns<=14.) {
        double CLs[4];
        CL_scan(CLs,target_ns,40000);
        
        vec_x.push_back(target_ns);
        vec_obs.push_back(CLs[0]);
        vec_exp[0].push_back(CLs[1]);
        vec_exp[1].push_back(CLs[2]);
        vec_exp[2].push_back(CLs[3]);
        target_ns += 0.25;
    }
    
    TGraph *gr_obs = new TGraph(vec_x.size(),vec_x.data(),vec_obs.data());
    gr_obs->SetLineWidth(3);
    gr_obs->SetLineColor(kBlue);
    
    TGraph *gr_exp = new TGraph(vec_x.size(),vec_x.data(),vec_exp[1].data());
    gr_exp->SetLineColor(kBlack);
    gr_exp->SetLineWidth(3);
    gr_exp->SetLineStyle(2);
    
    TGraph *gr_experr = new TGraph(vec_x.size()*2);
    for (int i=0; i<vec_x.size(); i++) {
        gr_experr->SetPoint(i,vec_x[i],vec_exp[2][i]);
        gr_experr->SetPoint(vec_x.size()*2-1-i,vec_x[i],vec_exp[0][i]);
    }
    gr_experr->SetFillColor(kGreen);
    
    TCanvas *c1 = new TCanvas("c1","",600,600);
    c1->SetMargin(0.14,0.06,0.13,0.07);
    c1->SetGrid();
    c1->SetLogy();
    
    TH2D *frame = new TH2D("frame","",10,0.,14.,10,0.01,1.2);
    frame->SetStats(false);
    frame->GetXaxis()->SetTitle("Signal Yield");
    frame->GetYaxis()->SetTitle("CL_{s} = 1-CL");
    frame->GetXaxis()->SetTitleSize(0.06);
    frame->GetYaxis()->SetTitleSize(0.06);
    frame->GetYaxis()->SetTitleOffset(1.0);
    frame->Draw("");
    
    gr_experr->Draw("fsame");
    gr_exp->Draw("same");
    gr_obs->Draw("same");
    
    TLine lin;
    lin.SetLineColor(kGray+2);
    lin.SetLineStyle(7);
    lin.SetLineWidth(3);
    lin.DrawLine(0.,0.05,14.,0.05);
    lin.DrawLine(0.,0.10,14.,0.10);
}