#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"
#include "/usr/local/aliroot/AliRoot_v4-17-Rev-09/include/AliPID.h"

//In this source code, we need to fill four histograms; the proton(p) + pi(-) and
//pi(-) + pi(+), both for all events and from different events. Throughout, these are
//numbered as follows:
//1. p + pi(-) from all events
//2. pi(+) + pi(-) from all events
//3. p + pi(-) from different events
//4. pi(+) + pi(-) from different events

//Function prototypes
void showMenu();
void drawHistogram(TH1D *h, const char *type);

//This function draws and saves the histogram, and is basically a modified version of drawOVerlay in final.C
void drawHistogram(TH1D *h, const char *type) {

TCanvas *c1 = new TCanvas();

h->SetLineColor(4);
h->Draw();

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

}

void showMenu() {

cout << "Which plot do you want to see and save?" << endl;
cout << "(0)Quit" << endl << "(1)Mass distribution" << endl <<"(2)Invariant mass of p + pi(-)" << endl;
cout << "(3)Invariant mass of pi(+) + pi(-)" << endl << "(4)Estimated background for p + pi(-)" << endl;
cout << "(5)Estimated background for pi(-) + pi(+)" << endl << "(6)Invariant mass of p + pi(-) with background subtracted" << endl;
cout << "(7)Invariant mass of pi(+) + pi(-) with background subtracted" << endl;

}

Int_t task2() {

//Using the class TLorentzVector makes the calculations easier, so this loads the library
//declares the necassary four vectors
gSystem->Load("libPhysics.so");
TLorentzVector v1;
TLorentzVector v2;

TH1D *hm = new TH1D("hm", "Mass", 500, 0, 1);
hm->GetXaxis()->SetTitle("Mass[GeV/c²]");
TH1D *h1 = new TH1D("h1", "Invariant mass of p + pi(-)", 500, 0.5, 3);
TH1D *h2 = new TH1D("h2", "Invariant mass of pi(-) + pi(+)", 500, 0, 2);
TH1D *hbg1 = new TH1D("hbg1", "Estimated background of p + pi(-)", 500, 0.5, 3);
TH1D *hbg2 = new TH1D("hbg2", "Estimated background of pi(+) + pi(-)", 500, 0, 2); 

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");

AliESDtrack *esdtrack = new AliESDtrack();

AliESDEvent *esdevent = new AliESDEvent();


esdevent->ReadFromTree(esdTree);

Int_t nesdevents = esdchain->GetEntries();

Double_t E1, E2, E3; 
Double_t px1, px2, px3; 
Double_t py1, py2, py3; 
Double_t pz1, pz2, pz3;

//The first loops finds and fills the invariant mass of the p + pi(-) systems
for(Int_t i = 0; i < nesdevents; i++) {
	esdchain->GetEntry(i);
	Int_t nesdtracks = esdevent->GetNumberOfTracks();

	for (Int_t j = 0; j < nesdtracks; j++) {
		esdtrack = esdevent->GetTrack(j);
		hm->Fill(esdtrack->GetMass());
		//Find the proton. If this is the track of the proton, set the varialbes equal to the properties
		//of the proton and enter the next loop which finds the pi(-) and fills the histogram
		if(esdtrack->GetMass() > 0.9 && esdtrack->GetMass() < 1 && esdtrack->Charge() > 0) {
			E1 = esdtrack->E();
			px1 = esdtrack->Px();
			py1 = esdtrack->Py();
			pz1 = esdtrack->Pz();
			for(Int_t k = 0; k < nesdtracks; k++) {
				esdtrack = esdevent->GetTrack(k);
				if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() < 0) {
					E2 = esdtrack->E();
					px2 = esdtrack->Px();
					py2 = esdtrack->Py();
					pz2 = esdtrack->Pz();
					v1.SetE(E1 + E2);
					v1.SetPx(px1 + px2);
					v1.SetPy(py1 + py2);
					v1.SetPz(pz1 + pz2);
					h1->Fill(v1.Mag());
					}
			}
		}
	}
}

//This loop finds and fills the invariant mass of the pi(+) + pi(-) systems
for(Int_t i = 0; i < nesdevents; i++) {
	esdchain->GetEntry(i);
	Int_t nesdtracks = esdevent->GetNumberOfTracks();

	for (Int_t j = 0; j < nesdtracks; j++) {
		esdtrack = esdevent->GetTrack(j);
		//Find the pi(+). If this is the track of the pi(+), set the variables equal to the properties
		//of the pi(+) and enter the next loop which finds the pi(-) and fills the histogram
		if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() > 0) {
				E3 = esdtrack->E();
				px3 = esdtrack->Px();
				py3 = esdtrack->Py();
				pz3 = esdtrack->Pz();
				for(Int_t k = 0; k < nesdtracks; k++) {
					esdtrack = esdevent->GetTrack(k);
					if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() < 0) {
						E2 = esdtrack->E();
						px2 = esdtrack->Px();
						py2 = esdtrack->Py();
						pz2 = esdtrack->Pz();
						v2.SetE(E2 + E3);
						v2.SetPx(px2 + px3);
						v2.SetPy(py2 + py3);
						v2.SetPz(pz2 + pz3);
						h2->Fill(v2.Mag());
						}
				}
		}
	}
}

