EWP-BplusToKstMuMu-AngAna/Code/Selection/Scripts/MakeNiceTMplots.cpp

//Create plots that compare the distributions for differnet truth-matching methods
//Renata Kopecna

#include "../GlobalFunctions.hh"
#include "../Design.hpp"
#include "../Paths.hpp"
#include <sys/stat.h> //mkdir
#include <boost/algorithm/string.hpp>

using namespace std;

TChain *GetTMfiles(int year, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){
    TChain * tree = new TChain("DecayTreeTruthMatched");
    if (B0 || K1 || Inc){ //possible TODO: year check
        tree->Add(GetInputFileBkg(year,"down",true,Reference,B0, K1, Inc).c_str());
        coutDebug("Getting the trees " + GetInputFileBkg(year,"down",true,Reference,B0, K1, Inc));
        tree->Add(GetInputFileBkg(year,"up",true,Reference,B0, K1, Inc).c_str());
        coutDebug("Getting the trees " + GetInputFileBkg(year,"up",true,Reference,B0, K1, Inc));
    }
    else{
        if (!checkYear(year,true,Reference,PHSP)) return NULL;
        coutDebug("Getting the trees " + GetInputFile(year,"down",true,Reference,false, false, false));
        tree->Add(GetInputFile(year,"down",true, true, Reference, PHSP, false).c_str());
        coutDebug("Getting the trees " + GetInputFile(year,"up",true,Reference,false, false, false));
        tree->Add(GetInputFile(year,"up",true, true, Reference, PHSP, false).c_str());
    }
    coutDebug("Tree obtained.");
    return tree;
}

TChain *GetTMfilesRun(int Run, bool Reference, bool PHSP){
    TChain * tree = new TChain("DecayTreeTruthMatched");
    for (auto yr: yearsMC(Reference,PHSP,Run)){
        coutDebug("Getting the trees " + GetInputFile(yr,"down",true,Reference,false, false, false));
        tree->Add(GetInputFile(yr,"down",true, true, Reference, PHSP, false).c_str());
        coutDebug("Getting the trees " + GetInputFile(yr,"up",true,Reference,false, false, false));
        tree->Add(GetInputFile(yr,"up",true, true, Reference, PHSP, false).c_str());
    }
    coutDebug("Tree obtained.");
    return tree;
}

//--------------------------------------------------------------------------------------------
// Compare what gamma types we have
//--------------------------------------------------------------------------------------------

int plotGammaTMcomparison(string firstBranch, string secondBranch,string year, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){

    gSystem->Load("Design_cpp.so");
    gROOT->SetBatch(kTRUE);

    coutInfo("Getting the TChain.");
    TChain *tree = GetTMfiles(std::stoi(year),Reference,PHSP,B0,K1,Inc);

    TH2D *h2_compare = new TH2D("h2_compare","h2_compare",6,0.5,6.5,6,0.5,6.5);
    string drawCommand = firstBranch + ":" + secondBranch + ">>" + h2_compare->GetName();

    coutDebug("Drawing this: " + drawCommand);
    tree->Draw(drawCommand.c_str());

    //design the histogram
    design_TH2F(h2_compare,secondBranch,firstBranch,"");
    h2_compare->GetXaxis()->SetTitleSize(0.03);
    h2_compare->GetYaxis()->SetTitleSize(0.03);
    h2_compare->GetXaxis()->SetTitleOffset(1.6);
    h2_compare->GetYaxis()->SetTitleOffset(3.0);

    h2_compare->GetXaxis()->SetBinLabel(1,"Both OK");
    h2_compare->GetXaxis()->SetBinLabel(2,"OK+Conv");
    h2_compare->GetXaxis()->SetBinLabel(3,"Both Conv");
    h2_compare->GetXaxis()->SetBinLabel(4,"OK+Rnd");
    h2_compare->GetXaxis()->SetBinLabel(5,"Conv+Rnd");
    h2_compare->GetXaxis()->SetBinLabel(6,"Both Rnd");

    h2_compare->GetYaxis()->SetBinLabel(1,"Both OK");
    h2_compare->GetYaxis()->SetBinLabel(2,"OK+Conv");
    h2_compare->GetYaxis()->SetBinLabel(3,"Both Conv");
    h2_compare->GetYaxis()->SetBinLabel(4,"OK+Rnd");
    h2_compare->GetYaxis()->SetBinLabel(5,"Conv+Rnd");
    h2_compare->GetYaxis()->SetBinLabel(6,"Both Rnd");

    //Save it in a canvas
    coutDebug("Saving the canvas.");
    TCanvas *canvas = new TCanvas("","",800, 750);
    canvas->cd();
    canvas->SetTopMargin(0.03);
    canvas->SetLeftMargin(0.18);
    canvas->SetRightMargin(0.03);
    canvas->SetBottomMargin(0.12);
    h2_compare->Draw("COLTEXT");

    string path = getTMplotsPath("compare_" + firstBranch + "_" + secondBranch, year, Reference, PHSP, B0, K1, Inc);
    canvas->SaveAs(path.c_str());
    replace(path,".eps",".root");
    canvas->SaveAs(path.c_str());

    tree->Clear();
    h2_compare->Clear();
    canvas->Clear();
    delete canvas;

    return 1;

}

