#include "TFile.h"
#include "TH1D.h"
#include "TCanvas.h"
#include "TStopwatch.h"
#include "TROOT.h"
#include "TSystem.h"
#include "TTree.h"
#include "TString.h"
#include <vector>
gROOT->LoadMacro("/home/ajay/pandaroot/gconfig/rootlogon.C");
rootlogon();
//copied from pnddata/PndDetectorList.h because it is not in the default include path
enum DetectorId {
/** kDCH must be the 1st id, and kHYP must be the last one. Please put new detectors in between!! **/
    kDCH,kDRC,kDSK,kEMC,kGEM,kLUMI,kMDT,kMVD,kRPC,kSTT,kTPC,kTOF,kFTS,kHYPG,kHYP};

void SaveAndUpdateHisto(TH1* currenthisto, TFile& storagefile)
{
	if (storagefile.Get(currenthisto->GetName()) != 0) {
		currenthisto->Add((TH1*)storagefile.Get(currenthisto->GetName()));
	}
	cout << currenthisto->GetName() << ": " << currenthisto->GetEntries() << endl;
	currenthisto->Write();
}

//void test2(TString fname="pid_llbar10^3.root",TString simfname="points_llbar10^3.root", int nevts=0)
void test2(TString fname="llbar_pid.root",TString simfname="llbar_Sim.root", int nevts=10)
{

gROOT->LoadMacro("$VMCWORKDIR/gconfig/rootlogon.C");rootlogon();

PndEventReader evr(fname);
if (nevts==0) nevts=evr.GetEntries();

// Event Statistics Histos
TH1F * hrecovertex = new TH1F("Resolution","",100,0,10);

TH2D* hacceptancephithetaprotonp = new TH2D("hacceptancephithetaprotonup", "Theta-Phi Acceptance p+", 90, -180, 180, 45, 0, 180);
TH2D* hacceptancephithetapim = new TH2D("hacceptancephithetapim", "Theta-Phi Acceptance pi-", 90, -180, 180, 45, 0, 180);
TH2D* hpthetaprotonp = new TH2D("hpthetaprotonp", "Theta vs P p+", 120, 0, 6, 90, 0, 180);
TH2D* hpthetapim = new TH2D("hpthetapim", "Theta vs P pi-", 120, 0, 6, 90, 0, 180);
TH1D* hprotonpperevent = new TH1D("hprotonpperevent","Pi+ per Event", 20, 0, 20);
TH1D* hpimperevent = new TH1D("hpimperevent","Pi- per Event", 20, 0, 20);
TH1D* hlambdarawmass = new TH1D("hlambdarawmass","lambda Raw Mass",100,1.065,1.165);
TH1D* hlambdamselmass = new TH1D("hlambdamselmass","lambda Selected Mass",100,0,2);
TH1D* hlambdavtxchi2 = new TH1D("hlambdavtxchi2","lambda Vertex Fit Chi2",200,0,100);

TH1D* hlambdavtxrejectmass = new TH1D("hlambdavtxrejectmass","lambda Rejected Mass",100,0,2);
TH1D* hlambdavtxbestmass = new TH1D("hlambdavtxbestmass","lambda Best Fit Mass",100,0,2);
TH2D* hpthetamup = new TH2D("hpthetamup", "Theta vs P mu+", 120, 0, 6, 90, 0, 180);
TH2D* hpthetamum = new TH2D("hpthetamum", "Theta vs P mu-", 120, 0, 6, 90, 0, 180);

TH2D* hplptpmc = new TH2D("hplptpmc", "Pt vs Pl mu+ MC", 160, -2, 6, 100, 0, 2);
TH2D* hplptmmc = new TH2D("hplptmmc", "Pt vs Pl mu- MC", 160, -2, 6, 100, 0, 2);
TH2D* hthetapthetammc = new TH2D("hpthetapthetammc", "Theta mu- vs Theta mu+ MC", 90, 0, 180, 90, 0, 180);
TH2D* hmjpsichi2 = new TH2D("hmjpsichi2", "chi2 vs J/psi mass", 80, 1, 2, 80, -1, 15);
//TH1D* hjpsivertexfit = new TH1D("hjpsivertexfit","J/psi Vertex Resolution (After Vertex Fit)",200,0,0.5);
TH1D* hjpsivertexfit = new TH1D("hjpsivertexfit","Lambda Vertex Resolution (After Vertex Fit)",100,-100,100);

//TCanvas *c =new TCanvas("c","combine",1000,600);
//TCanvas *c1 =new TCanvas("c1","combine",1000,600);
//c->Divide(2,2);
//c1->Divide(2,1);
// particle lists

TCandList protonpbase, pimbase;
TCandList protonp, pim;
TCandList lambdaraw;
TCandList lambdamsel;
TCandList lambdavtx;


TPidMassSelector* lMSel = new TPidMassSelector("lMSel", TRho::Instance()->GetPDG()->GetParticle(3122)->Mass(), 5);

Double_t lambdavtxfitchi2limit =18;
Double_t thetap;
Double_t thetam;

int minstthits = 1;
bool nosttcut = true; // set to false if only events with at least minstthits should be processed
DetectorId detid = kSTT; // set to kSTT or kTPC

int ievt=0,i=0,j=0,k=0;

TFile mcfile(simfname, "READ");
TTree* mctree = (TTree*)mcfile.Get("cbmsim");
TClonesArray* mcarray = new TClonesArray("PndMCTrack");
mctree->SetBranchAddress("MCTrack", &mcarray);


TVector3 lambdastartvertex;
TVector3 protonpmc3;
TVector3 pimmc3;

TLorentzVector ini(0,0,3,3.1422);

// check for acceptance of all charged particles
void checkacceptance(int mymcindex, int particletype, TClonesArray* mymcarray, DetectorId mydetid, int& myaccepted, TH2D* targethisto) {
	if ((PndMCTrack*)mymcarray->At(mymcindex) != 0) {
		PndMCTrack* mymctrack = (PndMCTrack*)mymcarray->At(mymcindex);
		// cout << "mcindex, particle type: " << mymcindex << ", " << mymctrack->GetPdgCode() << endl;
		if ((mymctrack->GetPdgCode() == particletype) && (mymctrack->GetMotherID() == -1)) {
			if (mymctrack->GetNPoints(mydetid) > 0) {
				++myaccepted;
				TVector3 momentum = mymctrack->GetMomentum();
				targethisto->Fill(TMath::RadToDeg()*momentum.Phi(), TMath::RadToDeg()*momentum.Theta());
			}
		} else {
		  //			cout << "mcindex " << mymcindex << " not of particletype " << particletype << endl;
		}
	}
}

// P-theta distribution
void makeptheta(int mymcindex, int particletype, TClonesArray* mymcarray, TH2D* targethisto) {
	if ((PndMCTrack*)mymcarray->At(mymcindex) != 0) {
		PndMCTrack* mymctrack = (PndMCTrack*)mymcarray->At(mymcindex);
		if ((mymctrack->GetPdgCode() == particletype) && (mymctrack->GetMotherID() == -1)) {
			TVector3 momentum = mymctrack->GetMomentum();
			targethisto->Fill(momentum.Mag(), TMath::RadToDeg()*momentum.Theta());
		} else {
		  //			cout << "mcindex " << mymcindex << " not of particletype " << particletype << endl;
		}
	}
}

// find number of reconstructable tracks and of STT hits for mu+, mu-
void processbaseparticles(TCandList& baselist, TCandList& outlist, int particletype, TClonesArray* mymcarray, bool mynosttcut, int myminstthits, int& mykilledtracks) {
  //   cout << "baselist.GetLength(): " << baselist.GetLength() << endl;
	for (int j = 0; j<baselist.GetLength(); ++j) {
		baselist[j].SetMass(TRho::Instance()->GetPDG()->GetParticle(particletype)->Mass());
//		cout << "Particle Type " << particletype << " SttHits: " << baselist[j].GetMicroCandidate().GetSttHits() << endl;
		int mcindex = baselist[j].GetMcIdx();
		//		cout << "mcindex: " << mcindex << endl;
		PndMCTrack* mctrack = 0;
		if (mcindex >= 0) mctrack = (PndMCTrack*)mymcarray->At(mcindex);
//		if (mctrack != 0) cout << "Particle Type " << mctrack->GetPdgCode() << " SttHits: " << baselist[j].GetMicroCandidate().GetSttHits() << endl;
		bool ctwashit = (mynosttcut) || (baselist[j].GetMicroCandidate().GetSttHits() >= myminstthits);
		if ((mctrack != 0) && (mctrack->GetPdgCode() == particletype) && (ctwashit)) outlist.Add(baselist[j]);
//              if ((mctrack != 0) && (mctrack->GetPdgCode() == particletype) && (mctrack->GetMotherID() != -1) && (ctwashit)) outlist.Add(baselist[j]);
		if ((!ctwashit) && (mctrack != 0) && (mctrack->GetMotherID() == -1)) ++mykilledtracks;
		if (!ctwashit) cout << "Particle Type " << particletype << " missed Central Tracker." << endl;
	}

}

// **** loop over all _events_

while (evr.GetEvent() && ievt<nevts) {
 
if (!(ievt%50)) //cout <<"evt "<<ievt<<endl;

  mctree->GetEntry(ievt);

  protonpbase.Cleanup();
  pimbase.Cleanup();

  lambdaraw.Cleanup();
  lambdamsel.Cleanup();
  lambdavtx.Cleanup();

evr.FillList(protonpbase,"ProtonVeryLoosePlus");
evr.FillList(pimbase,"PionVeryLooseMinus");

//lambdaraw.Combine(protonpbase,pimbase);
// check for acceptance of all charged particles
int withinacceptance = 0;
checkacceptance(2, 2212, mcarray, detid, withinacceptance, hacceptancephithetaprotonp);
checkacceptance(3,-211, mcarray, detid, withinacceptance, hacceptancephithetapim);

// P-theta distribution
makeptheta(2, 2212, mcarray, hpthetaprotonp);
makeptheta(3, -211, mcarray, hpthetapim);

int protonpkilled = 0;
int pimkilled = 0;

// find number of reconstructable tracks and of STT hits for pi+, pi-
processbaseparticles(protonpbase, protonp, 2212, mcarray, nosttcut, minstthits, protonpkilled);
processbaseparticles(pimbase, pim, -211, mcarray, nosttcut, minstthits, pimkilled);

//hprotonpperevent->Fill(protonp.GetLength());
//hpimperevent->Fill(pim.GetLength());

TLorentzVector protonpmc (0,0,0,TRho::Instance()->GetPDG()->GetParticle(2212)->Mass());
TLorentzVector pimmc (0,0,0,TRho::Instance()->GetPDG()->GetParticle(-211)->Mass());

//TLorentzVector totalP4 (0,0,0,0);

if ((((PndMCTrack*)mcarray->At(5))!=0) && (((PndMCTrack*)mcarray->At(2))->GetPdgCode() == 2212)) {
	protonpmc = ((PndMCTrack*)mcarray->At(5))->Get4Momentum();
//	cout << "mupmc: " << mupmc.Px() << " " << mupmc.Py() << " " << mupmc.Pz() << " " << mupmc.E() << " " << endl;
	lambdastartvertex = ((PndMCTrack*)mcarray->At(5))->GetStartVertex();
//	cout << "jpsistartvertex: " << jpsistartvertex.x() << " " << jpsistartvertex.y() << " " << jpsistartvertex.z() << " " << endl;
} else {
  	cout << "p+ not found!" << endl;
//	mupmc.SetXYZM(0, 0, 0, TRho::Instance()->GetPDG()->GetParticle(-13)->Mass());
	lambdastartvertex.SetXYZ(0, 0, 0);
}
//totalP4 = totalP4 + mupmc;

if ((((PndMCTrack*)mcarray->At(6))!=0) && (((PndMCTrack*)mcarray->At(6))->GetPdgCode() == -211)) {
	pimmc = ((PndMCTrack*)mcarray->At(6))->Get4Momentum();
//	cout << "mummc: " << mummc.Px() << " " << mummc.Py() << " " << mummc.Pz() << " " << mummc.E() << " " << endl;
} else {
  	cout << "pi- not found!" << endl;
//	mummc.SetXYZM(0, 0, 0, TRho::Instance()->GetPDG()->GetParticle(13)->Mass());
}

protonpmc3 = protonpmc.Vect();
pimmc3 = pimmc.Vect();

hplptpmc->Fill(protonpmc.Pz(),protonpmc.Perp());
hplptmmc->Fill(pimmc.Pz(),pimmc.Perp());

thetap = TMath::RadToDeg()*protonpmc3.Theta();
thetam = TMath::RadToDeg()*pimmc3.Theta();
hthetapthetammc->Fill(thetap,thetam);


// Combine proton+, pi- to lambda
lambdaraw.Combine(protonpbase,pimbase);
//cout << " lambdaraw TCandList entries: " << lambdaraw.GetLength() << endl;

for (j=0;j<lambdaraw.GetLength();++j) {
//	cout << "j: " << j << " lambda raw mass: " << lambdaraw[j].M() << endl;
	hlambdarawmass->Fill(lambdaraw[j].M());
}
//c->cd(1);  
//hlambdarawmass->Draw();

lambdamsel.Select(lambdaraw, lMSel);
//cout << "lambda MSel TCandList entries: " << lambdamsel.GetLength() << endl;
//hjpsimsellength->Fill(jpsimsel.GetLength());

if (lambdamsel.GetLength() < 1) {
// no valid mu+mu- combination, go to end of event loop
        ++ievt;
	continue;
	}

for (j=0;j<lambdamsel.GetLength();++j) {
//	cout << "j: " << j << " jpsi MSel mass: " << jpsimsel[j].M() << endl;
//	protonpsel = lambdamsel[j].Daughter(0)->P4();
//	pimsel = lambdamsel[j].Daughter(1)->P4();
//	cout << "protonpsel: " << protonpsel.Px() << " " << protonpsel.Py() << " " << protonpsel.Pz() << " " << protonpsel.E() << endl;
	//	cout << "mumsel: " << mumsel.Px() << " " << mumsel.Py() << " " << mumsel.Pz() << " " << mumsel.E() << endl;
hlambdamselmass->Fill(lambdamsel[j].M());
}
//c->cd(2); hlambdamselmass->Draw();

// vertex fit for lambda

int bestfitindex = -1;
Double_t bestfitchi2 = lambdavtxfitchi2limit;
Double_t bestfitmass = 0;
for (j=0; j < lambdamsel.GetLength(); ++j) {
	PndKinVtxFitter vtxfitter(lambdamsel[j]);
	//	vtxfitter.AddMassConstraint(TRho::Instance()->GetPDG()->GetParticle(433)->Mass());
	//	vtxfitter.AddMassConstraint(TRho::Instance()->GetPDG()->GetParticle(3122)->Mass());
	vtxfitter.Fit();
	TCandidate fitcand=*(vtxfitter.FittedCand(lambdamsel[j]));
       	TVector3 fitvertex=fitcand.Pos();
	// VAbsVertex  fitvertex=fitcand.Pos();
	//	cout << fitvertex.y() << "dedenadan" <<endl;
	//	cout << "j: " << j << " chi2: " << vtxfitter.GlobalChi2() << " fitted mass: " << fitcand.M() << endl;
//	hlambdavtxchi2->Fill(vtxfitter.GlobalChi2());
	if (vtxfitter.GlobalChi2()<bestfitchi2) {
		if (bestfitchi2 < lambdavtxfitchi2limit) {
			hlambdavtxrejectmass->Fill(bestfitmass);
		}
		bestfitchi2 = vtxfitter.GlobalChi2();
		bestfitindex = j;
		bestfitmass = fitcand.M();
	} else {
		hlambdavtxrejectmass->Fill(fitcand.M());
	}
	
}

//c->cd(3); hlambdavtxrejectmass->Draw();


//process best candidate

if (bestfitchi2 >= lambdavtxfitchi2limit || bestfitindex < 0) {
  // vertex fit; discard event
  //  cout << "bad vertex fit: best chi2: " << bestfitchi2 << endl; 
       ++ievt;
 	continue;
	}

//cout << "bestfitindex: " << bestfitindex << " bestfitchi2: " << bestfitchi2 << " bestfitmass: " << bestfitmass << endl; 
PndKinVtxFitter vtxfitter(lambdamsel[bestfitindex]);
//vtxfitter.AddMassConstraint(TRho::Instance()->GetPDG()->GetParticle(3122)->Mass());
vtxfitter.Fit();
TCandidate fitcand=*(vtxfitter.FittedCand(lambdamsel[bestfitindex]));
TVector3 fitvertex=fitcand.Pos();
// cout << "fitvertex" << fitvertex.Mag() <<endl;
lambdavtx.Add(fitcand);

if (vtxfitter.GlobalChi2()!=bestfitchi2 || fitcand.M()!=bestfitmass) {
        cout << "event "<< ievt << ": best fit values inconsistent!!" << endl;
	cout << "bestfitchi2: " << bestfitchi2 << " vtxfitter.GlobalChi2(): " << vtxfitter.GlobalChi2() << endl;
	cout << "bestfitmass: " << bestfitmass << " fitcand.M(): " << fitcand.M() << endl;
     
	}
//hjpsivtxbestchi2->Fill(vtxfitter.GlobalChi2());
// cout << fitcand.P()<< endl;
hlambdavtxbestmass->Fill(fitcand.M());
//c->cd(4);
//hlambdavtxbestmass->Draw();
hrecovertex->Fill(fitvertex.Mag());
//distance between fitted vertex and MC vertex
TVector3 vertexdisp = lambdastartvertex - fitvertex;
//cout << lambdastartvertex.Mag() - fitvertex.Mag() << "vertexdisp" << endl;
//cout << vertexdisp.Mag() << "vertexdisp" << endl;
// hjpsivertexfit->Fill(vertexdisp.Mag());
 hjpsivertexfit->Fill(vertexdisp.Mag()*(fitcand.M()/fitcand.P()));
// hjpsivertexfit->Fill(fitvertex.Mag());
 //hjpsivertexfit->Fill(lambdastartvertex.Mag());



// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

++ievt;

} // end main loop
 hrecovertex->Draw();
 // hjpsivertexfit->Draw();
} // end macro