class RhoCandList;
class RhoCandidate;
class PndAnaPidSelector;
class PndAnaPidCombiner;
class PndAnalysis;

#include "TStopwatch.h"

// *** routine to only keep PID matched candidates in list
int SelectTruePid(PndAnalysis *ana, RhoCandList &l)
{
	int removed = 0;
	
	for (int ii=l.GetLength()-1;ii>=0;--ii)
	{
		if ( !(ana->McTruthMatch(l[ii])) )
		{
			l.Remove(l[ii]);
			removed++;
		}
	}
	
	return removed;
}

void printCand(RhoCandidate *c)
{
	TLorentzVector lv=c->P4();
	
	cout <<c->PdgCode()<<" ("<<lv.X()<<"/"<<lv.Y()<<"/"<<lv.Z()<<"/"<<lv.E()<<")"<<endl;
}

void countDoubles(RhoCandList &l, int &n1, int &n2, int &n3)
{
	int n_smc  = 0;
	int n_strk = 0;
	int n_both = 0;
	double d = 0.00001;
	
	for (int i=0;i<l.GetLength()-1;++i)
	{
		for (int j=i+1;j<l.GetLength();++j)
		{
			TLorentzVector dl = l[i]->P4() - l[j]->P4();
		
			bool chkmc = (l[i]->GetMcTruth()==l[j]->GetMcTruth());
			bool chktrk = (fabs(dl.X())<d) && (fabs(dl.Y())<d) && (fabs(dl.Z())<d) && (fabs(dl.E())<d);
			if (chkmc) n_smc++;
			if (chktrk) n_strk++;
			if (chktrk && chkmc) n_both++;
		}	
	}
	n1 = n_strk;
	n2 = n_smc;
	n3 = n_both;
}

