#include <iostream>
#include <string>

#include "/usr/local/aliroot/AliRoot_v4-17-Rev-09/include/AliESDtrack.h"
#include "/usr/local/aliroot/AliRoot_v4-17-Rev-09/include/AliESDEvent.h"
#include "/usr/local/aliroot/AliRoot_v4-17-Rev-09/include/AliESDfriendTrack.h"

using namespace std;

//Function protoypes
void showMenu();
void drawOverlay(TH1D *h2, TH1D *h1, const char *type, const char *h1_name, const char *h2_name);

//This function shows the available plots and save the plotted histogram to disk
void showMenu() {

cout << "Please choose which plots you want to access." << endl;
cout << "(0)Quit, (1)Momentum, (2)Transverse momentum, (3)Theta angle," << endl;
cout << "(4)Phi angle, (5)TPC clusters per track, (6)Charged track multiplicity per event:";

}

void drawOverlay(TH1D *h1, TH1D *h2, const char *type, const char *h1_name, const char *h2_name) {
TCanvas *c1 = new TCanvas();
//Create the legend and later draw it on the same pad
L = new TLegend(0.55, 0.85, 0.75, 0.99);
//Start by drawing the histograms on the same pad
h1->SetLineColor(4);
h1->Draw();
h2->SetLineColor(2);
h2->Draw("sames");
//Move the statistics box of h2 so that the do not overlap
c1->Update();
TPaveStats *s = (TPaveStats*)h2->GetListOfFunctions()->FindObject("stats");
s->SetY1NDC (0.65);
s->SetY2NDC (0.80);

L->AddEntry(h1_name, "Offline", "l");
L->AddEntry(h2_name, "HLT", "l");
L->Draw();
c1->Update();

//Save the canvas as .png file
c1->SaveAs(type + TString(".png"));
cout << TString("Histogram saved as ") + type +TString(".png") << endl;

}

