// $Id$
//
//    File: JEventProcessor_CascadeStudies.cc
// Created: Thu Jul 11 10:46:01 EDT 2019
// Creator: staylor (on Linux ifarm1401.jlab.org 3.10.0-327.el7.x86_64 x86_64)
//

#include "JEventProcessor_CascadeStudies.h"
using namespace jana;

#define LAMBDA_CUT(mass) (fabs(mass-ParticleMass(Lambda))<0.012)
#define CASCADE_CUT(mass) (fabs(mass-ParticleMass(Xi0))<0.018)

// Routine used to create our JEventProcessor
#include <JANA/JApplication.h>
#include <JANA/JFactory.h>
extern "C"{
void InitPlugin(JApplication *app){
	InitJANAPlugin(app);
	app->AddProcessor(new JEventProcessor_CascadeStudies());
}
} // "C"


//------------------
// JEventProcessor_CascadeStudies (Constructor)
//------------------
JEventProcessor_CascadeStudies::JEventProcessor_CascadeStudies()
{

}

//------------------
// ~JEventProcessor_CascadeStudies (Destructor)
//------------------
JEventProcessor_CascadeStudies::~JEventProcessor_CascadeStudies()
{

}

//------------------
// init
//------------------
jerror_t JEventProcessor_CascadeStudies::init(void)
{
  // This is called once at program startup.
  japp->CreateLock("myCascadelock");

  gDirectory->mkdir("CascadeAnalysis")->cd();
  CascadeConfidenceLevel=new TH1F("CascadeConfidenceLevel","CL",100,0,1);
  CascadeChiSq=new TH1F("CascadeChisSq","#chi^{2}/ndf",1000,0,100);
  
  P2PimMass=new TH1F("Xi_P2PimMass","M(p#pi-#pi-)",300,1.2,1.5); 
  PPimMass=new TH1F("Xi_PPimMass","M(p#pi-)",300,1.,1.3); 
  LambdaPimMass=new TH1F("Xi_LambaPimMass","M(#Lambda-#pi-)",300,1.2,1.5); 
  
  gDirectory->mkdir("PID")->cd();
 
  Xi_KpTOFdt=new TH2F("Xi_KpTOFdt","#deltat(TOF)",100,0,10,200,-5,5);	
  Xi_KpFCALdt=new TH2F("Xi_KpFCALdt","#deltat(FCAL)",100,0,10,200,-5,5);
  Xi_KpBCALdt=new TH2F("Xi_KpBCALdt","#deltat(BCAL)",100,0,10,200,-5,5);
  Xi_KpdEdxCDC=new TH2F("Xi_KpdEdxCDC","dEdx vs p/M, proton cands.",1000,0,10,500,0,25.0);	
  Xi_KpdEdxFDC=new TH2F("Xi_KpdEdxFDC","dEdx vs p/M, proton cands.",1000,0,10,500,0,25.0); 
  Xi_KpEOverPFCAL=new TH2F("Xi_KpEOverPFCAL","E/p vs p",100,0,10,150,0,1.5);
  Xi_KpEOverPBCAL=new TH2F("Xi_KpEOverPBCAL","E/p vs p",100,0,10,150,0,1.5); 
  
  gDirectory->mkdir("purity")->cd();
  
  Xi_KpDenom=new TH2F("Xi_KpDenom","P vs #theta",100,0,100,100,0.05,10.05); 
  Xi_KpNumer_Kp=new TH2F("Xi_KpNumer_Kp","P vs #theta",100,0,100,100,0.05,10.05);
  Xi_KpNumer_Pip=new TH2F("Xi_KpNumer_Pip","P vs #theta",100,0,100,100,0.05,10.05);
  Xi_KpNumer_P=new TH2F("Xi_KpNumer_P","P vs #theta",100,0,100,100,0.05,10.05);
  Xi_KpNumer_Ep=new TH2F("Xi_KpNumer_Ep","P vs #theta",100,0,100,100,0.05,10.05);
 
  gDirectory->cd("../../../");

  CL_CUT=0.01;
  gPARMS->SetDefaultParameter("CASCADESTUDIES:CL_CUT",CL_CUT);

  



	// This is called once at program startup. 

	return NOERROR;
}

//------------------
// brun
//------------------
jerror_t JEventProcessor_CascadeStudies::brun(JEventLoop *eventLoop, int32_t runnumber)
{
	// This is called whenever the run number changes
	return NOERROR;
}

