#include "TH1.h"
#include "TGraph.h"
#include "TMultiGraph.h"
#include "TLimit.h"
#include "TLimitDataSource.h"
#include "TConfidenceLevel.h"
#include "TCanvas.h"

#include <iostream>

void Limit(void) {
  TH1F *signal = new TH1F("signal","signal",4, 0., 4.);
  TH1F *bkg = new TH1F("bkg","bkg",4, 0., 4.);
  TH1F *data = new TH1F("data","data",4, 0., 4.);
  
  //Bin1 muTau
  data->SetBinContent(1,5.0);
  data->SetBinError(1,sqrt(5.0));
  
  //Bin2 eleTau
  data->SetBinContent(2,3.0);
  data->SetBinError(2,sqrt(3.0));
  
  //Bin3 SR1
  data->SetBinContent(3,1.0);
  data->SetBinError(3,sqrt(1.0));
  
  //Bin4 SR2
  data->SetBinContent(4,2.0);
  data->SetBinError(4,sqrt(2.0));
  
  //Bin1 muTau bkg from PAS
  bkg->SetBinContent(1,8.59);
  bkg->SetBinError(1,(4.83));
  
  //Bin2 eleTau bkg from PAS
  bkg->SetBinContent(2,3.52);
  bkg->SetBinError(2,(3.39));
  
  //Bin3 SR1 bkg from PAS
  bkg->SetBinContent(3,1.58);
  bkg->SetBinError(3,(0.65));
  
  //Bin4 SR2 bkg from PAS
  bkg->SetBinContent(4,7.07);
  bkg->SetBinError(4,(2.25));
  
  float r_step = 0.01;
  int steps = 500;
  TH1F *cls = new TH1F("cls", "cls", steps, (0.5 * r_step), ((0.5 + steps) * r_step));
  TGraph *gr2 = new TGraph(steps);
  gr2->SetLineColor(4);
  TGraph *gr3 = new TGraph(steps);
  gr3->SetLineColor(5);
  TGraph *gr4 = new TGraph(steps);
  gr4->SetLineColor(3);
  
  for(int i = 1; i < (steps + 1); i++) {
    Double_t r = i * r_step;
    
    signal->SetBinContent(1, 0.767 * r);
    signal->SetBinContent(2, 0.266 * r);
    signal->SetBinContent(3, 0.48 * r);
    signal->SetBinContent(4, 0.42 * r);
    
    signal->SetBinError(1, 0.2 * signal->GetBinContent(1));
    signal->SetBinError(2, 0.2 * signal->GetBinContent(2));
    signal->SetBinError(3, 0.25 * signal->GetBinContent(3));
    signal->SetBinError(4, 0.25 * signal->GetBinContent(4));
    
    TLimitDataSource *mydatasource = new TLimitDataSource(signal,bkg,data);
    TConfidenceLevel *myconfidence = TLimit::ComputeLimit(mydatasource,50000);
    
    cls->Fill(r, myconfidence->CLs());
    gr2->SetPoint((i - 1), r, myconfidence->GetExpectedCLs_b(0));
    gr3->SetPoint((i - 1), r, myconfidence->GetExpectedCLs_b(1));
    gr4->SetPoint((i - 1), r, myconfidence->GetExpectedCLs_b(-1));
    
    std::cout << "  r    : " << r << std::endl;
    std::cout << "  CLs  : " << myconfidence->CLs() << std::endl;
    std::cout << "  Expected   CLs_b  : " << myconfidence->GetExpectedCLs_b(0) << std::endl;
    std::cout << "  Expected+1 CLs_b  : " << myconfidence->GetExpectedCLs_b(1) << std::endl;
    std::cout << "  Expected-1 CLs_b  : " << myconfidence->GetExpectedCLs_b(-1) << std::endl;
    
    delete mydatasource;
    delete myconfidence;
  }
  
  Int_t last = cls->FindLastBinAbove(0.05);
  std::cout << last << " "
            << cls->GetBinLowEdge(last) << " "
            << cls->GetBinWidth(last) << " "
            << cls->GetBinCenter(last) << " "
            << cls->GetBinContent(last) << " " << std::endl;
  
  TGraph *gr1 = new TGraph(cls);
  gr1->SetLineColor(2);
  
  TMultiGraph *mg = new TMultiGraph();
  mg->Add(gr1);
  mg->Add(gr2);
  mg->Add(gr3);
  mg->Add(gr4);
  
  TCanvas *c1 = new TCanvas("c1","observe+expected cls",200,10,700,500);
  c1->Divide(1, 2);
  
  c1->cd(1); gPad->SetGrid();
  cls->Draw();
  
  c1->cd(2); gPad->SetGrid();
  mg->Draw("AC");
  
  c1->cd(0);
  c1->Modified(); c1->Update();
}