class RhoCandList;
class RhoCandidate;
class PndAnaPidSelector;
class PndAnaPidCombiner;
class PndAnalysis;
class RhoTuple;
#include <map>

namespace andi {
	void qaFitter(PndRhoTupleQA * qa, TString pre, RhoCandidate * c, RhoTuple * n, RhoFitterBase * fit, bool skip = false) {
		if (0 == n) return;

		if (!skip) {
			qa->qaComp(pre, c, n);
			n->Column(pre + "Chi2", (float) fit->GetChi2(), 0.0f);
			n->Column(pre + "Ndf", (float) fit->GetNdf(), 0.0f);
			n->Column(pre + "Prob", (float) fit->GetProb(), 0.0f);
		}
	}
}


void ana(TString prefix = "", int nevts = 0) {
	// *** some variables
	int i = 0, j = 0, k = 0, l = 0;
	gStyle->SetOptFit(1011);
	bool doPreAnalysis = true;
	
	// *** the output file for FairRunAna
	TString OutFile = prefix + "output.root";  
					
	// *** the files coming from the simulation
	TString inPidFile  = prefix + "pid_complete_CH.root";    // this file contains the PndPidCandidates and McTruth
	TString inParFile  = prefix + "simparams.root";
	
	// *** PID table with selection thresholds; can be modified by the user
	TString pidParFile = TString(gSystem->Getenv("VMCWORKDIR"))+"/macro/params/all.par";	
	
	// *** initialization
	FairLogger::GetLogger()->SetLogToFile(kFALSE);
	FairRunAna* fRun = new FairRunAna();
	FairRuntimeDb* rtdb = fRun->GetRuntimeDb();
	fRun->SetInputFile(inPidFile);
	
	// *** setup parameter database 	
	FairParRootFileIo* parIO = new FairParRootFileIo();
	parIO->open(inParFile);
	FairParAsciiFileIo* parIOPid = new FairParAsciiFileIo();
	parIOPid->open(pidParFile.Data(), "in");
	
	rtdb->setFirstInput(parIO);
	rtdb->setSecondInput(parIOPid);
	rtdb->setOutput(parIO);  
	
	fRun->SetOutputFile(OutFile);
	fRun->Init(); 
	
	// *** create an output file for all histograms
	TString anaOutput = prefix + "output_ana";
	anaOutput += "-";
	anaOutput += nevts;
	anaOutput += ".root";
	TFile * out = TFile::Open(anaOutput, "RECREATE");
	
	// *** create ntuples for J/psi and psi(2S)
	// RhoTuple *njpsi = new RhoTuple("njpsi","J/psi Analysis");
	// RhoTuple *npsip = new RhoTuple("npsip","psi' Analysis");
	RhoTuple * ntpDp = new RhoTuple("ntpDp", "Dp Analysis");
	if (doPreAnalysis) RhoTuple * ntpPreParticles = new RhoTuple("ntpPre", "Pre Analysis");
	RhoTuple * ntpDm = new RhoTuple("ntpDm", "Dm Analysis");
	RhoTuple * nmc  = new RhoTuple("nmc",  "Mctruth info");
	
	// *** the data reader object
	PndAnalysis* theAnalysis = new PndAnalysis();
	if (nevts == 0 || nevts > theAnalysis->GetEntries()) nevts = theAnalysis->GetEntries();
	
	// *** RhoCandLists for the analysis
	RhoCandList kminus, piplus, piplus2, kplus, piminus, piminus2, dpluslist, dminuslist, all, mclist;
	
	// *** Mass selector for the jpsi cands
	double m0_dplus = TDatabasePDG::Instance()->GetParticle("D+")->Mass();   // Get nominal PDG mass of D+
	std::cout << "mass d+ = " << m0_dplus << std::endl;
	RhoMassParticleSelector * dPlusMassSelector = new RhoMassParticleSelector("dp", m0_dplus, 0.3);  // 0.3 is the full width of the window!
	RhoMassParticleSelector * dMinusMassSelector = new RhoMassParticleSelector("dm", m0_dplus, 0.3);
	
	// *** Pid Selection Algorithms
	TString pidSelection = "PidAlgoIdealCharged";  // Use MC info to set P=1 for ISPARTICLE and P=0 for ISNOTPARTICLE
	
	// *** the lorentz vector of the initial psi(2S)
	double pbarmom = 6.5;
	double m0_p = TDatabasePDG::Instance()->GetParticle("proton")->Mass();   // Get nominal PDG mass of the proton
	TLorentzVector ini(0, 0, pbarmom, sqrt(m0_p*m0_p + pbarmom*pbarmom) + m0_p);
	// TLorentzVector ini(0, 0, 6.231552, 7.240065);

	PndRhoTupleQA qa(theAnalysis, pbarmom); 

	// ***
	// the event loop
	// ***
	while (theAnalysis->GetEvent() && i++ < nevts) {
		if ((i%100) == 0) cout << "evt " << i << endl;
		// cout << "evt " << i << endl;
			
		// *** get MC list and store info in ntuple
	 	theAnalysis->FillList(mclist, "McTruth");

		qa.qaMcList("",   mclist, nmc);
		nmc->DumpData();

		// *** Setup event shape object
		theAnalysis->FillList(all,  "All", pidSelection);
		PndEventShape evsh(all, ini, 0.05, 0.1);
		
		// *** Select with no PID info ('All'); type and mass are set 		
		// theAnalysis->FillList(muplus,  "MuonAllPlus",	pidSelection);
		// theAnalysis->FillList(muminus, "MuonAllMinus",	pidSelection);
		// theAnalysis->FillList(piplus,  "PionAllPlus",	pidSelection);
		// theAnalysis->FillList(piminus, "PionAllMinus",	pidSelection);
		theAnalysis->FillList(kplus,  "KaonBestPlus", pidSelection);  // Best = take only best probablity for Kaon
		theAnalysis->FillList(kminus, "KaonBestMinus", pidSelection);
		theAnalysis->FillList(piplus,  "PionBestPlus", pidSelection);
		theAnalysis->FillList(piminus, "PionBestMinus", pidSelection);
			
		// *****
		// *** combinatorics for D+ -> K- Pi+ Pi+ and the antiparticle's process
		// *****

		// dpluslist.Cleanup();
		dpluslist.Combine(kminus, piplus, piplus);
		dpluslist.Select(dPlusMassSelector);
		dpluslist.SetType(411);

		int nDplusMc = theAnalysis->McTruthMatch(dpluslist);

		// dminuslist.Cleanup();
		dminuslist.Combine(kplus, piminus, piminus);
		dminuslist.Select(dMinusMassSelector);
		dminuslist.SetType(-411);
		int nDminusMc = theAnalysis->McTruthMatch(dminuslist);
		
		if (doPreAnalysis) {
			ntpPreParticles->Column("ev",  (Float_t) i);
			ntpPreParticles->Column("nDplus", (Int_t) dpluslist.GetLength());
			ntpPreParticles->Column("nDminus", (Int_t) dminuslist.GetLength());
			ntpPreParticles->Column("nKplus", (Int_t) kplus.GetLength());
			ntpPreParticles->Column("nKminus", (Int_t) kminus.GetLength());
			ntpPreParticles->Column("nPiplus", (Int_t) piplus.GetLength());
			ntpPreParticles->Column("nPiminus", (Int_t) piminus.GetLength());

			ntpPreParticles->DumpData();
		}
		// ***
		// Needed for ntuple: how good is the fit of the current d meson. is it the best fit of all the reconstructed d mesons of the current event? or is there a d meson candidate with a better fit.
		// 1 - of all candidates, this is the best-fitting candidate
		// 2 - of all candidates, there's one better fitting candidate; this is the 2nd best...
		// ..
		// nCand - this is the worst fitting candidate of all candidates
		// 
		// negative number: same scheme as before, but the chi2 probability cut (> 0.01) is not passed
		// ***
		std::map<double, int> dP_vtx_chi2ForIndex_good, dP_vtx_chi2ForIndex_bad;
		std::map<double, int> dP_mass_chi2ForIndex_good, dP_mass_chi2ForIndex_bad;
		for (j = 0; j < dpluslist.GetLength(); ++j) {
			PndKinVtxFitter vertexFitter(dpluslist[j]);
			// vertexFitter.SetVerbose();
			vertexFitter.Fit();
			bool failVtxProb = TMath::IsNaN(vertexFitter.GetProb());
			bool failVtxChi2 = TMath::IsNaN(vertexFitter.GetChi2());
			if (!failVtxChi2 && !failVtxProb) {
				// std::cout << "Prob = " << vertexFitter.GetProb() << std::endl;
				// std::cout << ", Chi2 = " << vertexFitter.GetChi2() << std::endl;
				if (vertexFitter.GetProb() > 0.01) {
					dP_vtx_chi2ForIndex_good[vertexFitter.GetChi2()] = j;
					// std::cout << "i = " << i << ", j = " << j << " -- Case 1: Prob > 0.01 (" << vertexFitter.GetProb() << "). Chi2 = " << vertexFitter.GetChi2() << " saved " << j << std::endl;
				} else {  // <= 0.01
					dP_vtx_chi2ForIndex_bad[vertexFitter.GetChi2()] = j;
					// std::cout << "i = " << i << ", j = " << j << " -- Case 2: Prob < 0.01 (" << vertexFitter.GetProb() << "). Chi2 = " << vertexFitter.GetChi2() << " saved " << j << std::endl;
				}
				RhoCandidate * dplusfit = dpluslist[j]->GetFit();
				PndKinFitter massFitter(dplusfit);
				massFitter.AddMassConstraint(m0_dplus);
				massFitter.Fit();
				bool failMassProb = TMath::IsNaN(massFitter.GetProb());
				bool failMassChi2 = TMath::IsNaN(massFitter.GetChi2());
				if (!failMassProb && !failMassChi2) {
					if (massFitter.GetProb() > 0.01) {
						dP_mass_chi2ForIndex_good[massFitter.GetChi2()] = j;
					} else {
						dP_mass_chi2ForIndex_bad[massFitter.GetChi2()] = j;
					}
				}
			}

		}
		std::map<double, int>::iterator it_good, it_bad;
		std::map<int, int> dP_vtx_indexOfBestFit;
		int runningI = 0;
		for (it_good = dP_vtx_chi2ForIndex_good.begin(); it_good != dP_vtx_chi2ForIndex_good.end(); it_good++, runningI++) {
			dP_vtx_indexOfBestFit[it_good->second] = runningI + 1;
		}
		runningI = 0;
		for (it_bad = dP_vtx_chi2ForIndex_bad.begin(); it_bad != dP_vtx_chi2ForIndex_bad.end(); it_bad++, runningI++) {
			// std::cout << "it->first = " << it_bad->first << ", it->second = " << it_bad->second << std::endl;
			dP_vtx_indexOfBestFit[it_bad->second] = - (runningI + 1);
		}

		std::map<int, int> dP_mass_indexOfBestFit;
		runningI = 0;
		for (it_good = dP_mass_chi2ForIndex_good.begin(); it_good != dP_mass_chi2ForIndex_good.end(); it_good++, runningI++) {
			dP_mass_indexOfBestFit[it_good->second] = runningI + 1;
		}
		runningI = 0;
		for (it_bad = dP_mass_chi2ForIndex_bad.begin(); it_bad != dP_mass_chi2ForIndex_bad.end(); it_bad++, runningI++) {
			// std::cout << "it->first = " << it_bad->first << ", it->second = " << it_bad->second << std::endl;
			dP_mass_indexOfBestFit[it_bad->second] = - (runningI + 1);
		}
		// ***
		// *** do all kind of analysis and store in N-tuple
		// ***
		for (j = 0; j < dpluslist.GetLength(); ++j) {
		    // some general info about event (actually written for each candidate!)
			ntpDp->Column("ev",  (Float_t) i);
			ntpDp->Column("cand",  (Float_t) j);
			ntpDp->Column("evI",  (Int_t) i);
			ntpDp->Column("candI",  (Int_t) j);
			ntpDp->Column("ncandDplus",  (Float_t) dpluslist.GetLength());
			ntpDp->Column("ncandDplusI",  (Int_t) dpluslist.GetLength());
			ntpDp->Column("nmctDplus",  (Float_t) nDplusMc);
			ntpDp->Column("nmctDplusI",  (Int_t) nDplusMc);
			
			// store info about initial 4-vector
			qa.qaP4("beam", ini, ntpDp);
			
			// store information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
			qa.qaComp("Dplus", dpluslist[j], ntpDp);
			
			// store info about event shapes
			qa.qaEventShapeShort("es", &evsh, ntpDp);

			PndKinVtxFitter vertexFitter(dpluslist[j]);  // should actually be a vertex fit
			// kinfit.AddMassConstraint(m0_dplus);
			vertexFitter.Fit();
			// // std::cout << "Fit Chi2/NDF: " << kinfit.GetChi2() << "/" << kinfit.GetNdf() << ", Fit Prob: " << kinfit.GetProb() << std::endl;
			RhoCandidate * dplusfit = dpluslist[j]->GetFit();
			
			andi::qaFitter(&qa, "fVtxDplus", dplusfit, ntpDp, &vertexFitter);
			ntpDp->Column("fVtxHowGood", (Int_t) dP_vtx_indexOfBestFit[j]);
			qa.qaMcDiff("fVtxDplus", dplusfit, ntpDp);

			PndKinFitter massFitter(dplusfit);
			massFitter.AddMassConstraint(m0_dplus);
			massFitter.Fit();
			RhoCandidate * dplusfit_mass = dplusfit->GetFit();
			andi::qaFitter(&qa, "fMassDplus", dplusfit_mass, ntpDp, &massFitter);
			ntpDp->Column("fMassHowGood", (Int_t) dP_mass_indexOfBestFit[j]);
			

			qa.qaMcDiff("Dplus", dpluslist[j], ntpDp);
			
			// store the 4-vector of the truth matched candidate (or a dummy, if not matched to keep ntuple consistent)
			RhoCandidate *truth = dpluslist[j]->GetMcTruth();		
			TLorentzVector lv;
			if (0x0 != truth) lv = truth->P4();
			qa.qaP4("trDplus", lv, ntpDp);
			
			ntpDp->DumpData();
		}
		// D-
		std::map<double, int> dM_vtx_chi2ForIndex_good, dM_vtx_chi2ForIndex_bad;
		std::map<double, int> dM_mass_chi2ForIndex_good, dM_mass_chi2ForIndex_bad;
		for (j = 0; j < dminuslist.GetLength(); ++j) {
			PndKinVtxFitter vertexFitter(dminuslist[j]);
			// vertexFitter.SetVerbose();
			vertexFitter.Fit();
			bool failVtxProb = TMath::IsNaN(vertexFitter.GetProb());
			bool failVtxChi2 = TMath::IsNaN(vertexFitter.GetChi2());
			if (!failVtxChi2 && !failVtxProb) {
				// std::cout << "Prob = " << vertexFitter.GetProb() << std::endl;
				// std::cout << ", Chi2 = " << vertexFitter.GetChi2() << std::endl;
				if (vertexFitter.GetProb() > 0.01) {
					dM_vtx_chi2ForIndex_good[vertexFitter.GetChi2()] = j;
					// std::cout << "i = " << i << ", j = " << j << " -- Case 1: Prob > 0.01 (" << vertexFitter.GetProb() << "). Chi2 = " << vertexFitter.GetChi2() << " saved " << j << std::endl;
				} else {  // <= 0.01
					dM_vtx_chi2ForIndex_bad[vertexFitter.GetChi2()] = j;
					// std::cout << "i = " << i << ", j = " << j << " -- Case 2: Prob < 0.01 (" << vertexFitter.GetProb() << "). Chi2 = " << vertexFitter.GetChi2() << " saved " << j << std::endl;
				}
				RhoCandidate * dminusfit = dminuslist[j]->GetFit();
				PndKinFitter massFitter(dminusfit);
				massFitter.AddMassConstraint(m0_dplus);
				massFitter.Fit();
				bool failMassProb = TMath::IsNaN(massFitter.GetProb());
				bool failMassChi2 = TMath::IsNaN(massFitter.GetChi2());
				if (!failMassProb && !failMassChi2) {
					if (massFitter.GetProb() > 0.01) {
						dM_mass_chi2ForIndex_good[massFitter.GetChi2()] = j;
					} else {
						dM_mass_chi2ForIndex_bad[massFitter.GetChi2()] = j;
					}
				}
			}

		}
		std::map<double, int>::iterator it_good, it_bad;
		std::map<int, int> dM_vtx_indexOfBestFit;
		int runningI = 0;
		for (it_good = dM_vtx_chi2ForIndex_good.begin(); it_good != dM_vtx_chi2ForIndex_good.end(); it_good++, runningI++) {
			dM_vtx_indexOfBestFit[it_good->second] = runningI + 1;
		}
		runningI = 0;
		for (it_bad = dM_vtx_chi2ForIndex_bad.begin(); it_bad != dM_vtx_chi2ForIndex_bad.end(); it_bad++, runningI++) {
			// std::cout << "it->first = " << it_bad->first << ", it->second = " << it_bad->second << std::endl;
			dM_vtx_indexOfBestFit[it_bad->second] = - (runningI + 1);
		}

		std::map<int, int> dM_mass_indexOfBestFit;
		runningI = 0;
		for (it_good = dM_mass_chi2ForIndex_good.begin(); it_good != dM_mass_chi2ForIndex_good.end(); it_good++, runningI++) {
			dM_mass_indexOfBestFit[it_good->second] = runningI + 1;
		}
		runningI = 0;
		for (it_bad = dM_mass_chi2ForIndex_bad.begin(); it_bad != dM_mass_chi2ForIndex_bad.end(); it_bad++, runningI++) {
			// std::cout << "it->first = " << it_bad->first << ", it->second = " << it_bad->second << std::endl;
			dM_mass_indexOfBestFit[it_bad->second] = - (runningI + 1);
		}
		for (j = 0; j < dminuslist.GetLength(); ++j) {
			ntpDm->Column("ev",		 (Float_t) i);
			ntpDm->Column("cand",	 (Float_t) j);
			ntpDm->Column("evI",		 (Int_t) i);
			ntpDm->Column("candI",	 (Int_t) j);
			ntpDm->Column("ncandDminus", (Float_t) dminuslist.GetLength());
			ntpDm->Column("ncandDminusI", (Int_t) dminuslist.GetLength());
			ntpDm->Column("nmctDminus",  (Float_t) nDminusMc);
			ntpDm->Column("nmctDminusI",  (Int_t) nDminusMc);

			qa.qaP4("beam", ini, ntpDm);

			qa.qaComp("Dminus", dminuslist[j], ntpDm);
			qa.qaEventShapeShort("es", &evsh, ntpDm);


			PndKinVtxFitter vertexFitter(dminuslist[j]);  // should actually be a vertex fit
			// kinfit.AddMassConstraint(m0_dplus);
			vertexFitter.Fit();
			// // std::cout << "Fit Chi2/NDF: " << kinfit.GetChi2() << "/" << kinfit.GetNdf() << ", Fit Prob: " << kinfit.GetProb() << std::endl;
			RhoCandidate * dminusfit = dminuslist[j]->GetFit();
			
			andi::qaFitter(&qa, "fVtxDminus", dminusfit, ntpDm, &vertexFitter);
			ntpDm->Column("fVtxHowGood", (Int_t) dM_vtx_indexOfBestFit[j]);
			qa.qaMcDiff("fVtxDminus", dminusfit, ntpDm);

			PndKinFitter massFitter(dminusfit);
			massFitter.AddMassConstraint(m0_dplus);
			massFitter.Fit();
			RhoCandidate * dminusfit_mass = dminusfit->GetFit();
			andi::qaFitter(&qa, "fMassDminus", dminusfit_mass, ntpDm, &massFitter);
			ntpDm->Column("fMassHowGood", (Int_t) dM_mass_indexOfBestFit[j]);
			
			qa.qaMcDiff("Dminus", dminuslist[j], ntpDm);

			RhoCandidate *truth = dminuslist[j]->GetMcTruth();		
			TLorentzVector lv;
			if (truth) lv = truth->P4();
			qa.qaP4("trDminus", lv, ntpDm);

			ntpDm->DumpData();
		}	
	}
	
	// *** write out all the histos
	out->cd();

	ntpDp->GetInternalTree()->Write();
	ntpDm->GetInternalTree()->Write();
	nmc->GetInternalTree()->Write();
	if (doPreAnalysis) ntpPreParticles->GetInternalTree()->Write();
		
	out->Save();
}