TTree* eicreconTree = NULL;

TFile* outFile = NULL;
TTree* outTree = NULL;

float mc_p;
float mc_eta;
float mc_phi;

int   rec_track;
float track_p;
float track_eta;
float track_phi;

float h_endcap_clust_E;
float h_endcap_clust_eta;
float h_endcap_clust_phi;

float e_endcap_clust_E;
float e_endcap_clust_eta;
float e_endcap_clust_phi;

float barrel_clust_E;
float barrel_clust_eta;
float barrel_clust_phi;

void SetInputBranchAddresses();
void CreateOutputTree(TString outFileName);
void ResetVariables();

using namespace std;

void getSingleElectronRecon(TString inFileName, TString outFileName) {

	TFile* eicreconFile = new TFile(inFileName);
	eicreconTree = (TTree*)eicreconFile->Get("events");

	int Nevents = eicreconTree->GetEntries();
	int ev(0);

	//Create Array Reader
    TTreeReader tr(eicreconTree);

	TTreeReaderArray<float> px_mc(tr,"GeneratedParticles.momentum.x");
	TTreeReaderArray<float> py_mc(tr,"GeneratedParticles.momentum.y");
	TTreeReaderArray<float> pz_mc(tr,"GeneratedParticles.momentum.z");
	TTreeReaderArray<float> E_mc(tr,"GeneratedParticles.energy");
	TTreeReaderArray<float> M_mc(tr,"GeneratedParticles.mass");
	TTreeReaderArray<int>   PDG_mc(tr,"GeneratedParticles.PDG");
	TTreeReaderArray<int>   type_mc(tr,"GeneratedParticles.type");

	TTreeReaderArray<float> px_track(tr,"ReconstructedChargedParticles.momentum.x");
	TTreeReaderArray<float> py_track(tr,"ReconstructedChargedParticles.momentum.y");
	TTreeReaderArray<float> pz_track(tr,"ReconstructedChargedParticles.momentum.z");
	TTreeReaderArray<float> M_track(tr,"ReconstructedChargedParticles.mass");
	TTreeReaderArray<int> PDG_track(tr,"ReconstructedChargedParticles.PDG");

	TTreeReaderArray<float> h_endcap_E(tr,"EcalEndcapPMergedClusters.energy");
	TTreeReaderArray<float> h_endcap_theta(tr,"EcalEndcapPMergedClusters.intrinsicTheta");
	TTreeReaderArray<float> h_endcap_phi(tr,"EcalEndcapPMergedClusters.intrinsicPhi");
	
	TTreeReaderArray<float> e_endcap_E(tr,"EcalEndcapNMergedClusters.energy");
	TTreeReaderArray<float> e_endcap_theta(tr,"EcalEndcapNMergedClusters.intrinsicTheta");
	TTreeReaderArray<float> e_endcap_phi(tr,"EcalEndcapNMergedClusters.intrinsicPhi");

	TTreeReaderArray<float> barrel_E(tr,"EcalBarrelSciGlassClusters.energy");
	TTreeReaderArray<float> barrel_theta(tr,"EcalBarrelSciGlassClusters.intrinsicTheta");
	TTreeReaderArray<float> barrel_phi(tr,"EcalBarrelSciGlassClusters.intrinsicPhi");

	CreateOutputTree(outFileName);
	
	while (tr.Next()) {

		if(ev%10000 == 0) cout << ev/1000 << "k / " << Nevents/1000 << "k events processed" << endl;
		ResetVariables();

		TVector3 electron_mc(0., 0., 0.);
		for(int imc = 0; imc < PDG_mc.GetSize(); imc++) {
			if(PDG_mc[imc] == 11 && type_mc[imc] == 1) {
				electron_mc.SetXYZ(px_mc[imc], py_mc[imc], pz_mc[imc]);
				break;
			}
		}

		mc_p = electron_mc.Mag();
		mc_eta = electron_mc.Eta();
		mc_phi = electron_mc.Phi();

		//Assume a single reconstructed track
		rec_track = 0;
		TVector3 electron_track(0., 0., 0.);
		if(PDG_track.GetSize()==1) {
			electron_track.SetXYZ(px_track[0], py_track[0], pz_track[0]);
			rec_track = 1;
		}

		track_p = electron_track.Mag();
		track_eta = electron_track.Eta();
		track_phi = electron_track.Phi();

		if(h_endcap_E.GetSize()>0)
			h_endcap_clust_E = h_endcap_E[0];
		if(e_endcap_E.GetSize()>0)
			e_endcap_clust_E = e_endcap_E[0];

		for(int icl = 0; icl < barrel_E.GetSize(); icl++) {
			barrel_clust_E += barrel_E[icl];
		}

		outTree->Fill();
		ev++;
	}

	outFile->cd();
	outTree->Write(outTree->GetName(), 2);

}

void CreateOutputTree(TString outFileName) {

	outFile = new TFile(outFileName, "RECREATE");
	outTree = new TTree("T", "Reconstructed and generated information from EICRecon");
	
	outTree->Branch("mc_p",			&mc_p);
	outTree->Branch("mc_eta",		&mc_eta);
	outTree->Branch("mc_phi",		&mc_phi);

	outTree->Branch("rec_track",		&rec_track);
	outTree->Branch("track_p",		&track_p);	
	outTree->Branch("track_eta",		&track_eta);	
	outTree->Branch("track_phi",		&track_phi);

	outTree->Branch("h_endcap_clust_E",	&h_endcap_clust_E);	
	outTree->Branch("h_endcap_clust_eta",	&h_endcap_clust_eta);
	outTree->Branch("h_endcap_clust_phi",	&h_endcap_clust_phi);

	outTree->Branch("e_endcap_clust_E",	&e_endcap_clust_E);	
	outTree->Branch("e_endcap_clust_eta",	&e_endcap_clust_eta);
	outTree->Branch("e_endcap_clust_phi",	&e_endcap_clust_phi);

	outTree->Branch("barrel_clust_E",	&barrel_clust_E);	
	outTree->Branch("barrel_clust_eta",	&barrel_clust_eta);
	outTree->Branch("barrel_clust_phi",	&barrel_clust_phi);

}

void ResetVariables() {

	// output variables
	mc_p = 0.;
	mc_eta = 0.;
	mc_phi = 0.;
			
	rec_track = 0.;      
	track_p = 0.;
	track_eta = 0.;
	track_phi = 0.;
			
	h_endcap_clust_E = 0.;
	e_endcap_clust_E = 0.;
	barrel_clust_E = 0.;
}