// Copyright 2014 Andreas Herten
#include <vector>
#include <algorithm>
#include <iostream>

#include "TFile.h"
#include "TROOT.h"
#include "TCanvas.h"
#include "TH1.h"
#include "TH2.h"
#include "TTree.h"
#include "TBranch.h"
#include "TString.h"
#include "THStack.h"
#include "../../common.cpp"

void divideCanvas(TCanvas * canvas, TObjArray * array, int rows) {
	int length = array->GetEntries();
	int columns;
	(length%rows ? columns = length/rows + 1 : columns = length/rows);
	canvas->Divide(rows, columns);
}

void divideCanvasAndDraw(TCanvas * canvas, TObjArray * array, TString drawoptions = "", bool drawLegends = true, int rows = 2) {
	TString datatype = array->Last()->ClassName();
	if (datatype.EqualTo("TH1D")) {
		divideCanvas(canvas, array, rows);
		int length = array->GetEntries();
		for (int i = 0; i < length; ++i) {
			canvas->cd(i+1);
			((TH1D*)array->At(i))->Draw(drawoptions);
			if (drawLegends) {
				TLegend * leg = andi::plainLegend();
				leg->AddEntry(((TH1D*)(array->At(i))), ((TH1D*)(array->At(i)))->GetTitle(), "l");
				leg->Draw();
			}
		}
	}
}

void divideCanvasAndDrawAndFit(TCanvas * canvas, TObjArray * array, TString drawoptions = "", bool drawLegends = false, int rows = 2, Bool_t verbose = false) {
	TString datatype = array->Last()->ClassName();
	if (datatype.EqualTo("TH1D")) {
		divideCanvasAndDraw(canvas, array, drawoptions, drawLegends, rows);
		int length = array->GetEntries();
		for (int i = 0; i < length; ++i) {
			canvas->cd(i+1);
			TF1 * fit = andi::gaussFit((TH1D*)(array->At(i)), true);
			std::cout << fit->GetTitle() << std::endl;
			// fit->Draw("SAME");
		}
	}
}

void scaleHistogram(TH1 * hist, Float_t scalefactor) {
	hist->Scale(1/scalefactor);
	hist->GetYaxis()->SetTitle("Rel. Units");
}