void emc_anatest(int dtau=5, int dt=0, int nevts=0)
//void emc_anatest(const char* inPidFile, const char* outAnaFile, const char* inParFile, int nevts=0)
{
        TDatabasePDG::Instance()->AddParticle("pbarpSystem","pbarpSystem",1.876544,kFALSE,0.1,0,"",88888);
        TStopwatch timer;
	// *** some variables
	int i=0,j=0, k=0, l=0, nPhotonsPerTB=0, nProg=0;
	gStyle->SetOptFit(1011);
/*
TString prefixIn = "dtdtau_test/hc/hc7g5k_PID_DIST_HR";
TString prefixOut = "dtdtau_test/hc/hc7g5k_ANA_DIST_HR";
TString connect = "_";
TString suffix = "M1.root";
TString sdtau;
TString sdt;*/
	
	// *** the output file for FairRunAna
	TString OutFile="ana_test.root";  
					
	// *** the files coming from the simulation
//	TString inPidFile  = prefixIn+sdtau.Itoa(dtau,10)+connect+sdt.Itoa(dt,10)+suffix;   // this file contains the PndPidCandidates and McTruth
	TString inPidFile  = "PID_DIST_HRauto_hc.root";   // this file contains the PndPidCandidates and McTruth
	TString inParFile  = "simparams_hc7g5kG4.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 outAnaFile = prefixOut+sdtau.Itoa(dtau,10)+connect+sdt.Itoa(dt,10)+suffix;
	TString outAnaFile = "ANA_DIST_HRauto_hc.root"; 
	TFile *out = TFile::Open(outAnaFile,"RECREATE");

	TStopwatch fTimer;
	double time;
	int remtime;
	
	// *** create some histograms
	
	TH1F *hetacm_all = new TH1F("hetacm_all","#eta_{c} mass (all)",400,0,4);
	TH1F *hhcm_all  = new TH1F("hhcm_all","h_{c} mass (all)",500,0,5);
	TH1F *hpi0m_all = new TH1F("hpi0m_all","#pi^{0} mass (all)",200,0,2);
	TH1F *hetam_all = new TH1F("hetam_all","#eta mass (all)",200,0,2);
	
	TH1F *hetacm_4cf  = new TH1F("hetacm_4cf","#eta_{c} mass (4C fit)",400,0,4); 
	TH1F *hetacm_mcf  = new TH1F("hetacm_mcf","#eta_{c} mass (mass constraint fit)",400,0,4); 
	TH1F *hhcm_4cf  = new TH1F("hhcm_4cf","h_{c} mass (4C fit)",500,0,5);
	
	TH1F *hhc_chi2_4c   = new TH1F("hhc_chi2_4c",  "h_{c}: #chi^{2} 4C fit",100,0,250); 
	TH1F *hetac_chi2_mf  = new TH1F("hetac_chi2_mf", "#eta_{c}: #chi^{2} mass fit",100,0,10); 

	TH1F *hhc_prob_4c   = new TH1F("hhc_prob_4c",  "h_{c}: Prob 4C fit",100,0,1); 
	TH1F *hetac_prob_mf  = new TH1F("hetac_prob_mf", "#eta_{c}: Prob mass fit",100,0,1); 

	TH1I *hPhotonsPerTB = new TH1I("hPhotonsPerTB","photon multiplicity per timebunch",40,0,40);

	//
	// Now the analysis stuff comes...
	//
	
	
	// *** the data reader object
	PndAnalysis* theAnalysis = new PndAnalysis();
	//theAnalysis->SetVerbose(3);
	if (nevts==0) nevts= theAnalysis->GetEntries();
	
	// *** RhoCandLists for the analysis
	RhoCandList gamma, pi0, eta, etac, hc;
	
	// *** Mass selector for the cands
	double minE = 0.030; // Minimal energy for gammas
	double m0_etac = TDatabasePDG::Instance()->GetParticle(441)->Mass();   // Get nominal PDG mass of the etac
	double m0_eta = TDatabasePDG::Instance()->GetParticle(221)->Mass();   // Get nominal PDG mass of the eta
	double m0_pi0 = TDatabasePDG::Instance()->GetParticle(111)->Mass();   // Get nominal PDG mass of the pi0
	RhoMassParticleSelector *etacMassSel=new RhoMassParticleSelector("etacMassSel",m0_etac,1.0);
	RhoMassParticleSelector *SelEta=new RhoMassParticleSelector("SelEta",m0_eta,0.05); // "name", central_value, 2*width
	RhoMassParticleSelector *SelPion=new RhoMassParticleSelector("SelPion",m0_pi0,0.04);
	RhoEnergyParticleSelector *SelEnergy = new RhoEnergyParticleSelector("SelEnergy",10.+minE,20.);
	
	// *** the lorentz vector of the initial hc
	TLorentzVector ini(0, 0, 5.61, sqrt(5.61*5.61+9.3827203e-01*9.3827203e-01)+9.3827203e-01);

	pi0.SetType(111);
	eta.SetType(221);
	gamma.SetType(22);
	etac.SetType(441);
	hc.SetType(10443);
	
	// ***
	// the event loop
	// ***
	
	while (theAnalysis->GetEvent() && i++<nevts)
	{
	fTimer.Start(kFALSE);

		if (i-1==nProg*nevts/100) {
		time = fTimer.RealTime();
		fTimer.Continue();
		nProg += 1;
		remtime = (100-nProg)*time/nProg;
		cout << "\r                                                                  " << std::flush;
		cout << "\r[INFO\t] Analysis: " << nProg << "\% completed (about " << remtime << " seconds remaining)" << std::flush;	
		}
		//if ((i%100)==0) cout<<"\r[INFO\t] " << i << " timebunches processed. " << std::flush;

		// *** Select with no PID info ('All'); type and mass are set 		
		theAnalysis->FillList(gamma,  "Neutral");
		gamma.Select(SelEnergy);
//cout << "Combining " << gamma.GetLength() << " photon candidates... " << endl;
		nPhotonsPerTB = gamma.GetLength();

		if (nPhotonsPerTB > 35) {
			cout << "\n[WARNING] Too many photon candidates (" << nPhotonsPerTB << ") in timebunch " << i << " - SKIPPING\n" << endl;
			continue; // HACK to skip timebunches with too much clusters -> causing overflow with combinatorics, and crashes
		}
		hPhotonsPerTB->Fill(nPhotonsPerTB);
		//if (nPhotonsPerTB<7) continue; // SPEEDUP: don't bother with finding pi0s and eta's if there aren't enough photons to make an h_c

		// *** combinatorics to get pi0's and eta's
		pi0.Combine(gamma,gamma);
		pi0.Select(SelPion);
		eta.Combine(gamma,gamma);
		eta.Select(SelEta);

		// *** fill histo's with pi0 and eta candidates (just as a check)
//cout << pi0.GetLength() << " pions and " << eta.GetLength() << " etas found." << endl;
		
		if (pi0.GetLength() == 0) continue; // dont bother with the rest if you didn't find any pi0s

		for (j=0;j<pi0.GetLength();++j) 
		{
			hpi0m_all->Fill( pi0[j]->M() );
		}

		for (j=0;j<eta.GetLength();++j) 
		{
			hetam_all->Fill( eta[j]->M() );
		}
		
		// *** combinatorics for etac -> pi0 pi0 eta
		etac.Combine(pi0, pi0, eta);
		
		// ***
		// *** do the TRUTH MATCH for etac
		// ***
//cout << etac.GetLength() << " eta_c candidates found." << endl;
		for (j=0;j<etac.GetLength();++j) 
		{
			hetacm_all->Fill( etac[j]->M() );
		}
		
		// *** some rough mass selection
		etac.Select(etacMassSel);
		
		// *** combinatorics for hc -> etac gamma
		hc.Combine(etac, gamma);

		// ***
		// *** do the TRUTH MATCH for hc
		// ***
//cout << hc.GetLength() << " h_c candidates found." << endl;
		for (j=0;j<hc.GetLength();++j) 
		{
			hhcm_all->Fill( hc[j]->M() );
		}	

//cout << "Trying 4C fit." << endl;		
		// ***
		// *** do 4C FIT (initial hc system)
		// ***
		for (j=0;j<hc.GetLength();++j) 
		{
			PndKinFitter fitter(hc[j]);	// instantiate the kin fitter for hc
			fitter.Add4MomConstraint(ini);	// set 4 constraint
			fitter.Fit();		            // do fit
			
			double chi2_4c = fitter.GetChi2();	// get chi2 of fit
			double prob_4c = fitter.GetProb();	// access probability of fit
			hhc_chi2_4c->Fill(chi2_4c);
			hhc_prob_4c->Fill(prob_4c);			
			
			if ( prob_4c > 0.01 )			// when good enough, fill some histo
			{
				RhoCandidate *jfit = hc[j]->Daughter(0)->GetFit();	// get fitted etac
				
				hetacm_4cf->Fill(jfit->M());

				RhoCandidate *hfit = hc[j]->GetFit();	// get fitted hc
				hhcm_4cf->Fill(hfit->M() );
			}
		}		
		
//cout << "Trying mass constraint fit." << endl;		
		// ***
		// *** do MASS CONSTRAINT FIT (etac)
		// ***
		for (j=0;j<etac.GetLength();++j) 
		{
			PndKinFitter mfitter(etac[j]);		// instantiate the PndKinFitter for etac
			mfitter.AddMassConstraint(m0_etac);	// add the mass constraint
			mfitter.Fit();						// do fit
			
			double chi2_m = mfitter.GetChi2();	// get chi2 of fit
			double prob_m = mfitter.GetProb();	// access probability of fit
			hetac_chi2_mf->Fill(chi2_m);
			hetac_prob_mf->Fill(prob_m);			
			
			if ( prob_m > 0.01 )				// when good enough, fill some histo
			{
				RhoCandidate *jfit = etac[j]->GetFit();	// access the fitted cand
				hetacm_mcf->Fill(jfit->M());
			}
		}		
		
	}
		
	// *** write out all the histos
	out->cd();
	
	hetacm_all->Write();
	hhcm_all->Write();
	hpi0m_all->Write();
	hetam_all->Write();

	hetacm_4cf->Write();
	hetacm_mcf->Write();
	hhcm_4cf->Write();

	hhc_chi2_4c->Write();
	hetac_chi2_mf->Write();
			
	hhc_prob_4c->Write();
	hetac_prob_mf->Write();

	hPhotonsPerTB->Write();
		
	out->Save();
        
        timer.Stop();
        Double_t rtime = timer.RealTime();
        Double_t ctime = timer.CpuTime();
        cout << endl << endl;
        cout << "Macro finished successfully." << endl;
		cout << "Input file was " << inPidFile << endl;
		cout << "Output file is " << outAnaFile << endl;
        cout << "Real time " << rtime << " s, CPU time " << ctime << " s" << endl;
        cout << endl;
        exit(0);
	
}