#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();

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 lambda_lambdabar(TString fname="llbar_pid.root",TString simfname="llbar_Sim.root", int nevts=100)
{
  
  gROOT->LoadMacro("$VMCWORKDIR/gconfig/rootlogon.C");rootlogon();
  
  PndEventReader evr(fname);
  // if (nevts==0) nevts=evr.GetEntries();
  
  // Event Statistics Histos

  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 Global Chi2",100,0,1);
  
  TH1D* hlambdavtxrejectmass = new TH1D("hlambdavtxrejectmass","lambda Rejected Mass",100,0,2);
  TH1D* hlambdavtxbestmass = new TH1D("hlambdavtxbestmass","lambda Best Fit Mass",100,0,2);
  
  
  TCandList protonpbase, pimbase;
  TCandList protonp, pim;
  TCandList lambdaraw;
  TCandList lambdamsel;
  TCandList lambdavtx;
  

  TPidMassSelector* lMSel = new TPidMassSelector("lMSel", TRho::Instance()->GetPDG()->GetParticle(3122)->Mass(), 0.05);
  
  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) {
    
    for (int j = 0; j<baselist.GetLength(); ++j) {
      baselist[j].SetMass(TRho::Instance()->GetPDG()->GetParticle(particletype)->Mass());

      int mcindex = baselist[j].GetMcIdx();
      //		cout << "mcindex: " << mcindex << endl;
      PndMCTrack* mctrack = 0;
      if (mcindex >= 0) mctrack = (PndMCTrack*)mymcarray->At(mcindex);
	
      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");
    


    if ((((PndMCTrack*)mcarray->At(2))!=0) && (((PndMCTrack*)mcarray->At(2))->GetPdgCode() == 2212)) {
      protonpmc = ((PndMCTrack*)mcarray->At(2))->Get4Momentum();
      
      lambdastartvertex = ((PndMCTrack*)mcarray->At(2))->GetStartVertex();
      
    } else {
      cout << "p+ not found!" << endl;
      lambdastartvertex.SetXYZ(0, 0, 0);
    }
    

    if ((((PndMCTrack*)mcarray->At(3))!=0) && (((PndMCTrack*)mcarray->At(3))->GetPdgCode() == -211)) {
      pimmc = ((PndMCTrack*)mcarray->At(3))->Get4Momentum();
      
    } else {
      cout << "pi- not found!" << endl;
      
    }
    

    // Combine proton+, pi- to lambda
    
    lambdaraw.Combine(protonpbase,pimbase);
    
    for (j=0;j<lambdaraw.GetLength();++j) {
      
    hlambdarawmass->Fill(lambdaraw[j].M());
    }
    
    lambdamsel.Select(lambdaraw, lMSel);
    
    if (lambdamsel.GetLength() < 1) {
      
      ++ievt;
      continue;
    }
    
    for (j=0;j<lambdamsel.GetLength();++j) 
      {
	hlambdamselmass->Fill(lambdamsel[j].M());
      }
 

// vertex fit for lambda

    int bestfitindex = -1;
    Double_t bestfitchi2 = lambdavtxfitchi2limit;
    Double_t bestfitmass = 0;
    for (j=0; j < lambdaraw.GetLength(); ++j) {
      PndKinVtxFitter vtxfitter(lambdaraw[j]);
      
      vtxfitter.Fit();
      TCandidate fitcand=*(vtxfitter.FittedCand(lambdaraw[j]));
      TVector3 fitvertex=fitcand.Pos();
      cout <<"print=="<<vtxfitter.GlobalChi2()<<endl;
      hlambdavtxchi2->Fill(vtxfitter.GlobalChi2());

      if (vtxfitter.GlobalChi2()<bestfitchi2) {
	/*	if (bestfitchi2 < lambdavtxfitchi2limit) 

	{
	  hlambdavtxrejectmass->Fill(bestfitmass);
	  }*/
		bestfitchi2 = vtxfitter.GlobalChi2();
		bestfitindex = j;
		bestfitmass = fitcand.M();
                hlambdavtxrejectmass->Fill(bestfitmass);
      }/* else {
	hlambdavtxrejectmass->Fill(fitcand.M());
	}*/
    }
    
    

    //process best candidate
    
    if (bestfitchi2 >= lambdavtxfitchi2limit || bestfitindex < 0) {
      ++ievt;
      continue;
	}
    
    
    PndKinVtxFitter vtxfitter(lambdaraw[bestfitindex]);
    vtxfitter.Fit();
    TCandidate fitcand=*(vtxfitter.FittedCand(lambdaraw[bestfitindex]));
    
    
    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;
	
    }
        hlambdavtxbestmass->Fill(fitcand.M());
 
 
// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
 
 ++ievt;
 
  } // end main loop
  TCanvas *c =new TCanvas("c","lambda",1000,800);
  c->Divide(1,1);
  // c->cd(1); hlambdarawmass->Draw();
  // c->cd(2); hlambdamselmass->Draw();
  // c->cd(3); hlambdavtxrejectmass->Draw();
  // c->cd(4); hlambdavtxbestmass->Draw();
  c->cd(1);  hlambdavtxchi2->Draw();
  // c->cd(2);  hlambdavtxbestmass->Draw();
  // c->cd(3);  hlambdavtxrejectmass->Draw();
  } // end macro