//Create tables of event numbers for B+->Kst+mumu
//david gerick

#include "../../GlobalFunctions.hh"
#include "../../HeidelbergFitter/LHCbStyle.h"

//#include "MVA_b2kmm.cc"

using namespace std;
using namespace RooFit ;

//////////////////////////////////////////////////////
///  GetEventNumbers()
///    Load all tuples and print the event numbers in a table for latex
///    Selection steps:
///    - 0: After stripping
///    - 1: After pre-selection
///    - 2: After BDT cut

Int_t GetEventNumbers(bool MC = false, bool Ref = false, bool PHSP = false, Int_t SelectionStep = 0) {

    if(SelectionStep < 0 || SelectionStep > 2){
      std::cout << "[ERROR]\t\tOnly use Selection Steps between 0 - 2 for after Stripping (0), Pre-Selection (1) or BDT cut (2)." << std::endl;
      return 0;
    }

    gStyle -> SetOptStat(0);
    gROOT->SetBatch(kTRUE);
    LHCbStyle();
   
    std::vector<Int_t> years;
    years.push_back(2011);
    years.push_back(2012);
    years.push_back(2015);
    years.push_back(2016);
    years.push_back(2017);
    years.push_back(2018);

    std::vector<std::string> Mag;
    if(SelectionStep == 0){
      Mag.push_back("down");
      Mag.push_back("up");
    }

    std::vector<bool> isLL;
    if(Kst2Kspiplus){
      isLL.push_back(true);
      isLL.push_back(false);
    }

    std::vector<UInt_t> EvNum[2][2]; //[Mag][DDLL].at(year)
    TChain * t;

    for(UInt_t y = 0; y < years.size(); y++){
      std::cout << "[LOADING]\tYear " << years.at(y) << std::endl;
      for(UInt_t m = 0; m < (SelectionStep == 0 ? 2 : 1); m++){
        for(UInt_t d = 0; d < (Kst2Kspiplus ? 2 : 1); d++){
          //init new tree:
          if(SelectionStep == 0)
            t = new TChain(Kst2Kpluspi0Resolved ? "b2KstKpi0mumuResolvedTuple/DecayTree" : "b2KstKs0pimumu_Tuple/DecayTree");
          else if(SelectionStep == 1)
            t = new TChain("DecayTree");
          else if(SelectionStep == 2)
            t = new TChain("SelectionOutput");

          //load all tuples:
          if(SelectionStep == 0)
            t->Add(Form("%s/data/%i%s/%i*_B2Kstmumu*.root", thePath.c_str(), years.at(y), Mag.at(m).c_str(), years.at(y)));
          else if(SelectionStep == 1){
            if(Kst2Kspiplus)
              t->Add(Form("%s/data/%i_KshortPiplus_%s_BDTinput.root", thePath.c_str(), years.at(y), d == 0 ? "DD" : "LL"));
            else
              t->Add(Form("%s/data/%i*BDTinput.root", thePath.c_str(), years.at(y)));
          }
          else if(SelectionStep == 2){
            t->Add(Form("%s/data/*BDToutputSelection_Run*%i.root", thePath.c_str(), years.at(y)));
          }

          if(Kst2Kspiplus){
            if(SelectionStep == 1)
              EvNum[m][d].push_back(t->GetEntries());
            else{
              EvNum[m][0].push_back(t->GetEntries("Ks_pi_minus_TRACK_Type != 3"));
              EvNum[m][1].push_back(t->GetEntries("Ks_pi_minus_TRACK_Type == 3"));
              delete t;
              break; //do not run DDLL loop for a second time!
            }
          }
          else
            EvNum[m][0].push_back(t->GetEntries());

          delete t;
        }
      }
    }

    //print LaTeX table:
    std::cout << "\\begin{tabular}{lc";
    for(UInt_t m = 0; m < Mag.size(); m++)
      std::cout << "c";
    if(Kst2Kspiplus)
      std::cout << "c";
    std::cout << "} \\hline" << std::endl;
    std::cout << "Year\t";
    if(Kst2Kspiplus)
      std::cout << "&track type\t";
    if(Mag.size() == 0)
      std::cout << "&both magnet polarities\\\\" << std::endl;
    else
      std::cout << "&magnet down\t&magnet up\t&combined\\\\" << std::endl;
    std::cout << "\\hline\\hline" << std::endl;
    for(UInt_t y = 0; y < years.size(); y++){
      std::cout << years.at(y);
      for(UInt_t d = 0; d < (Kst2Kspiplus ? 2 : 1); d++){
        std::cout << " \t";
        if(Kst2Kspiplus)
          std::cout << (d == 0 ? "&\\texttt{DD}\t" : "&\\texttt{LL}\t");
        if(Mag.size() == 0)
          std::cout << "&" << EvNum[0][d].at(y) << "\\\\" << std::endl;
        else
          std::cout << "&" << EvNum[0][d].at(y) << "\t&" << EvNum[1][d].at(y) << "\t&" << EvNum[0][d].at(y)+EvNum[1][d].at(y) << "\\\\" << std::endl; 
      }
      if(Kst2Kspiplus){
        if(Mag.size() == 0)
          std::cout << "\t&total\t&" << EvNum[0][0].at(y)+EvNum[0][1].at(y) << "\\\\" << std::endl;
        else
          std::cout << "\t&total\t&" << EvNum[0][0].at(y)+EvNum[0][1].at(y) 
                    << "\t&" << EvNum[1][0].at(y)+EvNum[1][1].at(y) 
                    << "\t&" << EvNum[0][0].at(y)+EvNum[0][1].at(y)+EvNum[1][0].at(y)+EvNum[1][1].at(y) << "\\\\" << std::endl;
      }
      std::cout << "\\hline" << std::endl;
    }
    std::cout << "\\end{tabular}" << std::endl;

    return 1;

}