int plotGammaTMcomparisonAllBranches(string year, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){

    string firstBranch = "TM_gammas";
    vector<string> secondBranch = {"TM_gammas_noPi0", "TM_gammas_noPi0_noBplus", "TM_gammas_noPi0_noKst_noBplus"};
    for (auto secBr : secondBranch){
        if (plotGammaTMcomparison(firstBranch,secBr,year,Reference,PHSP,B0,K1,Inc)==0) return 0;
    }
    return 1;
}

int plotGammaTMcomparisonAllYears(bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){
    vector<string> years = (B0||K1||Inc) ? yearsBkgMC(Reference,B0,K1,Inc,12) : yearsMC(Reference,PHSP,12);
    for (auto yr: years){
        if (plotGammaTMcomparisonAllBranches(yr,Reference,PHSP,B0,K1,Inc)==0) return 0;
    }
    return 1;
}
int plotGammaTMcomparisonAll(){
    if (plotGammaTMcomparisonAllYears( 0, 0, 0, 0, 0)==0) return 0;
    if (plotGammaTMcomparisonAllYears( 1, 0, 0, 0, 0)==0) return 0;
    if (plotGammaTMcomparisonAllYears( 0, 1, 0, 0, 0)==0) return 0;
    if (plotGammaTMcomparisonAllYears( 1, 0, 0, 0, 1)==0) return 0;
    return 1;
}

//--------------------------------------------------------------------------------------------

