#include <iostream>

using namespace std;

Double_t FindCut(TTree* Sig, Double_t ratio);
void FormatAxis(TAxis * axis);
void SetStyle();
void SetMeetingStyle();

void project() {
    

    // Open data file
    TFile* file = new TFile("3.root");

    // Read in trees from the file
    TTree* Sig  = (TTree*)file->Get("tree_mc_sig");
    TTree* Bkg  = (TTree*)file->Get("tree_mc_bkg");
    TTree* Data = (TTree*)file->Get("tree_data");

    // Initialize canvas for plotting and saving
    SetStyle();
    SetMeetingStyle();
    TCanvas* c = new TCanvas("can","can",2500,1800);
    gStyle->SetOptStat(kFALSE);
    gStyle->SetOptFit(0000);

    /*                                *
     *     Task 1: Event Selection    *
     *                                */

    cout<<endl<<"========================================================"<<endl;
    cout<<" Task 1: Event Selection"<<endl;
    // Find the 90% percent cut on ann
    Double_t cut_anna = FindCut( Sig, 0.9 );
    //Double_t cut_anna = 0.8;
    // Apply ann cut on both Sig and Bkg
    char temp[50];
    sprintf(temp,"ann>=%f",cut_anna);
    TCut cut1 = temp;
    Long64_t cut_Sig = Sig->Draw("ann",cut1);
    Long64_t cut_Bkg = Bkg->Draw("ann",cut1);
    Double_t Significance = sqrt( 2*((cut_Sig+cut_Bkg)*log(1.0+(cut_Sig+0.0)/cut_Bkg)-cut_Sig) );

    // Optimization
    //c->SetGrid();
    const unsigned o = 50;
    Double_t ini = 0.;
    Double_t end = 1.;
    Long64_t cut_Sig_o,cut_Bkg_o;
    Double_t ann_o[o],Significance_o[o];
    for(int i=0; i<o; i++) {
        
        ann_o[i] = ini+(end-ini)/o*i;
        sprintf(temp,"ann>=%f", ann_o[i]);
        TCut cut2 = temp;
        cut_Sig_o = Sig->Draw("ann",cut2);
        cut_Bkg_o = Bkg->Draw("ann",cut2);
        Significance_o[i] = sqrt( 2*((cut_Sig_o+cut_Bkg_o)*log(1.0+(cut_Sig_o+0.0)/cut_Bkg_o)-cut_Sig_o) );
    }
    
    TGraph* gr = new TGraph(o,ann_o,Significance_o);
    gr->SetLineColor(kBlue);
    gr->SetLineWidth(4);
    gr->SetMarkerStyle(24);
    gr->SetMarkerColor(2);
    gr->GetXaxis()->SetTitle("ann");
    gr->GetYaxis()->SetTitle("Significance"); 
    FormatAxis(gr->GetXaxis());
    FormatAxis(gr->GetYaxis());
    gr->Draw("AC*");

    c->SaveAs("Significance.png");

    // Print Summary
    cout<<"==> By applying cut \" ann >= "<<cut_anna<<" \"..."<<endl;
    cout<<"==> Results: "<<endl;
    cout<<"             | ---- "<< setw(15)<< "Total"<<setw(15)<<"Cut"<<setw(15)<<"Sup."<<endl;
    cout<<"             | Sig: "<< setw(15)<< Sig->GetEntries() << setw(15) << (cut_Sig+0.0)/Sig->GetEntries() << setw(15) << (Sig->GetEntries()-cut_Sig+0.0)/Sig->GetEntries()<<endl;
    cout<<"             | Bkg: "<< setw(15)<< Bkg->GetEntries() << setw(15) << (cut_Bkg+0.0)/Bkg->GetEntries() << setw(15) << (Bkg->GetEntries()-cut_Bkg+0.0)/Bkg->GetEntries()<<endl;
    cout<<"             | ---- "<<endl;
    cout<<"             | Significance: "<<Significance<<endl;
    cout<<"             | ---- "<<endl;

    /*                                *
     *     Task 2: fit to data        *
     *                                */

    cout<<"========================================================"<<endl;
    cout<<" Task 2: Fit to Data"<<endl;
    cout<<"--------------------------------------------------------"<<endl;
    // Draw the mass distribution of data applying the cut
    TH1D* massh = new TH1D("massh","mass distribution",160,0,160);
    Data->Draw("mass>>massh",cut1);

    /*
     * There are several methods to fit; instead of guessing the initial value of parameters,
     * I prefer the codes which are not commented. The difference between these two fit methods is,
     * the first one you need to guess the initial value,
     * and the second you let ROOT decide the initial value.
     * Either way works.
     */

    /*
    // first type 
    TF1* fitfunc = new TF1("fitfunc","[0]*[5]*TMath::Gaus(x,[1],[2])  + (1-[0])*exp([3]+[4]*x)",0,160);
    fitfunc->SetParameters( 0.5, 80., 5., 8.0, -0.024,1000);
    fitfunc->SetParLimits(1,75,85);
    */

    // second type
    Double_t par[6];
    TF1* f1 = new TF1("f1","gaus",75,85);
    TF1* f2 = new TF1("f2","expo",0,160);
    TF1* fitfunc = new TF1("fitfunc","[5]*gaus(0)+(1-[5])*expo(3)",0,160);
    massh->Fit("f1","QR");
    massh->Fit("f2","QR+");
    f1->GetParameters(&par[0]);
    f2->GetParameters(&par[3]);
    par[5] = 0.5;
    fitfunc->SetParameters(par);

    // Fitting data!!
    massh->Fit("fitfunc","Q");

    // fetching fitting results
    Double_t mean = fitfunc->GetParameter(1);
    Double_t std  = fitfunc->GetParameter(2);
    
    Double_t mean_err = fitfunc->GetParError(1);
    Double_t std_err  = fitfunc->GetParError(2);

    cout<<" mean: "<<setprecision(5)<<mean<<setw(5)<<", err: "<<mean_err<<endl;
    cout<<" std:  "<<setprecision(5)<<std <<setw(5)<<", err: "<<std_err <<endl;
    cout<<"--------------------------------------------------------"<<endl;

    // calculate local/global significance
    TF1* f3 = new TF1("f3","(1-[0])*exp([1]+[2]*x)",0,160);
    f3->SetParameters(fitfunc->GetParameter(5),fitfunc->GetParameter(3),fitfunc->GetParameter(4));

    cout<<" Significance: local = "<< (fitfunc->Integral(75,85) - f3->Integral(75,85))/sqrt(f3->Integral(75,85));
    cout<< ", global = " << (fitfunc->Integral(0,150) - f3->Integral(0,150))/sqrt(f3->Integral(0,150)) << endl;

    
    // saving plot
    massh->SetLineColor(kBlue);
    massh->SetLineWidth(4);
    massh->SetMarkerStyle(24);
    massh->SetMarkerColor(2);
    massh->GetXaxis()->SetTitle("Mass [GeV]");
    massh->GetYaxis()->SetTitle("Yields [GeV/bin]"); 
    FormatAxis(massh->GetXaxis());
    FormatAxis(massh->GetYaxis());
    massh->Draw();

    fitfunc->SetLineColor(kRed);
    fitfunc->SetLineWidth(3);
    fitfunc->Draw("same");

    f3->Draw("same");

    TLatex l;
    l.SetTextSize(0.04);
    l.DrawLatex(57,2300,"#splitline{#mu = 80.357 #pm 0.109}{#sigma = 3.0388 #pm 0.111}");

    TArrow ar;
    ar.SetLineWidth(4);
    ar.SetFillColor(kBlack);
    ar.SetAngle(40);
    ar.DrawArrow(80,2000,80,1150,0.03,"|>");

    // Writing text
    TLatex * text1 = new TLatex(0.335,0.86, "iStep Homework");
    text1->SetNDC();
    text1->SetTextFont(62);
    text1->SetTextSize(0.06);
    //text1->SetTextAlign(33);
    //text1->Draw();

    // Plotting Legends
    TLegend* leg = new TLegend(0.8, 0.83, 0.95, 0.93);
    leg->AddEntry("massh","Data","l");
    leg->AddEntry("fitfunc","Fit Function","lfp");
    leg->Draw();

    c->SaveAs("mass.png");

    cout<<"========================================================"<<endl;
}


