//Inlusive sample check
//Renata Kopecna

#include "../GlobalFunctions.hh"
#include "../Paths.hpp"
#include "../Utils.hpp"
#include "../Design.hpp"

int plotTM(int Run, bool Kstar){
    //Load MC data from TFiles
    TChain * tree   = new TChain("DecayTreeTruthMatched");
    if (Run == 1){
        tree->Add("/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/Inclusive_MC_BDToutput_Run1_2011.root");
        tree->Add("/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/Inclusive_MC_BDToutput_Run1_2012.root");
    }
    else if (Run == 2) tree->Add("/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/Inclusive_MC_BDToutput_Run2_2016.root");

    TH1D *h_all;
    TH1D *h_notTM;
    if (Kstar){
        h_all   = new TH1D ("h_all",  "h_all",  100, PDGMASS.K_STAR_PLUS-K_star_plus_mass_difference, PDGMASS.K_STAR_PLUS+K_star_plus_mass_difference);
        h_notTM = new TH1D ("h_notTM","h_notTM",100, PDGMASS.K_STAR_PLUS-K_star_plus_mass_difference, PDGMASS.K_STAR_PLUS+K_star_plus_mass_difference);
    }
    else{
        h_all   = new TH1D ("h_all",  "h_all",  100, cut_B_plus_M_low, cut_B_plus_M_high);
        h_notTM = new TH1D ("h_notTM","h_notTM",100, cut_B_plus_M_low, cut_B_plus_M_high);
    }

   double q_bins[7] = {0.1e6, 4.0e6, 8.0e6, 11.0e6, 12.5e6, 15.0e6, 20.0e6};
   TH1D *q_all   = new TH1D ("q_all",  "q_all",  6, q_bins);
   TH1D *q_notTM = new TH1D ("q_notTM","q_notTM",6, q_bins);

   string MLPcut = "(MLPresponse > " + to_string(getTMVAcut(Run)) + " && IsAloneAtNotTM <= " + to_string(getTMVAcut(Run)) + ")";
   string TMcut = "(TMed && TM_gammas<6)";
   string Psi2sCut = "(abs(B_plus_TRUEID)==521 &&" ;
   Psi2sCut.append("abs(J_psi_TRUEID)==100443 && abs(J_psi_MC_MOTHER_ID)==521 && ");
   Psi2sCut.append("abs(K_star_plus_TRUEID)==323 && abs(K_star_plus_MC_MOTHER_ID)==521 && ");
   Psi2sCut.append("abs(K_plus_TRUEID)==321 && abs(K_plus_MC_MOTHER_ID)==323 && abs(K_plus_MC_GD_MOTHER_ID)==521 && ");
   Psi2sCut.append("(abs(pi_zero_resolved_TRUEID)==111 ||  abs(pi_zero_resolved_TRUEID)==22))");

   string BstarCut = "(abs(B_plus_TRUEID)==521 && ";
   BstarCut.append("abs(K_star_plus_TRUEID)==521 && abs(K_star_plus_MC_MOTHER_ID)==523 &&");
   BstarCut.append("abs(K_plus_TRUEID)==321 && abs(K_plus_MC_MOTHER_ID)==521 && ");
   BstarCut.append("(abs(pi_zero_resolved_TRUEID)==111 ||  abs(pi_zero_resolved_TRUEID)==22))");


   string BCut = "(B_plus_TRUEID==521 && ";
   BCut.append("K_star_plus_TRUEID==323 && K_star_plus_MC_MOTHER_ID==521 && ");
   BCut.append("K_plus_TRUEID==321 && K_plus_MC_MOTHER_ID==323 && K_plus_MC_GD_MOTHER_ID==521 && ");
   BCut.append("(pi_zero_resolved_TRUEID==111 ||  pi_zero_resolved_TRUEID==22))");


   string fullCut = MLPcut + " && !(" + TMcut + " || " + Psi2sCut + " || " + BstarCut + ")";
   cout << fullCut << endl;

   if(Kstar){
       tree->Draw("K_star_plus_M_scaled>>h_all",MLPcut.c_str());
       tree->Draw("K_star_plus_M_scaled>>h_notTM",fullCut.c_str());
   }
   else{
       tree->Draw("B_plus_M_DTF>>h_all",MLPcut.c_str());
       tree->Draw("B_plus_M_DTF>>h_notTM",fullCut.c_str());
   }
   tree->Draw("Q2>>q_all",MLPcut.c_str());
   tree->Draw("Q2>>q_notTM",fullCut.c_str());

   TCanvas *c = c_canvas("Inclusive sample check");
   c->cd();
   design_TH1D(h_all,recoMassKplus,"Entries (a.u.)", kBlack);
   design_TH1D(h_notTM,recoMassKplus,"Entries (a.u.)", kRed);
   h_all->Draw();
   h_notTM->Draw("SAME");
   if (Run == 1) addAnyTag(0.6,0.85,"MC 2011+2012",kBlack,0.8);
   else          addAnyTag(0.6,0.85,"MC 2016",kBlack,0.8);
   addAnyTag(0.6,0.79,"B#rightarrow XJ/#psi",kBlack,0.8);
   string name = "Inclusive_MC_Run_" +to_string(Run) + std::string(Kstar? "_Kstar" : "");;
   c->SaveAs(string(name + ".root").c_str());
   c->SaveAs(string(name + ".eps").c_str());


   TCanvas *cq = c_canvas("Inclusive sample check");
   cq->cd();
   design_TH1D(q_all,"q^{2} [MeV]","Entries (a.u.)", kBlack);
   design_TH1D(q_notTM,"q^{2} [MeV]","Entries (a.u.)", kRed);
   addAnyTag(0.6,0.85,"Simulation",kBlack,0.9);
   addAnyTag(0.6,0.79,"Run " + to_string(Run),kBlack,0.9);
   q_all->Draw();
   q_notTM->Draw("SAME");
   name = "Inclusive_Q2_MC_Run_" +to_string(Run) + std::string(Kstar? "_Kstar" : "");
   cq->SaveAs(string(name + ".root").c_str());
   cq->SaveAs(string(name + ".eps").c_str());


   return 1;
}