// Compare different gamma contributions in the B mass (TODO: pi mass?)
//--------------------------------------------------------------------------------------------
int plotAllGammaContributions(string branch, string year, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false, bool log = true){

    gSystem->Load("Design_cpp.so");
    gROOT->SetBatch(kTRUE);

    coutInfo("Getting the TChain.");
    TChain *tree = GetTMfiles(std::stoi(year),Reference,PHSP,B0,K1,Inc);

    TH1D *AllTMed        = BmassHistogram("AllTMed",        "", "B^{+} mass [MeV]", kBlack);
    TH1D *TMedWithGamma1 = BmassHistogram("TMedWithGamma1", "", "B^{+} mass [MeV]", kRed);
    TH1D *TMedWithGamma2 = BmassHistogram("TMedWithGamma2", "", "B^{+} mass [MeV]", kGreen);
    TH1D *TMedWithGamma3 = BmassHistogram("TMedWithGamma3", "", "B^{+} mass [MeV]", kBlue);
    TH1D *TMedWithGamma4 = BmassHistogram("TMedWithGamma4", "", "B^{+} mass [MeV]", kMagenta);
    TH1D *TMedWithGamma5 = BmassHistogram("TMedWithGamma5", "", "B^{+} mass [MeV]", kAzure);
    TH1D *TMedWithGamma6 = BmassHistogram("TMedWithGamma6", "", "B^{+} mass [MeV]", kOrange);

    //Fill the histograms, just using DTF cause screw it
    string cond = "TMed==1&&"+branch;
    tree->Draw("B_plus_M_DTF>>AllTMed","TMed==1");
    tree->Draw("B_plus_M_DTF>>TMedWithGamma1",string(cond+"==1").c_str());
    tree->Draw("B_plus_M_DTF>>TMedWithGamma2",string(cond+"==2").c_str());
    tree->Draw("B_plus_M_DTF>>TMedWithGamma3",string(cond+"==3").c_str());
    tree->Draw("B_plus_M_DTF>>TMedWithGamma4",string(cond+"==4").c_str());
    tree->Draw("B_plus_M_DTF>>TMedWithGamma5",string(cond+"==5").c_str());
    tree->Draw("B_plus_M_DTF>>TMedWithGamma6",string(cond+"==6").c_str());

    //Set the legend
    TLegend *leg = new TLegend(0.63,0.51,0.92,0.95);
    leg->AddEntry(AllTMed,       "TMed all",        "l");
    leg->AddEntry(TMedWithGamma1,"TMed+both OK",    "l");
    leg->AddEntry(TMedWithGamma2,"TMed+OK+Conv",    "l");
    leg->AddEntry(TMedWithGamma3,"TMed+both Conv",  "l");
    leg->AddEntry(TMedWithGamma4,"TMed+OK+Rnd",     "l");
    leg->AddEntry(TMedWithGamma5,"TMed+Conv+Rnd",   "l");
    leg->AddEntry(TMedWithGamma6,"TMed+both Rnd",   "l");

    //make the plot
    coutDebug("Saving the canvas.");
    TCanvas *canvas = c_canvas("");
    canvas->cd();
    if (log) canvas->SetLogy();
    AllTMed->Draw();
    TMedWithGamma1->Draw("SAME");
    TMedWithGamma2->Draw("SAME");
    TMedWithGamma3->Draw("SAME");
    TMedWithGamma4->Draw("SAME");
    TMedWithGamma5->Draw("SAME");
    TMedWithGamma6->Draw("SAME");
    leg->Draw("SAME");

    string path = getTMplotsPath("TMed_AllGammaContributions"+ string(log ? "_log" : ""), year, Reference, PHSP, B0, K1, Inc);
    canvas->SaveAs(path.c_str());
    replace(path,".eps",".root");
    canvas->SaveAs(path.c_str());

    tree->Clear();
    AllTMed->Clear();
    TMedWithGamma1->Clear();
    TMedWithGamma2->Clear();
    TMedWithGamma3->Clear();
    TMedWithGamma4->Clear();
    TMedWithGamma5->Clear();
    TMedWithGamma6->Clear();
    canvas->Clear();
    delete canvas;

    coutInfo("Done plotting.");

    return 1;
}

int plotAllGammaContributionsAllYears(bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){
    vector<string> years = (B0||K1||Inc) ? yearsBkgMC(Reference,B0,K1,Inc,12) : yearsMC(Reference,PHSP,12);
    for (auto yr: years){
        if (plotAllGammaContributions("TM_gammas",yr,Reference,PHSP,B0,K1,Inc, true)==0) return 0;
        if (plotAllGammaContributions("TM_gammas",yr,Reference,PHSP,B0,K1,Inc, false)==0) return 0;
    }
    return 1;
}
int plotAllGammaContributionsAll(){
    if (plotAllGammaContributionsAllYears( 0, 0, 0, 0, 0)==0) return 0;
    if (plotAllGammaContributionsAllYears( 1, 0, 0, 0, 0)==0) return 0;
    if (plotAllGammaContributionsAllYears( 0, 1, 0, 0, 0)==0) return 0;
    if (plotAllGammaContributionsAllYears( 1, 0, 0, 0, 1)==0) return 0;
    return 1;
}