//...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...

Double_t FindCut(TTree* Sig, Double_t ratio) {
    /* This Function is used to find the cut point that keep cerntain ratio of the Signal. */

    cout<<"--------------------------------------------------------"<<endl;
    Double_t eps = 0.0001; // tolerance
    Long64_t sum = 0; // number of events
    Double_t cut = 0; // the cut point

    Long64_t n_entries = Sig->GetEntries();
    Sig->Draw("ann>>h1(20000,0,1)","");
    TH1D* h1 = (TH1D*)gPad->GetPrimitive("h1");
    Int_t nbinsx = h1->GetXaxis()->GetNbins();

    cout<<"==> Applying Binary Search..."<<endl;
    cout<<"==> Parameters: "<<endl;
    cout<<"                ratio:       "<<ratio<<endl;
    cout<<"                tolarence:   "<<eps<<endl;
    // binary search to find the certain point
    Int_t bin_max = nbinsx;
    Int_t bin_min = 0;
    Int_t bin_mid = 0.5*(bin_max+bin_min);
    while( bin_min < bin_max ) {
        sum = h1->Integral(bin_mid,nbinsx);
        if( abs(sum/(n_entries+0.0) - ratio) <= eps ) {
            cout<<"==> Cut Point founded!..."<<endl;
            break;
        }
        else if( (sum/(n_entries+0.0) < ratio - eps) ) {
            bin_max = bin_mid;
        }
        else if( (sum/(n_entries+0.0) > ratio + eps) ) {
            bin_min = bin_mid; 
        }
        bin_mid = ( (Int_t)(0.5*(bin_max+bin_min)) == bin_mid ) ? bin_mid+1 : (Int_t)(0.5*(bin_max+bin_min)) ;
    }
    cut = h1->GetBinCenter(bin_mid);
    cout<<"==> Result(ann>=): "<<cut<<", with "<<sum/(n_entries+0.0)*100<<"% signal kept..."<<endl;
    cout<<"--------------------------------------------------------"<<endl;
    return cut;
}