void eval(TString filename = "D_K_Pi_Pi_fast.root_ana.root") {
	andi::setCustomStyle();
	TString macroBasename = "eval.C";

	TFile * file = new TFile(filename);
	TTree * tupleDplus = (TTree*)(file->Get("ntp1"));
	TTree * tupleDminus = (TTree*)(file->Get("ntp2"));
	TTree * trueTuples = (TTree*)(file->Get("nmc"));

	Int_t nOfEntriesDplus = tupleDplus->GetEntries();
	Int_t nOfEntriesDminus = tupleDminus->GetEntries();
	Int_t nOfGeneratedEvents = trueTuples->GetEntries();

	Int_t ncandDplus = 0;
	Int_t ncandDminus = 0;
	tupleDplus->SetBranchAddress("ncandDplusI", &ncandDplus);
	tupleDminus->SetBranchAddress("ncandDminusI", &ncandDminus);


	TH1D * nOfDpMesons = new TH1D("nOfDpMesons", "nOfDpMesons", 5, 0, 5);
	nOfDpMesons->SetLineColor(kBlue);
	TH1D * nOfDmMesons = new TH1D("nOfDmMesons", "nOfDmMesons", 5, 0, 5);
	nOfDmMesons->SetLineColor(kRed);

	andi::DInfoContainer dPlus;
	andi::setBranchAddresses(tupleDplus, dPlus, "Dplus");
	andi::DInfoContainer dPlusFit;
	andi::setBranchAddresses(tupleDplus, dPlusFit, "fDplus");
	andi::DInfoContainer trDplus;
	andi::setBranchAddresses(tupleDplus, trDplus, "trDplus", false);

	andi::DInfoContainer dMinus;
	andi::setBranchAddresses(tupleDminus, dMinus, "Dminus");


	// ###############################
	// D misc momentum properties (both D+ and D-)
	// ###############################

	TObjArray histsMomentumDmDp;
	histsMomentumDmDp.Add(new TH1D("Dmpx", "D^{-} p_{x}", 100, -0.5, 0.5));
	((TH1D*)(histsMomentumDmDp.Last()))->SetLineColor(kRed);
	tupleDminus->Project("Dmpx", "Dminuspx");
	histsMomentumDmDp.Add(new TH1D("Dmpy", "D^{-} p_{y}", 100, -0.5, 0.5));
	((TH1D*)(histsMomentumDmDp.Last()))->SetLineColor(kRed);
	tupleDminus->Project("Dmpy", "Dminuspy");
	histsMomentumDmDp.Add(new TH1D("Dppx", "D^{+} p_{x}", 100, -0.5, 0.5));
	((TH1D*)(histsMomentumDmDp.Last()))->SetLineColor(kBlue);
	tupleDplus->Project("Dppx", "Dpluspx");
	histsMomentumDmDp.Add(new TH1D("Dppy", "D^{p} p_{y}", 100, -0.5, 0.5));
	((TH1D*)(histsMomentumDmDp.Last()))->SetLineColor(kBlue);
	tupleDplus->Project("Dppy", "Dpluspy");

	TObjArray histsMomentum2DmDp;
	histsMomentum2DmDp.Add(new TH1D("Dmpt", "D^{-} p_{t}", 100, 0, 1.5));
	((TH1D*)(histsMomentum2DmDp.Last()))->SetLineColor(kRed);
	tupleDminus->Project("Dmpt", "Dminuspt");
	histsMomentum2DmDp.Add(new TH1D("Dmp", "D^{-} p", 100, 0, 6));
	((TH1D*)(histsMomentum2DmDp.Last()))->SetLineColor(kRed);
	tupleDminus->Project("Dmp", "Dminusp");
	histsMomentum2DmDp.Add(new TH1D("Dppt", "D^{+} p_{t}", 100, 0, 1.5));
	((TH1D*)(histsMomentum2DmDp.Last()))->SetLineColor(kBlue);
	tupleDplus->Project("Dppt", "Dpluspt");
	histsMomentum2DmDp.Add(new TH1D("Dpp", "D^{+} p", 100, 0, 6));
	((TH1D*)(histsMomentum2DmDp.Last()))->SetLineColor(kBlue);
	tupleDplus->Project("Dpp", "Dplusp");


	// ###############################
	// D misc other properties (both D+ and D-)
	// ###############################

	TH1D * histDmEnergy = new TH1D("histDmEnergy", "D^{-} Energy", 100, 0, 7);
	histDmEnergy->SetLineColor(kRed);
	tupleDminus->Project("histDmEnergy", "Dminuse");
	TH1D * histDmMass = new TH1D("histDmMass", "D^{-} Mass", 100, 0, 2.7);
	histDmMass->SetLineColor(kRed);
	tupleDminus->Project("histDmMass", "Dminusm");
	TH1D * histDmCharge = new TH1D("histDmCharge", "D^{-} Charge", 100, -1, 1.1);
	histDmCharge->SetLineColor(kRed);
	tupleDminus->Project("histDmCharge", "Dminuschg");	
	TH1D * histDmPdgId = new TH1D("histDmPdgId", "D^{-} PDG ID", 100, -500, 500);
	histDmPdgId->SetLineColor(kRed);
	tupleDminus->Project("histDmPdgId", "Dminuspdg");	
	TH1D * histDpEnergy = new TH1D("histDpEnergy", "D^{+} Energy", 100, 0, 7);
	histDpEnergy->SetLineColor(kBlue);
	tupleDplus->Project("histDpEnergy", "Dpluse");
	TH1D * histDpMass = new TH1D("histDpMass", "D^{+} Mass", 100, 0, 2.7);
	histDpMass->SetLineColor(kBlue);
	tupleDplus->Project("histDpMass", "Dplusm");
	TH1D * histDpCharge = new TH1D("histDpCharge", "D^{+} Charge", 100, -1, 1.1);
	histDpCharge->SetLineColor(kBlue);
	tupleDplus->Project("histDpCharge", "Dpluschg");
	TH1D * histDpPdgId = new TH1D("histDpPdgId", "D^{+} PDG ID", 100, -500, 500);
	histDpPdgId->SetLineColor(kBlue);
	tupleDplus->Project("histDpPdgId", "Dpluspdg");


	// ###############################
	// D+ differences to MC truth
	// ###############################

	TH1D * histDpResPt = new TH1D("histDpResPt", "D^{+} #Deltap_{t}", 100, -0.06, 0.06);
	histDpResPt->SetLineColor(kBlue);
	histDpResPt->GetXaxis()->SetTitle("(p_{t} - p_{t}^{MC}) / GeV/c");
	histDpResPt->GetXaxis()->SetTitleOffset(1.15);
	histDpResPt->GetXaxis()->SetTitleSize(0.035);
	histDpResPt->GetYaxis()->SetTitle("counts");
	tupleDplus->Project("histDpResPt", "Dpluspt - trDpluspt");

	TH1D * histDpResPtFit = new TH1D("histDpResPtFit", "D^{+} #Deltap_{t} (mass constraint fit)", 100, -0.06, 0.06);
	histDpResPtFit->SetLineColor(kBlue);
	histDpResPtFit->GetXaxis()->SetTitle("(p_{t} - p_{t}^{MC}) / GeV/c");
	histDpResPtFit->GetXaxis()->SetTitleOffset(1.15);
	histDpResPtFit->GetXaxis()->SetTitleSize(0.035);
	histDpResPtFit->GetYaxis()->SetTitle("counts");
	tupleDplus->Project("histDpResPtFit", "fDpluspt - trDpluspt");

	TH1D * histDpResM = new TH1D("histDpResM", "D^{+} #Deltam", 100, -0.1, 0.1);
	histDpResM->SetLineColor(kBlue);
	histDpResM->GetXaxis()->SetTitle("(m - m^{MC}) / GeV/c^2");
	histDpResM->GetXaxis()->SetTitleOffset(0.95);
	histDpResM->GetYaxis()->SetTitle("counts");
	tupleDplus->Project("histDpResM", "Dplusm - trDplusm");

	TH1D * histDpResP = new TH1D("histDpResP", "D^{+} #Deltap", 100, -0.2, 0.2);
	histDpResP->SetLineColor(kBlue);
	histDpResP->GetXaxis()->SetTitle("(p - p^{MC}) / GeV/c");
	histDpResP->GetXaxis()->SetTitleOffset(0.95);
	histDpResP->GetYaxis()->SetTitle("counts");
	tupleDplus->Project("histDpResP", "Dplusmcdiffp");
	// histDpResP->Fill((dPlus.m.p - trDplus.m.p));

	// ###############################
	// D+ positions
	// ###############################

	TH2F * histDpXvsY = new TH2F("histDpXvsY", "D^{+} x vs. y", 100, -0.02, 0.02, 100, -0.02, 0.02);
	histDpXvsY->GetXaxis()->SetTitle("x / cm");
	histDpXvsY->GetYaxis()->SetTitle("y / cm");
	tupleDplus->Project("histDpXvsY", "Dplusd0x:Dplusd0y");

	TH1D * histDpX = new TH1D("histDpX", "D^{+} x", 100, -0.02, 0.02);
	tupleDplus->Project("histDpX", "Dplusd0x");
	histDpX->SetLineColor(kBlue);
	histDpX->GetXaxis()->SetTitle("x / cm");
	histDpX->GetYaxis()->SetTitle("counts");
	TH1D * histDpY = new TH1D("histDpY", "D^{+} y", 100, -0.02, 0.02);
	tupleDplus->Project("histDpY", "Dplusd0y");
	histDpY->SetLineColor(kBlue);
	histDpY->GetXaxis()->SetTitle("y / cm");
	histDpY->GetYaxis()->SetTitle("counts");
	TH1D * histDpZ = new TH1D("histDpZ", "D^{+} z", 100, -0.02, 0.20);
	tupleDplus->Project("histDpZ", "Dplusd0z");
	histDpZ->SetLineColor(kBlue);
	histDpZ->GetXaxis()->SetTitle("z / cm");
	histDpZ->GetYaxis()->SetTitle("counts");

	// ###############################
	// Loops, if necessary
	// ###############################

	int nOfExactlyOneDpOneDm = 0;

	// D+
	for (int i = 0; i < nOfEntriesDplus; ++i) {
		tupleDplus->GetEntry(i);
		if (i%100 == 0) std::cout << "Processing event " << i << std::endl;
		
		nOfDpMesons->Fill(ncandDplus);

		// if (ncandDminus == 1 && ncandDplus == 1) nOfExactlyOneDpOneDm++;

		// ((TH1D *)(histsMomentumDmDp[2]))->Fill(dPlus.m.px);
		// ((TH1D *)(histsMomentumDmDp[3]))->Fill(dPlus.m.py);
		// ((TH1D *)(histsMomentum2DmDp[2]))->Fill(dPlus.m.pt);
		// ((TH1D *)(histsMomentum2DmDp[3]))->Fill(dPlus.m.p);

		// histDpEnergy->Fill(dPlus.m.E);
		// histDpMass->Fill(dPlus.m.m);
		// histDpCharge->Fill(dPlus.m.chg);
		// histDpPdgId->Fill(dPlus.m.pdg);

		// histDpResPt->Fill((dPlus.m.pt - trDplus.m.pt));
		// histDpResPtFit->Fill((dPlusFit.m.pt - trDplus.m.pt));
		// histDpResM->Fill((dPlus.m.m - trDplus.m.m));
	}
	// D-
	for (int i = 0; i < nOfEntriesDminus; ++i) {
		tupleDminus->GetEntry(i);

		nOfDmMesons->Fill(ncandDminus);

		// ((TH1D *)(histsMomentumDmDp[0]))->Fill(dMinus.m.px);
		// ((TH1D *)(histsMomentumDmDp[1]))->Fill(dMinus.m.py);
		// ((TH1D *)(histsMomentum2DmDp[0]))->Fill(dMinus.m.pt);
		// ((TH1D *)(histsMomentum2DmDp[1]))->Fill(dMinus.m.p);

		// histDmEnergy->Fill(dMinus.m.E);
		// histDmMass->Fill(dMinus.m.m);
		// histDmCharge->Fill(dMinus.m.chg);
		// histDmPdgId ->Fill(dMinus.m.pdg);
	}

	// ###############################
	// Displaying and Printing
	// ###############################

	nOfDpMesons->Scale(1/(float)nOfGeneratedEvents);
	nOfDmMesons->Scale(1/(float)nOfGeneratedEvents);

	TCanvas * c1 = new TCanvas("c1", "Comparison of D Multiplicities", 0, 0, 800, 500);
	nOfDpMesons->GetXaxis()->SetTitle("#(D^{#pm}/event)");
	nOfDpMesons->GetYaxis()->SetTitle("# events (rel.)");
	nOfDpMesons->Draw("HIST");
	nOfDmMesons->Draw("HIST SAME");
	std::cout << "#(Exactly 1 D- and D+ at the same time) = " << nOfExactlyOneDpOneDm << " (" << (float)nOfExactlyOneDpOneDm/(float)nOfGeneratedEvents * 100 << "%)" << std::endl;
	// gStyle->SetOptStat(0);
	TLegend * leg = andi::plainLegend();
	// // leg->SetHeader("Efficiencies");
	leg->AddEntry(nOfDpMesons, "# D^{+}/event", "l");
	leg->AddEntry(nOfDmMesons, "# D^{-}/event", "l");
	leg->Draw();
	nOfDpMesons->GetXaxis()->SetTitleOffset(0.95);
	andi::saveCanvas(c1, macroBasename);


	TCanvas * c2 = new TCanvas("c2", "Momenta of D Mesons", 0, 0, 800, 500);
	for (int i = 0; i < 4; i++) ((TH1D*)(histsMomentumDmDp[i]))->Scale(1/(float)nOfGeneratedEvents);
	divideCanvasAndDraw(c2, &histsMomentumDmDp, "HIST");


	TCanvas * c3 = new TCanvas("c3", "Momenta of D Mesons (2)", 0, 0, 800, 500);
	for (int i = 0; i < 4; i++) ((TH1D*)(histsMomentum2DmDp[i]))->Scale(1/(float)nOfGeneratedEvents);
	divideCanvasAndDraw(c3, &histsMomentum2DmDp, "HIST");


	TObjArray dMinusProperties;
	dMinusProperties.Add(histDmEnergy);
	dMinusProperties.Add(histDmMass);
	dMinusProperties.Add(histDmCharge);
	dMinusProperties.Add(histDmPdgId);
	for (int i = 0; i < 4; i++) ((TH1D*)(dMinusProperties[i]))->Scale(1/(float)nOfGeneratedEvents);
	TCanvas * c4 = new TCanvas("c4", "Other properties of of D Minus Mesons", 0, 0, 800, 500);
	
	divideCanvasAndDraw(c4, &dMinusProperties, "HIST");


	TObjArray dPlusProperties;
	dPlusProperties.Add(histDpEnergy);
	dPlusProperties.Add(histDpMass);
	dPlusProperties.Add(histDpCharge);
	dPlusProperties.Add(histDpPdgId);
	for (int i = 0; i < 4; i++) ((TH1D*)(dPlusProperties[i]))->Scale(1/(float)nOfGeneratedEvents);
	TCanvas * c5 = new TCanvas("c5", "Other properties of of D Plus Mesons", 0, 0, 800, 500);
	
	divideCanvasAndDraw(c5, &dPlusProperties, "HIST");

	// scaleHistogram(histDpResPt, nOfGeneratedEvents);
	TCanvas * c6 = new TCanvas("c6", "D Plus pt resolution", 0, 0, 800, 500);
	histDpResPt->Draw("HIST");
	TF1 * fitDpResPt = andi::gaussFit(histDpResPt, true);
	fitDpResPt->Draw("SAME");
	andi::makePadTitleAndDraw(histDpResPt);
	andi::saveCanvas(c6, macroBasename);

	TCanvas * c7 = new TCanvas("c7", "D Plus pt resolution (mass constraint fit)", 0, 0, 800, 500);
	histDpResPtFit->Draw("HIST");
	TF1 * fitDpResPtFit = andi::gaussFit(histDpResPtFit, true);
	fitDpResPtFit->Draw("SAME");
	andi::makePadTitleAndDraw(histDpResPtFit);
	andi::saveCanvas(c7, macroBasename);

	TCanvas * c8 = new TCanvas("c8", "D Plus m resolution", 0, 0, 800, 500);
	histDpResM->Draw("HIST");
	TF1 * fitDpResM = andi::gaussFit(histDpResM, true);
	fitDpResM->Draw("SAME");
	andi::makePadTitleAndDraw(histDpResM);
	andi::saveCanvas(c8, macroBasename);

	TCanvas * c9 = new TCanvas("c9", "D Plus p resolution", 0, 0, 800, 500);
	histDpResP->Draw("HIST");
	TF1 * fitDpResP = andi::gaussFit(histDpResP, true);
	fitDpResP->Draw("SAME");
	andi::makePadTitleAndDraw(histDpResP);
	andi::saveCanvas(c9, macroBasename);

	TCanvas * c10 = new TCanvas("c10", "D Plus position distribution", 0, 0, 800, 500);
	histDpXvsY->Draw("COLZ");
	andi::makePadTitleAndDraw((TH1D*)histDpXvsY);
	andi::saveCanvas(c10, macroBasename);

	TCanvas * c11 = new TCanvas("c11", "D Plus x position", 0, 0, 800, 500);
	histDpX->Draw("COLZ");
	TF1 * fitDpX = andi::gaussFit(histDpX, true);
	fitDpX->Draw("SAME");
	andi::makePadTitleAndDraw(histDpX);
	andi::saveCanvas(c11, macroBasename);

	TCanvas * c12 = new TCanvas("c12", "D Plus y position", 0, 0, 800, 500);
	histDpY->Draw("COLZ");
	TF1 * fitDpY = andi::gaussFit(histDpY, true);
	fitDpY->Draw("SAME");
	andi::makePadTitleAndDraw(histDpY);
	andi::saveCanvas(c12, macroBasename);

	TCanvas * c13 = new TCanvas("c13", "D Plus z position", 0, 0, 800, 500);
	histDpZ->Draw("COLZ");
	// TF1 * fitDpZ = andi::gaussFit(histDpZ, true);
	// fitDpZ->Draw("SAME");
	andi::makePadTitleAndDraw(histDpZ);
	andi::saveCanvas(c13, macroBasename);


	TFile * out = TFile::Open("eval.root", "RECREATE");
	out->cd();
	nOfDpMesons->Write();
	histsMomentumDmDp.Write();
	histsMomentum2DmDp.Write();
	dMinusProperties.Write();
	dPlusProperties.Write();
	histDpResPt->Write();
	histDpResPtFit->Write();
	histDpResM->Write();
	histDpResP->Write();
	c8->Write();
	histDpXvsY->Write();
	histDpX->Write();
	histDpY->Write();
	histDpZ->Write();
}
int main() {
	eval();
	return 0;
}