//------------------
// evnt
//------------------
jerror_t JEventProcessor_CascadeStudies::evnt(JEventLoop *loop, uint64_t eventnumber)
{
  vector<const DChargedTrack*>ctracks;
  loop->Get(ctracks); 
  if (ctracks.size()!=5) return VALUE_OUT_OF_RANGE;

  vector<const DParticleID *> pid_algorithms;
  loop->Get(pid_algorithms);
  if(pid_algorithms.size()<1){
    _DBG_<<"Unable to get a DParticleID object! NO PID will be done!"<<endl;
    return RESOURCE_UNAVAILABLE;
  }
  const DParticleID *pid_algorithm = pid_algorithms[0];
  
  vector<const DBeamPhoton*>beamphotons;
  loop->Get(beamphotons);
  
  vector<const DMCThrown *>throwns;
  loop->Get(throwns);

  japp->WriteLock("myCascadelock");

  vector<unsigned int>neglist;
  vector<unsigned int>poslist;
  for (int i=0;i<5;i++){
    if (ctracks[i]->Get_BestTrackingFOM()->charge()<0){
      neglist.push_back(i);
    }
    else poslist.push_back(i);
  }
  if (neglist.size()==2){     
    DAnalysisUtilities *dAnalysisUtilities=new DAnalysisUtilities(loop);
    DKinFitUtils_GlueX *dKinFitUtils = new DKinFitUtils_GlueX(loop);
    DKinFitter *dKinFitter = new DKinFitter(dKinFitUtils);   
    dKinFitUtils->Reset_NewEvent();
    dKinFitter->Reset_NewEvent();
    
    unsigned int posindexes[6][3];
    unsigned int negindexes[2][2];
       
    posindexes[0][0]=poslist[0];
    posindexes[0][1]=poslist[1];
    posindexes[0][2]=poslist[2]; 
    posindexes[1][0]=poslist[0];
    posindexes[1][1]=poslist[2];
    posindexes[1][2]=poslist[1];
    posindexes[2][0]=poslist[1];
    posindexes[2][1]=poslist[2];
    posindexes[2][2]=poslist[0];
    posindexes[3][0]=poslist[1];
    posindexes[3][1]=poslist[0];
    posindexes[3][2]=poslist[2]; 
    posindexes[4][0]=poslist[2];
    posindexes[4][1]=poslist[0];
    posindexes[4][2]=poslist[1];
    posindexes[5][0]=poslist[2];
    posindexes[5][1]=poslist[1];
    posindexes[5][2]=poslist[0];
    
    negindexes[0][0]=neglist[0];
    negindexes[0][1]=neglist[1];  
    negindexes[1][0]=neglist[1];
    negindexes[1][1]=neglist[0]; 
  
    for (unsigned int i=0;i<beamphotons.size();i++){    
      for (int j=0;j<2;j++){
	const DChargedTrackHypothesis *piminus_hyp=ctracks[negindexes[j][0]]->Get_Hypothesis(PiMinus);
	if (piminus_hyp==NULL) continue;
	const DChargedTrackHypothesis *piminus2_hyp=ctracks[negindexes[j][1]]->Get_Hypothesis(PiMinus);
	if (piminus2_hyp==NULL) continue;
	
	const DTrackTimeBased *piminus_track=piminus_hyp->Get_TrackTimeBased();
	if (piminus_track==NULL) continue;
	const DTrackTimeBased *piminus2_track=piminus2_hyp->Get_TrackTimeBased();
	if (piminus2_hyp==NULL) continue;
	
	for (int k=0;k<6;k++){
	  const DChargedTrackHypothesis *proton_hyp=ctracks[posindexes[k][0]]->Get_Hypothesis(Proton);
	  if (proton_hyp==NULL) continue;
	  const DChargedTrackHypothesis *kplus_hyp1=ctracks[posindexes[k][1]]->Get_Hypothesis(KPlus);
	  if (kplus_hyp1==NULL) continue;
	  const DChargedTrackHypothesis *kplus_hyp2=ctracks[posindexes[k][2]]->Get_Hypothesis(KPlus);
	  if (kplus_hyp2==NULL) continue;
	   
	  const DTrackTimeBased *proton_track=proton_hyp->Get_TrackTimeBased();
	  if (proton_track==NULL) continue;
	  const DTrackTimeBased *kplus_track1=kplus_hyp1->Get_TrackTimeBased();
	  if (kplus_track1==NULL) continue;
	  const DTrackTimeBased *kplus_track2=kplus_hyp2->Get_TrackTimeBased();
	  if (kplus_track2==NULL) continue;
	  
	  double t0_rf=kplus_hyp1->t0();
	  double dt_st=beamphotons[i]->time()-t0_rf; 
	  
	  double weight=1.;
	  bool got_beam_photon=false;
	  if (fabs(dt_st)>6.&&fabs(dt_st)<18.){
	    weight=-1./6.;
	    got_beam_photon=true;
	   }
	  if (fabs(dt_st)<2.){
	    got_beam_photon=true;
	  }
	  if (got_beam_photon==false) continue;

	  //--------------------------------
	  // Kinematic fit
	  //-------------------------------- 
	  dKinFitter->Reset_NewFit();
      
	  set<shared_ptr<DKinFitParticle>> InitialParticles, FinalParticles;
	  
	  shared_ptr<DKinFitParticle>myBeam=dKinFitUtils->Make_BeamParticle(beamphotons[i]);
	  shared_ptr<DKinFitParticle>myTarget=dKinFitUtils->Make_TargetParticle(Proton);
	  
	  InitialParticles.insert(myBeam);  
	  InitialParticles.insert(myTarget);
	   
	  shared_ptr<DKinFitParticle>myKPlus1=dKinFitUtils->Make_DetectedParticle(kplus_track1);
	  shared_ptr<DKinFitParticle>myProton=dKinFitUtils->Make_DetectedParticle(proton_track);
	  shared_ptr<DKinFitParticle>myPiMinus=dKinFitUtils->Make_DetectedParticle(piminus_track);  
	  shared_ptr<DKinFitParticle>myKPlus2=dKinFitUtils->Make_DetectedParticle(kplus_track2);
	  shared_ptr<DKinFitParticle>myPiMinus2=dKinFitUtils->Make_DetectedParticle(piminus2_track); 
	  
	  FinalParticles.insert(myKPlus1);
	  FinalParticles.insert(myProton);
	  FinalParticles.insert(myPiMinus);
	  FinalParticles.insert(myKPlus2);
	  FinalParticles.insert(myPiMinus2);
	   
	  //MAKE LAMBDA
	  set<shared_ptr<DKinFitParticle>> locFromInitialState2, locFromFinalState2;
	  locFromFinalState2.insert(myPiMinus);  
	  locFromFinalState2.insert(myProton);
	  shared_ptr<DKinFitParticle> myLambda = dKinFitUtils->Make_DecayingParticle(AntiLambda, locFromInitialState2, locFromFinalState2);
	  
	  // Decay Vertex constraint
	  set<shared_ptr<DKinFitParticle>> locFullConstrainParticles2, locNoConstrainParticles2;
	  locFullConstrainParticles2.insert(myPiMinus);  
	  locFullConstrainParticles2.insert(myProton);
	  locNoConstrainParticles2.insert(myLambda);
	  
	  DVector3 second_vertex2;
	  dAnalysisUtilities->Calc_DOCAVertex(myPiMinus.get(),myProton.get(),second_vertex2);
		  
	  shared_ptr<DKinFitConstraint_Vertex> locDecayVertexConstraint2 = dKinFitUtils->Make_VertexConstraint(locFullConstrainParticles2, locNoConstrainParticles2,second_vertex2);
	  dKinFitter->Add_Constraint(locDecayVertexConstraint2);
	  
	  // make constraint
	  shared_ptr<DKinFitConstraint_P4>locP4Constraint = dKinFitUtils->Make_P4Constraint(InitialParticles, FinalParticles);
	  // set constraint
	  dKinFitter->Add_Constraint(locP4Constraint);
	  
	  // PERFORM THE KINEMATIC FIT
	  dKinFitter->Fit_Reaction();
	  
	  // GET THE FIT RESULTS
	  double locConfidenceLevel = dKinFitter->Get_ConfidenceLevel();
	  CascadeConfidenceLevel->Fill(locConfidenceLevel);
	  CascadeChiSq->Fill(dKinFitter->Get_ChiSq()/double(dKinFitter->Get_NDF()));
      
	  if (locConfidenceLevel>CL_CUT){
	    if (CascadeAnalysis(dKinFitter,
				ctracks[posindexes[k][0]] /* proton guess */,
				ctracks[posindexes[k][1]] /* kplus guess */,
				ctracks[posindexes[k][2]] /* kplus guess */,
				ctracks[negindexes[j][0]], // pi-
				ctracks[negindexes[j][1]],weight)){
	      FillPIDHistos(t0_rf,kplus_hyp1,kplus_track1,kplus_hyp2,
			    kplus_track2,proton_hyp,proton_track,piminus_hyp,
			    piminus_track,piminus2_hyp,piminus2_track,
			    pid_algorithm);

	    }
	  }
	  
	}
      }
    }

    delete dKinFitter;
    delete dKinFitUtils;
    delete dAnalysisUtilities;
  }
	
  japp->Unlock("myCascadelock");
  
  return NOERROR;
}