//The next two loops estimate the background of p + pi(-) with event mixing
for(Int_t i = 9; i < 10; i++){
	esdchain->GetEntry(i);
	nesdtracks = esdevent->GetNumberOfTracks();
		
		for(Int_t j = 0; j < nesdtracks; j++) {
			esdtrack = esdevent->GetTrack(j);
			if(esdtrack->GetMass() > 0.9 && esdtrack->GetMass() < 1 && esdtrack->Charge() > 0) {
				E1 = esdtrack->E();
				px1 = esdtrack->Px();
				py1 = esdtrack->Py();
				pz1 = esdtrack->Pz();
				cout << "Proton found in event 9!" << endl;
			}
		}
}

for(Int_t i = 0; i < nesdevents; i++){
	esdchain->GetEntry(i);
	nesdtracks = esdevent->GetNumberOfTracks();
		for(Int_t j = 0; j < nesdtracks; j++) {
			esdtrack = esdevent->GetTrack(j);
			if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() < 0 && i != 9) {
					E2 = esdtrack->E();
					px2 = esdtrack->Px();
					py2 = esdtrack->Py();
					pz2 = esdtrack->Pz();
					v1.SetE(E1 + E2);
					v1.SetPx(px1 + px2);
					v1.SetPy(py1 + py2);
					v1.SetPz(pz1 + pz2);
					hbg1->Fill(v1.Mag());
					}
		}
}

//The next two loops do the same as the two previous ones, only for the pion system
for(Int_t i = 13; i < 14; i++){
	esdchain->GetEntry(i);
	nesdtracks = esdevent->GetNumberOfTracks();
		
		for(Int_t j = 0; j < nesdtracks; j++) {
			esdtrack = esdevent->GetTrack(j);
			if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() > 0) {
				E3 = esdtrack->E();
				px3 = esdtrack->Px();
				py3 = esdtrack->Py();
				pz3 = esdtrack->Pz();
				cout << "Pion found in event 13!" << endl;
			}
		}
}

for(Int_t i = 0; i < nesdevents; i++){
	esdchain->GetEntry(i);
	nesdtracks = esdevent->GetNumberOfTracks();
		for(Int_t j = 0; j < nesdtracks; j++) {
			esdtrack = esdevent->GetTrack(j);
			if(esdtrack->GetMass() > 0.12 && esdtrack->GetMass() < 0.15 && esdtrack->Charge() < 0 && i != 13) {
					E2 = esdtrack->E();
					px2 = esdtrack->Px();
					py2 = esdtrack->Py();
					pz2 = esdtrack->Pz();
					v2.SetE(E3 + E2);
					v2.SetPx(px3 + px2);
					v2.SetPy(py3 + py2);
					v2.SetPz(pz3 + pz2);
					hbg2->Fill(v2.Mag());
					}
		}
}

//Copy the histogram, normalize away from where the signal is expected and subtract the background
TH1D *h1N = (TH1D*)h1->Clone();
h1N->SetName("h1N");
h1N->SetTitle("p + pi(-), background subtracted");

TH1D *hbg1N = (TH1D*)hbg1->Clone(); 
hbg1N->SetName("hbg1N");

//Normalize the two spectra away from the peak
Double_t scale = 1/h1N->Integral(300,500);
h1N->Scale(scale);

Double_t scale2 = 1/hbg1N->Integral(300,500);
hbg1N->Scale(scale2);

h1N->Add(hbg1N, -1);

//Do the same for the pi(+) + pi(-) system
TH1D *h2N = (TH1D*)h2->Clone();
h2N->SetName("h2N");
h2N->SetTitle("pi(+) + pi(-),background subtracted");

TH1D *hbg2N = (TH1D*)hbg2->Clone(); 
hbg1N->SetName("hbg2N");

//Normalize the two spectra away from the peak
Double_t scale3 = 1/h2N->Integral(350,500);
h1N->Scale(scale3);

Double_t scale4 = 1/hbg2N->Integral(350,500);
hbg1N->Scale(scale4);

//Subtract the background from the copy
h2N->Add(hbg2N, -1);


//Show the available plots and let the user choose which one to see and save
showMenu();
Int_t choice;
cin >> choice;
while(choice != 0) {
	while(choice != 1 && choice != 2 && choice != 3 && choice != 4
		&& choice != 5 && choice != 6 && choice != 7 && choice != 0) {
		cout << "This is not a valid choice. Please try again:";
		cin >> choice;
		}
if(choice == 1)
drawHistogram(hm, "Mass distribution");
if(choice == 2)
drawHistogram(h1, "Plot1");
if(choice == 3)
drawHistogram(h2, "Plot2");
if(choice == 4)
drawHistogram(hbg1, "Estimated background for plot 1");
if(choice == 5)
drawHistogram(hbg2, "Estimated background for plot 2");
if(choice == 6)
drawHistogram(h1N, "p + pi(-),background subtracted");
if(choice == 7)
drawHistogram(h2N, "pi(+) + pi(-),background subtracted");

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

return 0;
}