//--------------------------------------------------------------------------------------------
// Compare different gamma contributions in costhetak //--------------------------------------------------------------------------------------------
int plotCosAngle(string branch, string year, int Run, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false, string angle = "thetak", bool cos = true, bool norm = false, bool fancy = true){
    gSystem->Load("Design_cpp.so");
    gROOT->SetBatch(kTRUE);
    TH1::SetDefaultSumw2(kTRUE);
    coutInfo("Getting the TChain.");
    TChain *tree ;
    if (Run ==0) tree = GetTMfiles(std::stoi(year),Reference,PHSP,B0,K1,Inc);
    else tree = GetTMfilesRun(Run,Reference,PHSP);

    TH1D *AllTMed        = hist_cosAngle(angle, "AllTMed",        "", kBlack,   cos, norm);
    TH1D *TMedWithGamma1 = hist_cosAngle(angle, "TMedWithGamma1", "", kRed,     cos, norm);
    TH1D *TMedWithGamma2 = hist_cosAngle(angle, "TMedWithGamma2", "", kOrange,   cos, norm);
    TH1D *TMedWithGamma3 = hist_cosAngle(angle, "TMedWithGamma3", "", kBlue,    cos, norm);
    TH1D *TMedWithGamma4 = hist_cosAngle(angle, "TMedWithGamma4", "", kMagenta, cos, norm);
    TH1D *TMedWithGamma5 = hist_cosAngle(angle, "TMedWithGamma5", "", kAzure,   cos, norm);
    TH1D *TMedWithGamma6 = hist_cosAngle(angle, "TMedWithGamma6", "", kGreen+2,  cos, norm);
    TH1D *TMedTrue       = hist_cosAngle(angle, "TMedTrue",       "", kRed,  cos, norm);
    TH1D *TMedTrueRnd    = hist_cosAngle(angle, "TMedRnd",        "", kBlue,  cos, norm);
    TH1D *Ratio          = hist_cosAngle(angle, "Ratio",        "", kBlack,  cos, norm);

    //Fill the histograms, just using DTF cause screw it
    boost::to_lower(angle);
    replace(angle,"_","");
    string var = "";
    if      (angle == "thetak") var = "B_plus_ThetaK";
    else if (angle == "thetal") var = "B_plus_ThetaL";
    else if (angle == "phi")    var = "B_plus_Phi";
    else{
        coutERROR("Wrong angle selected! Select from thetak, thetal and phi.");
        return 0;
    }
    string command = cos ? "TMatch::cos("+var+")" : var;
    string cond = "TMed==1&&"+branch;
    tree->Draw(string(command+">>AllTMed").c_str(),"TMed==1");
    tree->Draw(string(command+">>TMedWithGamma1").c_str(),string(cond+"==1").c_str());
    tree->Draw(string(command+">>TMedWithGamma2").c_str(),string(cond+"==2").c_str());
    tree->Draw(string(command+">>TMedWithGamma3").c_str(),string(cond+"==3").c_str());
    tree->Draw(string(command+">>TMedWithGamma4").c_str(),string(cond+"==4").c_str());
    tree->Draw(string(command+">>TMedWithGamma5").c_str(),string(cond+"==5").c_str());
    tree->Draw(string(command+">>TMedWithGamma6").c_str(),string(cond+"==6").c_str());

    //Add together both OK, one from conversion, both from conversion
    if (fancy){

        TMedTrue->Add(TMedWithGamma1);
        TMedTrue->Add(TMedWithGamma2);
        TMedTrue->Add(TMedWithGamma3);

        //Add together one random + one OK/ from conversion
        TMedTrueRnd->Add(TMedWithGamma4);
        TMedTrueRnd->Add(TMedWithGamma5);
    }


    //Set the legend
    TLegend *leg = new TLegend(0.68,0.75,0.92,0.95);
    leg->AddEntry(AllTMed,       "TMed all",        "l");
    leg->AddEntry(TMedWithGamma1,"TMed+both OK",    "l");
    leg->AddEntry(TMedWithGamma2,"TMed+OK+Conv",    "l");
    if (!norm)leg->AddEntry(TMedWithGamma3,"TMed+both Conv",  "l");
    leg->AddEntry(TMedWithGamma4,"TMed+OK+Rnd",     "l");
    if (!norm)leg->AddEntry(TMedWithGamma5,"TMed+Conv+Rnd",   "l");
    if (!norm)leg->AddEntry(TMedWithGamma6,"TMed+both Rnd",   "l");

    TLegend *legFancy;
    if (fancy){
        if (norm){
            if (cos){
                legFancy = new TLegend(0.25,0.65,0.61,0.9);
                legFancy->SetTextSize(0.05);
            }
            else{
                legFancy = new TLegend(0.55,0.65,0.93,0.9);
                legFancy->SetTextSize(0.05);
            }
        }
        else{
            legFancy = new TLegend(0.5,0.65,0.9,0.9);
            legFancy->SetTextSize(0.04);
        }
        legFancy->AddEntry(AllTMed,         "All #gamma#gamma","l");
        legFancy->AddEntry(TMedTrue,        "True #gamma#gamma","l");
        legFancy->AddEntry(TMedTrueRnd,     "True #gamma and random #gamma","l");
        if (!norm) legFancy->AddEntry(TMedWithGamma6,  "Random #gamma#gamma","l");
    }

    //Normalize
    double allEntries = AllTMed->GetEntries();
    if (norm){
        AllTMed->Scale(1.0/allEntries);
        TMedWithGamma1->Scale(1.0/TMedWithGamma1->GetEntries());
        TMedWithGamma2->Scale(1.0/TMedWithGamma2->GetEntries());
        TMedWithGamma3->Scale(1.0/TMedWithGamma3->GetEntries());
        TMedWithGamma4->Scale(1.0/TMedWithGamma4->GetEntries());
        TMedWithGamma5->Scale(1.0/TMedWithGamma5->GetEntries());
        TMedWithGamma6->Scale(1.0/TMedWithGamma6->GetEntries());
        if(fancy){
            TMedTrue->Scale(1.0/TMedTrue->GetEntries());
            TMedTrueRnd->Scale(1.0/TMedTrueRnd->GetEntries());
            Ratio->Divide(TMedTrue,TMedTrueRnd);
        }
        AllTMed->GetYaxis()->SetTitle("Counts/nEntries [a.u.]");
    }

    //make the plot
    coutDebug("Saving the canvas.");
    TCanvas *canvas = c_canvas("");
    canvas->cd();
    if (!norm) canvas->SetLogy();
    if (!norm) AllTMed->GetYaxis()->SetRangeUser(0.1,AllTMed->GetMaximum()*30);
    else AllTMed->GetYaxis()->SetRangeUser(0.0,AllTMed->GetMaximum()*2);
    coutDebug("Histogram maximum: "+to_string(AllTMed->GetMaximum()));
    AllTMed->Draw();
    TMedWithGamma1->Draw("SAME");
    TMedWithGamma2->Draw("SAME");
    if (!norm)TMedWithGamma3->Draw("SAME");
    TMedWithGamma4->Draw("SAME");
    if (!norm)TMedWithGamma5->Draw("SAME");
    if (!norm)TMedWithGamma6->Draw("SAME");
    leg->Draw("SAME");
    coutDebug("Canvas saved!");

    TCanvas *canvasFancy = c_canvas("fancy");
    if (fancy){ //TODO: legend for black line, add ratio of the normalized histograms
        coutDebug("Plotting fancy canvas");
        canvasFancy->cd();
        coutDebug("Histogram maximum: "+to_string(AllTMed->GetMaximum()));
        if (norm){
            //canvasFancy->SetBottomMargin(0.1);
            float lowerPlotHeight = 0.3;
            TPad *pad1 = new TPad("pad1", "plot",0.0,lowerPlotHeight,1.0,1.0,0);
            TPad *pad2 = new TPad("pad2", "ratio",0.0,0.0,1.0,lowerPlotHeight,0);
            pad1->Draw();
            pad2->Draw();

            //modifiy pad for the ratio
            pad2->Clear();
            pad2->SetBorderSize  (0);
            pad2->SetTopMargin   (1e-6);
            pad2->SetBottomMargin(0.3);
            pad2->cd();
            Ratio->SetMarkerSize(1.0);
            Ratio->GetYaxis()->SetRangeUser(0.5 , 1.5);
            Ratio->GetYaxis()->SetNdivisions(505);
            Ratio->GetYaxis()->SetTitleOffset(0.5);
            Ratio->GetYaxis()->SetTitleSize(0.12);
            Ratio->GetYaxis()->SetLabelSize(0.1);
            Ratio->GetYaxis()->SetTitle("#gamma_{T}#gamma_{T}/#gamma_{T}#gamma_{R}");

            Ratio->GetXaxis()->SetTitleOffset(0.9);
            Ratio->GetXaxis()->SetTitleSize(0.15);
            Ratio->GetXaxis()->SetLabelSize(0.14);

            Ratio->Draw("E1");
            TLine * unityline = new TLine( cos ? -1.0 : 0.0, 1,  cos ? 1.0 : 3.15, 1);
            unityline->SetLineStyle(2);
            unityline->SetLineColor(4);
            unityline->Draw("SAME");


            //modifiy pad for main plot
            pad1->Clear();
            pad1->SetBorderSize  (0);
            pad1->SetBottomMargin(1e-6);
            pad1->SetTopMargin(pad1->GetTopMargin() / ( 1.0 - lowerPlotHeight) );
            pad1->cd();

            TMedTrue->GetYaxis()->SetRangeUser(0.0,TMedTrue->GetMaximum()*1.5);

            TMedTrue->GetYaxis()->SetTitleOffset(1.20);
            TMedTrue->GetYaxis()->SetLabelSize(0.05);

            TMedTrue->Draw();
            AllTMed->Draw("SAME");
            TMedTrueRnd->Draw("SAME");
            legFancy->Draw("SAME");

        }
        else{
            canvasFancy->SetLogy();
            AllTMed->GetYaxis()->SetRangeUser(0.1,AllTMed->GetMaximum()*30);
            AllTMed->Draw();
            TMedTrue->Draw("SAME");
            TMedTrueRnd->Draw("SAME");
            TMedWithGamma6->Draw("SAME");
            legFancy->Draw("SAME");
        }

    }

    string path = getTMplotsPath("TMed_"+angle+"_AllGammaContributions"+string(norm ? "_normalized" : ""), year, Reference, PHSP, B0, K1, Inc);
    if (Run != 0) replace(path,year,"Run_"+to_string(Run));
    if (cos) replace(path,angle,"cos"+angle);
    coutDebug("Saving plot into " + path);
    canvas->SaveAs(path.c_str());
    replace(path,".eps",".root");
    canvas->SaveAs(path.c_str());

    if (fancy){
        replace(path,".root","_fancy.eps");
        coutDebug("Saving fancy plot into " + path);
        canvasFancy->SaveAs(path.c_str());
        replace(path,".eps",".root");
        canvasFancy->SaveAs(path.c_str());
    }

    tree->Clear();
    AllTMed->Clear();
    TMedWithGamma1->Clear();
    TMedWithGamma2->Clear();
    TMedWithGamma3->Clear();
    TMedWithGamma4->Clear();
    TMedWithGamma5->Clear();
    TMedWithGamma6->Clear();
    TMedTrue->Clear();
    TMedTrueRnd->Clear();
    canvas->Clear();
    delete canvas;

    coutInfo("Done plotting.");

    return 1;
}