Int_t task1() {

//Make offline histograms
TH1D *hp = new TH1D("hp", "Momentum", 100, 0, 6);
hp->SetMaximum(400000);
TH1D *hpt = new TH1D("hpt", "Transverse", 100, 0, 6);
hpt->SetMaximum(450000);
TH1D *hth = new TH1D("hth", "Theta angle", 100, 0, 3);
TH1D *hph = new TH1D("hph", "Phi angle", 100, 0, 8);
hph->SetMaximum(40000);
hph->SetMinimum(10000);
TH1D *htpc = new TH1D("htpc", "TPC clusters per track", 100, 0, 200);
TH1D *hctm = new TH1D("hctm", "CTM per event", 100, 0, 100);

//Make HLT histograms
TH1D *hp_HLT = new TH1D("hp_HLT", "Momentum", 100, 0, 6);
hp->SetMaximum(400000);
TH1D *hpt_HLT = new TH1D("hpt_HLT", "Transverse", 100, 0, 6);
hpt_HLT->SetMaximum(450000);
TH1D *hth_HLT = new TH1D("hth_HLT", "Theta angle", 100, 0, 4);
//hth_HLT->SetMaximum(8000);
TH1D *hph_HLT = new TH1D("hph_HLT", "Phi angle", 100, 0, 8);
TH1D *htpc_HLT = new TH1D("htpc_HLT", "TPC clusters per track", 100, 0, 200);
hph_HLT->SetMaximum(40000);
hph_HLT->SetMinimum(10000);
TH1D *hctm_HLT = new TH1D("hctm_HLT", "CTM per event", 100, 0, 100);



//The next step is to extract the necassary data and fill the histograms
//Start by creating two chains, one with the offline trees and one with the HLT trees
TChain *esdchain = new TChain("esdTree");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.10/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.20/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.20/AliESDs2.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.30/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.40/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.50/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.60/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.70/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.80/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.90/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.10/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.20/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.30/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.40/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.50/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.60/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.70/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.80/AliESDs.root");
esdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.90/AliESDs.root");

TChain *HLTesdchain = new TChain("HLTesdTree");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.10/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.20/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.20/AliESDs2.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.30/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.40/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.50/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.60/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.70/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.80/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322040.90/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.10/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.20/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.30/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.40/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.50/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.60/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.70/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.80/AliESDs.root");
HLTesdchain->Add("/Home/stud5/ema069/PHYS291/115322/10000115322041.90/AliESDs.root");

//Create pointers to variables that points to the events and tracks we want to extract data from
AliESDEvent *esdevent = new AliESDEvent();
AliESDtrack *esdtrack = new AliESDtrack();
AliESDEvent *HLTesdevent = new AliESDEvent();
AliESDtrack *esdtrack = new AliESDtrack();

esdevent->ReadFromTree(esdTree);
HLTesdevent->ReadFromTree(HLTesdTree);

Int_t nesdevents = esdchain->GetEntries();
Int_t nHLTesdevents = HLTesdchain->GetEntries();

for (Int_t i = 0; i < nesdevents; i++) {
	esdchain->GetEntry(i);
	Int_t nesdtracks = esdevent->GetNumberOfTracks();
	//Fill CTM(charged track multiplicity) per event histogram, the tracks stems from charged particles
	hctm->Fill(nesdtracks);

	for (Int_t j = 0; j < nesdtracks; j++) {
		esdtrack = esdevent->GetTrack(j);

		//Fill historgrams with track properties
		hp->Fill(esdtrack->P());
		hpt->Fill(esdtrack->Pt());
		hth->Fill(esdtrack->Theta());
		hph->Fill(esdtrack->Phi());
		htpc->Fill(esdtrack->GetTPCNcls());
		
	}
}

//This loop does the same as the one above, but this takes data from the HLTesd tree
for (Int_t i = 0; i < nHLTesdevents; i++) {
	HLTesdchain->GetEntry(i);
	Int_t nHLTesdtracks = HLTesdevent->GetNumberOfTracks();
	hctm_HLT->Fill(nHLTesdtracks);

	for (Int_t j = 0; j < nHLTesdtracks; j++) {
		HLTesdtrack = HLTesdevent->GetTrack(j);

		hp_HLT->Fill(HLTesdtrack->P());
		hpt_HLT->Fill(HLTesdtrack->Pt());
		hth_HLT->Fill(HLTesdtrack->Theta());
		hph_HLT->Fill(HLTesdtrack->Phi());
		htpc_HLT->Fill(HLTesdtrack->GetTPCNcls());
		
	}
}

//The remaining part of the code displays the plots available, lets the user choose which plot that is 
//to be displayed. This plot is saved on the disk as a .root file. If the choice is not valid, the user
//enters a loop and stays there until something valid is typed in
showMenu();

Int_t choice;
cin >> choice;
while(choice != 1 && choice != 2 && choice != 3 && choice != 4
	&& choice != 5 && choice != 6 && choice != 0) {
	cout << "This is not a valid choice. Please try again:";
	cin >> choice;
}

while(choice != 0) {
	if(choice == 1)
	drawOverlay(hp, hp_HLT, "P", "hp", "hp_HLT");
	if(choice == 2)
	drawOverlay(hpt, hpt_HLT, "Pt", "hpt", "htp_HLT");
	if(choice == 3)
	drawOverlay(hth, hth_HLT, "Theta", "hth", "hth_HLT");
	if(choice == 4)
	drawOverlay(hph, hph_HLT, "Phi", "hph", "hph_HLT");
	if(choice == 5)
	drawOverlay(htpc, htpc_HLT, "TPC", "htpc", "htpc_HLT");
	if(choice == 6)
	drawOverlay(hctm, hctm_HLT, "CTM", "hctm", "hctm_HLT");
	showMenu();
	cin >> choice;
	while(choice != 1 && choice != 2 && choice != 3 && choice && choice != 4
		&& choice != 5 && choice != 6 && choice != 0) {
	cout << "This is not a valid choice. Please try again:";
	cin >> choice;
	}
}

cout << endl << "So long!";


return 0;

}