//*
#if !defined(__CINT__) || defined(__MAKECINT__)

// PANDA includes
#include "PndAnalysis.h"
#include "PndPidCandidate.h"
#include "Pnd4CFitter.h"
#include "PndKinFitter.h"
#include "PndKinVtxFitter.h"
#include "RhoCandidate.h"
#include "RhoCandList.h"
#include "RhoParticleSelectorBase.h"
#include "RhoGoodPhotonSelector.h"
#include "RhoMassParticleSelector.h"

#include "FairParRootFileIo.h"
#include "FairParAsciiFileIo.h"
#include "FairRecoCandidate.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "FairTrackParam.h"
#include "FairLogger.h"

#include <TDatabasePDG.h>
#include <TFile.h>
#include <TH1.h>
#include <TH2.h>
#include <TStyle.h>

#endif
//*/

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

// *** 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 Y4260_jpsi_ana(int nevts=0)
{
  // Analyze Y4260 decays into J/psi+/pi0+/pi0 (J/psi -> /mu/mu)

  // *** some variables
  int i=0, j=0, k=0, l=0;
  gStyle->SetOptFit(1011);
	
  // *** the output file for FairRunAna
  TString OutFile="outputY4260jpsi.root";  
					
  // *** the files coming from the simulation
  TString inPidFile  = "pid_complete_1.root";    // this file contains the PndPidCandidates and McTruth
  TString inParFile  = "simparams_1.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
  TFile *out = TFile::Open("output_ana_Y4260jpsi.root","RECREATE");
	
  // *** create some histograms
  TH1F *hjpsim_true = new TH1F("hjpsim_true","J/#psi mass from MC mu+mu-",200,0,4.5);
  TH1F *hpsim_true  = new TH1F("hpsim_true","#Y(4260) mass (full truth match)",200,2,7);
  TH1F *hpsim_bkg  = new TH1F("hpsim_bkg","#Y(4260) mass (no full truth match)",200,2,7);
  
  TH1F *hjpsim_all = new TH1F("hjpsim_all","J/#psi mass (all)",200,0,4.5);
  TH1F *hpsim_all  = new TH1F("hpsim_all","#psi(2S) mass (all)",200,2,7);
  TH1F *hpi0_all  = new TH1F("hpi0_all","#pi^0 mass (all)",200,0,0.2);
  
  TH1F *hjpsim_lpid = new TH1F("hjpsim_lpid","J/#psi mass (loose pid)",200,0,4.5);
  TH1F *hpsim_lpid  = new TH1F("hpsim_lpid","#psi(2S) mass (loose pid)",200,2,7);
  
  TH1F *hjpsim_tpid = new TH1F("hjpsim_tpid","J/#psi mass (tight pid)",200,0,4.5);
  TH1F *hpsim_tpid  = new TH1F("hpsim_tpid","#psi(2S) mass (tight pid)",200,2,7);
  
  TH1F *hjpsim_trpid = new TH1F("hjpsim_trpid","J/#psi mass (true pid)",200,0,4.5);
  TH1F *hpsim_trpid  = new TH1F("hpsim_trpid","#psi(2S) mass (true pid)",200,2,7);
  
  TH1F *hjpsim_ftm = new TH1F("hjpsim_ftm","J/#psi mass (full truth match)",200,0,4.5);
  TH1F *hpsim_ftm  = new TH1F("hpsim_ftm","#psi(2S) mass (full truth match)",200,2,7);
  TH1F *hpi0_ftm  = new TH1F("hpi0_ftm","#pi^0 mass (full truth match)",200,0,0.2);
  
  TH1F *hjpsim_nm = new TH1F("hjpsim_nm","J/#psi mass (no truth match)",200,0,4.5);
  TH1F *hpsim_nm  = new TH1F("hpsim_nm","#psi(2S) mass (no truth match)",200,2,7);
  TH1F *hpi0_nm  = new TH1F("hpi0_nm","#pi^0 mass (no truth match)",200,0,0.2);
  TH1F *hpi0_4cf  = new TH1F("hpi0_4cf","#pi^{0} mass (4C fit)",200,0,0.2);
  
  TH1F *hjpsim_diff = new TH1F("hjpsim_diff","J/#psi mass diff to truth",100,-2,2);
  TH1F *hpsim_diff  = new TH1F("hpsim_diff","#psi(2S) mass diff to truth",100,-2,2);
  TH1F *hpi0_diff  = new TH1F("hpi0_diff","#pi^0 mass diff to truth",100,-0.1,0.1);
  
  TH1F *hjpsim_vf   = new TH1F("hjpsim_vf","J/#psi mass (vertex fit)",200,0,4.5);
  TH1F *hpsim_vf   = new TH1F("hpsim_vf","Y4260 mass (vertex fit)",200,2,7);
  TH1F *hjpsim_4cf  = new TH1F("hjpsim_4cf","J/#psi mass (4C fit)",200,2.5,3.5);
  TH1F *hjpsim_mcf  = new TH1F("hjpsim_mcf","J/#psi mass (mass constraint fit)",200,2.5,3.5);
  
  TH1F *hjpsi_chi2_vf  = new TH1F("hjpsi_chi2_vf", "J/#psi: #chi^{2} vertex fit",100,0,10);
  TH1F *hpsi_chi2_vf  = new TH1F("hpsi_chi2_vf", "Y4260: #chi^{2} vertex fit",100,0,10);
  TH1F *hpsi_chi2_4c   = new TH1F("hpsi_chi2_4c",  "#psi(2S): #chi^{2} 4C fit",100,0,250);
  TH1F *hjpsi_chi2_mf  = new TH1F("hjpsi_chi2_mf", "J/#psi: #chi^{2} mass fit",100,0,10);
  
  TH1F *hjpsi_prob_vf  = new TH1F("hjpsi_prob_vf", "J/#psi: Prob vertex fit",100,0,0.4);
  TH1F *hpsi_prob_vf  = new TH1F("hpsi_prob_vf", "Y4260: Prob vertex fit",100,0,0.4);
  TH1F *hpsi_prob_4c   = new TH1F("hpsi_prob_4c",  "#psi(2S): Prob 4C fit",100,0,0.4);
  TH1F *hjpsi_prob_mf  = new TH1F("hjpsi_prob_mf", "J/#psi: Prob mass fit",100,0,0.4);
  
  TH2F *hvpos = new TH2F("hvpos","(x,y) projection of fitted decay vertex",100,-2,2,100,-2,2);
  
  TH1F *hpi0_mult  = new TH1F("hpi0_mult","#pi^{0} daughter multiplicity",10,0,10);
  TH1F *hpi0pi0  = new TH1F("hpi0pi0","Number of #pi^{0}#pi^{0} candidates",10,0,10);
  TH1F *hpi0pi0_new  = new TH1F("hpi0pi0_new","Number of #pi^{0}#pi^{0} candidates",10,0,10);
  TH1F *hpi0_prob_mf  = new TH1F("hpi0_prob_mf", "#pi^{0}: Prob mass fit",100,0,0.4);
  TH1F *hprob_4c  = new TH1F("hprob_4c", "#pi^{0}: Prob mass fit",100,0,0.4);
  TH1F *hpi0eta  = new TH1F("hpi0eta","Number of #pi^{0}#eta candidates",10,0,10);

  //
  // Now the analysis stuff comes...
  //
  
  // *** the data reader object
  PndAnalysis* theAnalysis = new PndAnalysis();
  if (nevts==0) nevts= theAnalysis->GetEntries();
  
  // *** RhoCandLists for the analysis
  RhoCandList muplus, muminus, piplus, piminus, jpsi, psi2s, mclist, photons, pi0s;
	
  // *** Mass selector for the jpsi cands
  double m0_jpsi = TDatabasePDG::Instance()->GetParticle("J/psi")->Mass();   // Get nominal PDG mass of the J/psi
  RhoMassParticleSelector *jpsiMassSel=new RhoMassParticleSelector("jpsi",m0_jpsi,1.0);
	
  // *** Photon selector
  RhoGoodPhotonSelector *photSel = new RhoGoodPhotonSelector("PhS");
  photSel->SetCriterion("loose");

  // *** Mass selector for the pi0
  double m0_pi0 = TDatabasePDG::Instance()->GetParticle("pi0")->Mass();   // Get nominal PDG mass of the pi0
  RhoMassParticleSelector *pi0MassSel=new RhoMassParticleSelector("pi0",m0_pi0,0.030);

  // *** the lorentz vector of the initial Y4260
  TLorentzVector ini(0, 0, 8.682, 9.671);
	
  // ***
  // the event loop
  // ***
  while (theAnalysis->GetEvent() && i++<nevts) {
    if ((i%100)==0) cout<<"evt " << i << endl;
				
    // Get MC truth info
    theAnalysis->FillList(mclist,  "McTruth");
  
    // Get number of J/psi decaying into mu+mu-
    for (int jmc = 0; jmc < mclist.GetLength(); ++jmc) {

      if (mclist[jmc]->PdgCode() != 443) continue; // select only J/psi
      if (mclist[jmc]->NDaughters() != 2) continue; // select only with 2 daughters
      Int_t ok = kTRUE;
      for (k = 0; k < mclist[jmc]->NDaughters(); ++k) 
	if (TMath::Abs(mclist[jmc]->Daughter(k)->PdgCode()) != 13) ok = kFALSE; // select mu
      if (ok) hjpsim_true->Fill(mclist[jmc]->Mass());
    }
    // Get number of pi0 decaying into gam gam
    for (int jmc = 0; jmc < mclist.GetLength(); ++jmc) {
      if (mclist[jmc]->PdgCode() != 111) continue; // select only pi0
      hpi0_mult->Fill(mclist[jmc]->NDaughters());
      /*
      if (mclist[jmc]->NDaughters() == 2) continue;
      if (mclist[jmc]->NDaughters() == 1) { cout << i << endl; exit(0); }
      for (k = 0; k < mclist[jmc]->NDaughters(); ++k) 
	cout << mclist[jmc]->NDaughters() << ": " << k << " " << mclist[jmc]->Daughter(k)->PdgCode() << endl;
      */
    }
				
    // *** Select with no PID info ('All'); type and mass are set 		
    /*
    theAnalysis->FillList(muplus,  "MuonAllPlus");
    theAnalysis->FillList(muminus, "MuonAllMinus");
    theAnalysis->FillList(piplus,  "PionAllPlus");
    theAnalysis->FillList(piminus, "PionAllMinus");
    */
    //*
    //cout << muplus.GetLength() << " " << muminus.GetLength() << " " << piplus.GetLength() << " " << piminus.GetLength() << endl;
    theAnalysis->FillList(muplus,  "MuonTightPlus", "PidAlgoMdtHardCuts");
    theAnalysis->FillList(muminus, "MuonTightMinus", "PidAlgoMdtHardCuts");
    //theAnalysis->FillList(muplus,  "MuonLoosePlus",  "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
    //theAnalysis->FillList(muminus, "MuonLooseMinus", "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
    ///theAnalysis->FillList(piplus,  "PionLoosePlus",  "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
    //theAnalysis->FillList(piminus, "PionLooseMinus", "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
    theAnalysis->FillList(piplus,  "PionAllPlus");
    theAnalysis->FillList(piminus, "PionAllMinus");
    //cout << muplus.GetLength() << " " << muminus.GetLength() << " " << piplus.GetLength() << " " << piminus.GetLength() << endl;
    theAnalysis->FillList(photons, "All");
    photons.Select(photSel);
    //*/
		
    // *** combinatorics for J/psi -> mu+ mu-
    jpsi.Combine(muplus, muminus);

    // *** combinatorics for pi0 -> gam gam
    pi0s.Combine(photons, photons);
		
    // ***
    // *** do the TRUTH MATCH for jpsi
    // ***
    jpsi.SetType(443);
				
    for (j=0;j<jpsi.GetLength();++j) {
      hjpsim_all->Fill( jpsi[j]->M() );
      
      if (theAnalysis->McTruthMatch(jpsi[j])) { 
	hjpsim_ftm->Fill( jpsi[j]->M() );
	hjpsim_diff->Fill( jpsi[j]->GetMcTruth()->M() - jpsi[j]->M() );
	if (jpsi[j]->NDaughters() != 2) exit(0);
	for (k = 0; k < jpsi[j]->NDaughters(); ++k) 
	  if (TMath::Abs(jpsi[j]->Daughter(k)->PdgCode()) != 13) exit(0); // select mu
      } else hjpsim_nm->Fill( jpsi[j]->M() );
    }

    // ***
    // *** do the TRUTH MATCH for pi0
    // ***
    pi0s.SetType(111);
    // *** some rough mass selection
    //pi0s.Select(pi0MassSel);
				
    for (j = 0; j < pi0s.GetLength(); ++j) {
      hpi0_all->Fill( pi0s[j]->M() );
      
      //if (theAnalysis->McTruthMatch(pi0s[j]) && pi0s[j]->NDaughters() == 2 ) { 
      if (theAnalysis->McTruthMatch(pi0s[j]) ) { 
	hpi0_ftm->Fill( pi0s[j]->M() );
	hpi0_diff->Fill( pi0s[j]->GetMcTruth()->M() - pi0s[j]->M() );
	//if (jpsi[j]->NDaughters() != 2) exit(0);
	//for (k = 0; k < pi0s[j]->NDaughters(); ++k) 
	//if (TMath::Abs(jpsi[j]->Daughter(k)->PdgCode()) != 13) exit(0); // select mu
      } else hpi0_nm->Fill( pi0s[j]->M() );
    }
		
    // ***
    // *** do VERTEX FIT (J/psi)
    // ***
    RhoCandList old(jpsi);
    jpsi.Cleanup();
    for (j = old.GetLength()-1; j >= 0; --j) {
      //for (j = 0; j < jpsi.GetLength(); ++j) {
      if (old[j] == NULL) continue;
      PndKinVtxFitter vtxfitter(old[j]);	// instantiate a vertex fitter
      vtxfitter.Fit();
      
      double chi2_vtx = vtxfitter.GetChi2();	// access chi2 of fit
      double prob_vtx = vtxfitter.GetProb();	// access probability of fit
      hjpsi_chi2_vf->Fill(chi2_vtx);
      hjpsi_prob_vf->Fill(prob_vtx);			
      
      if ( prob_vtx > 0.01 ) {
	// when good enough, fill some histos
	RhoCandidate *jfit = old[j]->GetFit();	// access the fitted cand
	TVector3 jVtx=jfit->Pos();	    // and the decay vertex position
	  
	hjpsim_vf->Fill(jfit->M());            
	hvpos->Fill(jVtx.X(),jVtx.Y());
	jpsi.Append(old[j]);
      }
    }
		
    // *** some rough mass selection
    //jpsi.Select(jpsiMassSel);
		
    // *** combinatorics for Y(4260) -> J/psi pi0 pi0
    psi2s.Combine(jpsi, pi0s, pi0s);
    //cout << psi2s.GetLength() << endl;
		
    // ***
    // *** do 4C FIT (initial Y4260 system)
    // ***
    old = psi2s;
    Int_t leng = old.GetLength() - 1;
    psi2s.Cleanup();
    for (j = leng; j >= 0; --j) {
      if (old[j] == NULL) continue;
      PndKinFitter fitter(old[j]);	// instantiate the kin fitter in psi(2S)
      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
      hpsi_chi2_4c->Fill(chi2_4c);
      hpsi_prob_4c->Fill(prob_4c);			
			
      Bool_t accept = kTRUE;
      Double_t mmm[3];
      if ( prob_4c > 0.001 ) {
	// when good enough, fill some histo
	RhoCandidate *jfit = old[j]->Daughter(0)->GetFit();	// get fitted J/psi
	hjpsim_4cf->Fill(jfit->M());
	if (jfit->M() < 3.06 || jfit->M() > 3.14) accept = kFALSE;
	mmm[0] = jfit->M();
	jfit = old[j]->Daughter(1)->GetFit();	// get fitted pi0
	hpi0_4cf->Fill(jfit->M());
	if (jfit->M() < 0.12 || jfit->M() > 0.15) accept = kFALSE;
	mmm[1] = jfit->M();
	jfit = old[j]->Daughter(2)->GetFit();	// get fitted pi0
	hpi0_4cf->Fill(jfit->M());
	if (jfit->M() < 0.12 || jfit->M() > 0.15) accept = kFALSE;
	mmm[2] = jfit->M();
      } else accept = kFALSE;
      if (accept) psi2s.Append(old[j]);
      //else 
      //cout << accept << " " << prob_4c << " " << mmm[0] << " " << mmm[1] << " " <<mmm[2] << endl;
    }		
    //cout << leng << " " << psi2s.GetLength() << endl;
    //psi2s.RemoveClones();
    hpi0pi0->Fill(psi2s.GetLength());
    //if (psi2s.GetLength() != 1) continue; // several candidates

    // ***
    // *** do the TRUTH MATCH for psi(2S)
    // ***
    //AZ psi2s.SetType(30443);
    psi2s.SetType(88881);

    for (j=0;j<psi2s.GetLength();++j) 
      {
	hpsim_all->Fill( psi2s[j]->M() );
	
	if (theAnalysis->McTruthMatch(psi2s[j])) 
	  {
	    hpsim_ftm->Fill( psi2s[j]->M() );
	    hpsim_diff->Fill( psi2s[j]->GetMcTruth()->M() - psi2s[j]->M() );
	  }
	else 
	  hpsim_nm->Fill( psi2s[j]->M() );
      }			

		
		// ***
		// *** do MASS CONSTRAINT FIT (J/psi)
		// ***
		for (j=0;j<jpsi.GetLength();++j) 
		{
			PndKinFitter mfitter(jpsi[j]);		// instantiate the PndKinFitter in psi(2S)
			mfitter.AddMassConstraint(m0_jpsi);	// 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
			//hjpsi_chi2_mf->Fill(chi2_m);
			//hjpsi_prob_mf->Fill(prob_m);			
			
			if ( prob_m > 0.01 )				// when good enough, fill some histo
			{
				RhoCandidate *jfit = jpsi[j]->GetFit();	// access the fitted cand
				hjpsim_mcf->Fill(jfit->M());
			}
		}		
		
		
		// ***
		// *** TRUE PID combinatorics
		// ***
		
		// *** do MC truth match for PID type
		SelectTruePid(theAnalysis, muplus);
		SelectTruePid(theAnalysis, muminus);
		SelectTruePid(theAnalysis, piplus);
		SelectTruePid(theAnalysis, piminus);
				
		// *** all combinatorics again with true PID
		jpsi.Combine(muplus, muminus);
		for (j=0;j<jpsi.GetLength();++j) hjpsim_trpid->Fill( jpsi[j]->M() );
		jpsi.Select(jpsiMassSel);
		
		psi2s.Combine(jpsi, piplus, piminus);
		for (j=0;j<psi2s.GetLength();++j) hpsim_trpid->Fill( psi2s[j]->M() );
		
		
		// ***
		// *** LOOSE PID combinatorics
		// ***
		
		// *** and again with PidAlgoMvd;PidAlgoStt;PidAlgoDrc and loose selection
		theAnalysis->FillList(muplus,  "MuonLoosePlus",  "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		theAnalysis->FillList(muminus, "MuonLooseMinus", "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		theAnalysis->FillList(piplus,  "PionLoosePlus",  "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		theAnalysis->FillList(piminus, "PionLooseMinus", "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		
		jpsi.Combine(muplus, muminus);
		for (j=0;j<jpsi.GetLength();++j) hjpsim_lpid->Fill( jpsi[j]->M() );
		jpsi.Select(jpsiMassSel);
		
		psi2s.Combine(jpsi, piplus, piminus);
		for (j=0;j<psi2s.GetLength();++j) hpsim_lpid->Fill( psi2s[j]->M() );
		
		
		// ***
		// *** TIGHT PID combinatorics
		// ***
		
		// *** and again with PidAlgoMvd;PidAlgoStt and tight selection
		theAnalysis->FillList(muplus,  "MuonTightPlus",  "PidAlgoMdtHardCuts");
		theAnalysis->FillList(muminus, "MuonTightMinus", "PidAlgoMdtHardCuts");
		theAnalysis->FillList(piplus,  "PionLoosePlus",  "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		theAnalysis->FillList(piminus, "PionLooseMinus", "PidAlgoMvd;PidAlgoStt;PidAlgoDrc");
		
		jpsi.Combine(muplus, muminus);
		for (j=0;j<jpsi.GetLength();++j) hjpsim_tpid->Fill( jpsi[j]->M() );
		jpsi.Select(jpsiMassSel);
		//*
		jpsi.SetType(443); 
		for (j = 0; j < jpsi.GetLength(); ++j) { 
		  //theAnalysis->McTruthMatch(jpsi[j]);
		  //cout << j << " jpsi " << theAnalysis->McTruthMatch(jpsi[j]) << " " << jpsi[j]->GetMcTruth() << endl;
		  //if (jpsi[j]->GetMcTruth()) cout << jpsi[j]->GetMcTruth()->PdgCode() << endl;
		}
		//*/
		
		psi2s.Combine(jpsi, piplus, piminus);
		for (j=0;j<psi2s.GetLength();++j) hpsim_tpid->Fill( psi2s[j]->M() );

		// AZ Select true combinations
		psi2s.SetType(88881);

		for (j = 0; j < psi2s.GetLength(); ++j) { 
		  //theAnalysis->McTruthMatch(psi2s[j]);
		  cout << j << " " << theAnalysis->McTruthMatch(psi2s[j]) << " " << psi2s[j]->GetMcTruth() << endl;
		  if (psi2s[j]->GetMcTruth()) cout << psi2s[j]->GetMcTruth()->PdgCode() << endl;
		  if (theAnalysis->McTruthMatch(psi2s[j])) hpsim_true->Fill( psi2s[j]->M() );
		  else hpsim_bkg->Fill( psi2s[j]->M() );
		}
	}
	
	// *** write out all the histos
	out->cd();
	
	out->Write();

	out->Save();
	
}