int plotAllCosAngleAllAngles(string year, int Run, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false, bool cos = false, bool norm = false){
    if (plotCosAngle("TM_gammas",year,Run,Reference,PHSP,B0,K1,Inc, "thetak", cos, norm)==0) return 0;
    if (plotCosAngle("TM_gammas",year,Run,Reference,PHSP,B0,K1,Inc, "thetal", cos, norm)==0) return 0;
    if (plotCosAngle("TM_gammas",year,Run,Reference,PHSP,B0,K1,Inc, "phi", cos, norm)==0) return 0;
    return 1;

}

int plotAllCosAngleAllYears(bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){
    vector<string> years = (B0||K1||Inc) ? yearsBkgMC(Reference,B0,K1,Inc,12) : yearsMC(Reference,PHSP,12);
    for (auto yr: years){
        if (plotAllCosAngleAllAngles(yr,0,Reference,PHSP,B0,K1,Inc, true,  true)==0) return 0;
        if (plotAllCosAngleAllAngles(yr,0,Reference,PHSP,B0,K1,Inc, false, true)==0) return 0;
        if (plotAllCosAngleAllAngles(yr,0,Reference,PHSP,B0,K1,Inc, true,  false)==0) return 0;
        if (plotAllCosAngleAllAngles(yr,0,Reference,PHSP,B0,K1,Inc, false, false)==0) return 0;
    }
    if (plotAllCosAngleAllAngles("2011",1,Reference,PHSP,false,false,false, false, false)==0) return 0;
    if (plotAllCosAngleAllAngles("2011",12,Reference,PHSP,false,false,false, false, false)==0) return 0;
    if (plotAllCosAngleAllAngles("2016",2,Reference,PHSP,false,false,false, false, false)==0) return 0;
    return 1;
}
int plotAllCosAngleAllRuns(bool Reference, bool PHSP){
    if (plotAllCosAngleAllAngles("2011",1,Reference,PHSP,false,false,false, false, false)==0) return 0;
    if (plotAllCosAngleAllAngles("2011",1,Reference,PHSP,false,false,false, false, true)==0) return 0;
    //if (plotAllCosAngleAllAngles("2011",12,Reference,PHSP,false,false,false, false, false)==0) return 0;
    if (plotAllCosAngleAllAngles("2016",2,Reference,PHSP,false,false,false, false, false)==0) return 0;
    if (plotAllCosAngleAllAngles("2016",2,Reference,PHSP,false,false,false, false, true)==0) return 0;
    return 1;
}