//...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...ooo0ooo...

void FormatAxis(TAxis * axis){
  //axis->SetLabelFont(42);
  //axis->SetLabelSize(0.06);
  //axis->SetLabelOffset(0.01);
  //axis->SetNdivisions(510);
  //axis->SetTitleFont(42);
  //axis->SetTitleColor(1);
  //axis->SetTitleSize(0.07);
  //axis->SetTitleOffset(1.15);
  axis->CenterTitle();
}

void SetStyle(){
  // No Canvas Border
  gStyle->SetCanvasBorderMode(0);
  gStyle->SetCanvasBorderSize(0);
  // White BG
  gStyle->SetCanvasColor(10);
  // Format for axes
  gStyle->SetLabelFont(42,"xyz");
  gStyle->SetLabelSize(0.05,"xyz");
  gStyle->SetLabelOffset(0.01,"xyz");
  gStyle->SetNdivisions(510,"xyz");
  gStyle->SetTitleFont(42,"xyz");
  gStyle->SetTitleColor(1,"xyz");
  gStyle->SetTitleSize(0.07,"xyz");
  gStyle->SetTitleOffset(1.15,"xyz");
  // No pad borders
  gStyle->SetPadBorderMode(0);
  gStyle->SetPadBorderSize(0);
  // White BG
  gStyle->SetPadColor(10);
  // Margins for labels etc.
  gStyle->SetPadLeftMargin(0.17);
  gStyle->SetPadBottomMargin(0.17);
  gStyle->SetPadRightMargin(0.05);
  gStyle->SetPadTopMargin(0.07);
  // No error bars in x direction
  gStyle->SetErrorX(0);
  // Format legend
  gStyle->SetLegendBorderSize(1);
}

void SetMeetingStyle(){
  gStyle->SetOptDate(0);
  gStyle->SetOptTitle(0);
  gStyle->SetOptStat(1111);
  gStyle->SetStatBorderSize(1);
  gStyle->SetStatColor(10);
  gStyle->SetStatFont(42);
  gStyle->SetStatFontSize(0.03);
  gStyle->SetOptFit(1111);
}