//------------------
// erun
//------------------
jerror_t JEventProcessor_CascadeStudies::erun(void)
{
	// This is called whenever the run number changes, before it is
	// changed to give you a chance to clean up before processing
	// events from the next run number.
	return NOERROR;
}

//------------------
// fini
//------------------
jerror_t JEventProcessor_CascadeStudies::fini(void)
{
	// Called before program exit after event processing is finished.
	return NOERROR;
}


bool JEventProcessor_CascadeStudies::CascadeAnalysis(DKinFitter *dKinFitter,
							const DChargedTrack *proton_hyp,
							const DChargedTrack *kplus_hyp1, 
							const DChargedTrack *kplus_hyp2,
								
							const DChargedTrack *piminus_hyp,
							const DChargedTrack *piminus_hyp2,				   
							double weight){
  DLorentzVector beam_kf;
  DLorentzVector proton_kf;
  vector<DLorentzVector>kps_kf,pims_kf;
  
  set<shared_ptr<DKinFitParticle>>myParticles=dKinFitter->Get_KinFitParticles();
  set<shared_ptr<DKinFitParticle>>::iterator locParticleIterator=myParticles.begin();
  for(; locParticleIterator != myParticles.end(); ++locParticleIterator){
    if ((*locParticleIterator)->Get_KinFitParticleType()==d_BeamParticle){
      beam_kf=(*locParticleIterator)->Get_P4();
    }
    else if ((*locParticleIterator)->Get_KinFitParticleType()==d_DetectedParticle){
      switch ((*locParticleIterator)->Get_PID()){
      case 2212:  // Proton
	proton_kf=((*locParticleIterator)->Get_P4());
	break;
      case 321: // K+
	kps_kf.push_back((*locParticleIterator)->Get_P4());
	break;
      case -211: // pi-
	pims_kf.push_back((*locParticleIterator)->Get_P4());
	break;
      default:
	break;
      }
    }
  }
  DLorentzVector ppim1_kf=proton_kf+pims_kf[0]; 
  DLorentzVector ppim2_kf=proton_kf+pims_kf[1];
  DLorentzVector p2pim_kf=ppim1_kf+pims_kf[1];
  P2PimMass->Fill(p2pim_kf.M(),weight);
  PPimMass->Fill(ppim1_kf.M(),weight);
  PPimMass->Fill(ppim2_kf.M(),weight);
  if (LAMBDA_CUT(ppim1_kf.M()) || LAMBDA_CUT(ppim2_kf.M())){
    LambdaPimMass->Fill(p2pim_kf.M(),weight);
    if (weight>0 && CASCADE_CUT(p2pim_kf.M())){
      for (int i=0;i<2;i++){
	double kplus_theta=kps_kf[i].Theta()*180./M_PI;
	double kplus_momentum=kps_kf[i].P();
	Xi_KpDenom->Fill(kplus_theta,kplus_momentum);

	const DTrackTimeBased *best_fom_track=NULL;
	if (i==0) best_fom_track=kplus_hyp1->Get_BestFOM()->Get_TrackTimeBased();
	else best_fom_track=kplus_hyp2->Get_BestFOM()->Get_TrackTimeBased();
	switch(best_fom_track->PID()){
	case PiPlus:
	  Xi_KpNumer_Pip->Fill(kplus_theta,kplus_momentum);
	  break;
	case KPlus:
	  Xi_KpNumer_Kp->Fill(kplus_theta,kplus_momentum);
	  break;
	case Proton:
	  Xi_KpNumer_P->Fill(kplus_theta,kplus_momentum);
	  break;
	case Positron:
	  Xi_KpNumer_Ep->Fill(kplus_theta,kplus_momentum);
	  break;
	default:
	  break;
	}
      }

      return true;
    }
  }

  return false;
}