int plotAllCosAngleAll(){
    if (plotAllCosAngleAllYears( 0, 0, 0, 0, 0)==0) return 0;
    if (plotAllCosAngleAllYears( 1, 0, 0, 0, 0)==0) return 0;
    if (plotAllCosAngleAllYears( 0, 1, 0, 0, 0)==0) return 0;
    if (plotAllCosAngleAllYears( 1, 0, 0, 0, 1)==0) return 0;
    return 1;
}

//--------------------------------------------------------------------------------------------
// Compare different gamma contributions for different TM methods //--------------------------------------------------------------------------------------------

int plotGammaTMdiffTM(string branch, string year, bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){

    //Load the files
    coutInfo("Getting the TChain.");
    TChain *tree = GetTMfiles(std::stoi(year),Reference,PHSP,B0,K1,Inc);

    TH1D *h_TMed               = generalHistogram("h_TMed",              "",6,0.5,6.5, branch.c_str(), "Counts a.u.", kRed);
    h_TMed->SetFillColor(kRed);
    TH1D *h_BKGCAT             = generalHistogram("h_BKGCAT",            "",6,0.5,6.5, branch.c_str(), "Counts a.u.", kBlack);
    h_BKGCAT->SetFillColor(kBlack);
    TH1D *h_TMed_noPi0         = generalHistogram("h_TMed_noPi0",        "",6,0.5,6.5, branch.c_str(), "Counts a.u.", kBlue);
    h_TMed_noPi0->SetFillColor(kBlue);
    TH1D *h_TMed_noPi0_noBplus = generalHistogram("h_TMed_noPi0_noBplus","",6,0.5,6.5, branch.c_str(), "Counts a.u.", kGreen);
    h_TMed_noPi0_noBplus->SetFillColor(kGreen);

    h_TMed->GetXaxis()->SetBinLabel(1,"Both OK");
    h_TMed->GetXaxis()->SetBinLabel(2,"OK+Conv");
    h_TMed->GetXaxis()->SetBinLabel(3,"Both Conv");
    h_TMed->GetXaxis()->SetBinLabel(4,"OK+Rnd");
    h_TMed->GetXaxis()->SetBinLabel(5,"Conv+Rnd");
    h_TMed->GetXaxis()->SetBinLabel(6,"Both Rnd");

    tree->Draw(string(branch+">>h_TMed").c_str(), "TMed==1");
    tree->Draw(string(branch+">>h_BKGCAT").c_str(), "TMedBKGCAT==1");
    tree->Draw(string(branch+">>h_TMed_noPi0").c_str(), "TMed_noPi0==1");
    tree->Draw(string(branch+">>h_TMed_noPi0_noBplus").c_str(), "TMed_noPi0_noBplus==1");

    //Set the legend
    TLegend *leg = new TLegend(0.68,0.61,0.92,0.95);
    leg->AddEntry(h_TMed,       "TRUEID",            "f");
    leg->AddEntry(h_BKGCAT,     "BKGCAT",            "f");
    leg->AddEntry(h_TMed_noPi0, "TRUEID wo #pi^{0}", "f");
    //leg->AddEntry(h_TMed_noPi0_noBplus, "TRUEID wo #pi^{0}, B^{+}", "l");

    //make the plot
    coutDebug("Saving the canvas.");

    TCanvas *canvas = c_canvas("");
    canvas->cd();

    //Create a TH stack
    THStack *hs = new THStack("hs","");
    hs->Add(h_TMed);
    hs->Add(h_BKGCAT);
    hs->Add(h_TMed_noPi0);
    //hs->Add(h_TMed_noPi0_noBplus);
    hs->Draw("nostackb"); //Draw it first so you can change the yaxis title....
    hs->GetYaxis()->SetTitle("Counts a.u.");
    hs->GetYaxis()->SetTitleOffset(1.6);
    hs->GetXaxis()->SetLabelSize(0.05);

    hs->Draw("nostackb");

    leg->Draw("SAME");

    string path = getTMplotsPath("GammaTM_compareTMmethods", year, Reference, PHSP, B0, K1, Inc);
    canvas->SaveAs(path.c_str());
    replace(path,".eps",".root");
    canvas->SaveAs(path.c_str());

    //Free the memory
    tree->Clear();
    h_TMed->Clear();
    h_BKGCAT->Clear();
    h_TMed_noPi0->Clear();
    h_TMed_noPi0_noBplus->Clear();
    canvas->Clear();
    delete canvas;

    coutInfo("Done plotting.");

    return 1;
}