void JEventProcessor_CascadeStudies::FillPIDHistos(double t0_rf,
						   const DChargedTrackHypothesis *proton_hyp,
						   const DTrackTimeBased *proton_track,
						   const DChargedTrackHypothesis *kplus_hyp1,
						   const DTrackTimeBased *kplus_track1,  
						   const DChargedTrackHypothesis *kplus_hyp2,
						   const DTrackTimeBased *kplus_track2, 
						   const DChargedTrackHypothesis *piminus_hyp1,
						   const DTrackTimeBased *piminus_track1, 
						   const DChargedTrackHypothesis *piminus_hyp2, 
						   const DTrackTimeBased *piminus_track2,
						   const DParticleID *pid_algorithm){  
  double p=kplus_track1->momentum().Mag();
  double dEdxAmp=1e6*kplus_track1->ddEdx_CDC_amp;	 
  double dEdx_fdc=1e6*kplus_track1->ddEdx_FDC;
  if (dEdxAmp>0){
    Xi_KpdEdxCDC->Fill(p/ParticleMass(KPlus),dEdxAmp);
  }
  if (dEdx_fdc>0){
    Xi_KpdEdxFDC->Fill(p/ParticleMass(KPlus),dEdx_fdc);
  }
  
  shared_ptr<const DBCALShowerMatchParams>bcalparms=kplus_hyp1->Get_BCALShowerMatchParams(); 
  shared_ptr<const DFCALShowerMatchParams>fcalparms=kplus_hyp1->Get_FCALShowerMatchParams();	  
  shared_ptr<const DTOFHitMatchParams>tofparms=kplus_hyp1->Get_TOFHitMatchParams();
	
  if (bcalparms!=NULL){
    double dt_bcal=bcalparms->dBCALShower->t-bcalparms->dFlightTime-t0_rf;
    Xi_KpBCALdt->Fill(p,dt_bcal);  
  }
  if (tofparms!=NULL){ 
    double dt=tofparms->dHitTime-tofparms->dFlightTime-t0_rf;
    Xi_KpTOFdt->Fill(p,dt);
  } 
  if (fcalparms!=NULL){
    double dt_fcal=fcalparms->dFCALShower->getTime()
      -fcalparms->dFlightTime-t0_rf;
    Xi_KpFCALdt->Fill(p,dt_fcal);  

    double E_over_p=fcalparms->dFCALShower->getEnergy()/p;
    Xi_KpEOverPFCAL->Fill(p,E_over_p);

  }

  p=kplus_track2->momentum().Mag();
  dEdxAmp=1e6*kplus_track2->ddEdx_CDC_amp;	 
  dEdx_fdc=1e6*kplus_track2->ddEdx_FDC;
  if (dEdxAmp>0){
    Xi_KpdEdxCDC->Fill(p/ParticleMass(KPlus),dEdxAmp);
  }
  if (dEdx_fdc>0){
    Xi_KpdEdxFDC->Fill(p/ParticleMass(KPlus),dEdx_fdc);
  }
  
  bcalparms=kplus_hyp2->Get_BCALShowerMatchParams(); 
  fcalparms=kplus_hyp2->Get_FCALShowerMatchParams();	  
  tofparms=kplus_hyp2->Get_TOFHitMatchParams();
	
  if (bcalparms!=NULL){
    double dt_bcal=bcalparms->dBCALShower->t-bcalparms->dFlightTime-t0_rf;
    Xi_KpBCALdt->Fill(p,dt_bcal);  
  }
  if (tofparms!=NULL){ 
    double dt=tofparms->dHitTime-tofparms->dFlightTime-t0_rf;
    Xi_KpTOFdt->Fill(p,dt);
  } 
  if (fcalparms!=NULL){
    double dt_fcal=fcalparms->dFCALShower->getTime()
      -fcalparms->dFlightTime-t0_rf;
    Xi_KpFCALdt->Fill(p,dt_fcal);  

    double E_over_p=fcalparms->dFCALShower->getEnergy()/p;
    Xi_KpEOverPFCAL->Fill(p,E_over_p);

  }  

}