int plotGammaTMdiffTMAllYears(bool Reference, bool PHSP, bool B0 = false, bool K1 = false, bool Inc = false){
    vector<string> years = (B0||K1||Inc) ? yearsBkgMC(Reference,B0,K1,Inc,12) : yearsMC(Reference,PHSP,12);
    for (auto yr: years){
        if (plotGammaTMdiffTM("TM_gammas",yr,Reference,PHSP,B0,K1,Inc)==0) return 0;
        if (plotGammaTMdiffTM("TM_gammas",yr,Reference,PHSP,B0,K1,Inc)==0) return 0;
    }
    return 1;
}

int plotGammaTMdiffTMAll(){
    if (plotGammaTMdiffTMAllYears( 0, 0, 0, 0, 0)==0) return 0;
    if (plotGammaTMdiffTMAllYears( 1, 0, 0, 0, 0)==0) return 0;
    if (plotGammaTMdiffTMAllYears( 0, 1, 0, 0, 0)==0) return 0;
    if (plotGammaTMdiffTMAllYears( 1, 0, 0, 0, 1)==0) return 0;
    return 1;
}

//--------------------------------------------------------------------------------------------
// Test functions //--------------------------------------------------------------------------------------------

bool testFunction(){
    setVerboseLevel(1);
    plotGammaTMdiffTM("TM_gammas","2011",0,0,0,0,0);
    return true;
}

void test(){
    if (testFunction()) cout << cGREEN << "All good!" << cRESET <<  endl;
    else coutERROR("Test failed!");
}