// $Id$
//
//    File: JEventProcessor_LambdaStudies.cc
// Created: Fri Jun 28 10:04:22 EDT 2019
// Creator: staylor (on Linux ifarm1401.jlab.org 3.10.0-327.el7.x86_64 x86_64)
//

#include "JEventProcessor_LambdaStudies.h"
using namespace jana;

#define LAMBDA_CUT 0.012

// 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_LambdaStudies());
}
} // "C"


//------------------
// JEventProcessor_LambdaStudies (Constructor)
//------------------
JEventProcessor_LambdaStudies::JEventProcessor_LambdaStudies()
{

}

//------------------
// ~JEventProcessor_LambdaStudies (Destructor)
//------------------
JEventProcessor_LambdaStudies::~JEventProcessor_LambdaStudies()
{
 
}

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

  gDirectory->mkdir("LambdaAnalysis")->cd();
  LambdaConfidenceLevel=new TH1F("LambdaConfidenceLevel","CL",100,0,1);
  LambdaChiSq=new TH1F("LambdaChisSq","#chi^{2}/ndf",1000,0,100);
  PimPMass= new TH1F("PimPMass","p#pi- mass",2000,1.,3.);

  gDirectory->mkdir("PID")->cd();

  KpTOFdt_Lambda=new TH2F("KpTOFdt_Lambda","#deltat(TOF)",100,0,10,200,-5,5);	
  KpFCALdt_Lambda=new TH2F("KpFCALdt_Lambda","#deltat(FCAL)",100,0,10,200,-5,5);
  
  KpTOFdtPull=new TH2F("KpTOFdtPull","#deltat(TOF)/#sigmat(TOF)",100,0,10,200,-5,5);	
  KpFCALdtPull=new TH2F("KpFCALdtPull","#deltat(FCAL)/#sigmat(FCAL)",100,0,10,200,-5,5);
  KpBCALdtPull=new TH2F("KpBCALdtPull","#deltat(BCAL)/#sigmat(BCAL)",100,0,10,200,-5,5);
  KpTOFdtCL=new TH1F("KpTOFdtCL","CL for TOF",100,0,1);
  KpFCALdtCL=new TH1F("KpFCALdtCL","CL for FCAL",100,0,1);
  KpBCALdtCL=new TH1F("KpBCALdtCL","CL for BCAL",100,0,1);
  
  KpBCALdt_Lambda=new TH2F("KpBCALdt_Lambda","#deltat(BCAL)",100,0,10,200,-5,5);
  KpdEdxCDC_Lambda=new TH2F("KpdEdxCDC_Lambda","dEdx vs p/M, K+ cands.",1000,0,10,500,0,25.0);	
  KpdEdxFDC_Lambda=new TH2F("KpdEdxFDC_Lambda","dEdx vs p/M, K+ cands.",1000,0,10,500,0,25.0);
  KpdEdxCDCPull_Lambda=new TH2F("KpdEdxCDCPull_Lambda","#delta(dE/dx)/#sigma(dE/dx) vs p/M, K+ cands.",1000,0,10,100,-5,5.0);
  KpdEdxCDCCL_Lambda=new TH1F("KpdEdxCDCCL_Lambda","CL for CDC dEdx",100,0,1);
  KpdEdxFDCPull_Lambda=new TH2F("KpdEdxFDCPull_Lambda","#delta(dE/dx)/#sigma(dE/dx) vs p/M, K+ cands.",1000,0,10,100,-5,5.0);
  KpdEdxFDCCL_Lambda=new TH1F("KpdEdxFDCCL_Lambda","CL for FDC dEdx",100,0,1);
  
  PTOFdt_Lambda=new TH2F("PTOFdt_Lambda","#deltat(TOF)",100,0,10,200,-5,5);	
  PFCALdt_Lambda=new TH2F("PFCALdt_Lambda","#deltat(FCAL)",100,0,10,200,-5,5);
  PBCALdt_Lambda=new TH2F("PBCALdt_Lambda","#deltat(BCAL)",100,0,10,200,-5,5);
  PdEdxCDC_Lambda=new TH2F("PdEdxCDC_Lambda","dEdx vs p/M, proton cands.",1000,0,10,500,0,25.0);	
  PdEdxFDC_Lambda=new TH2F("PdEdxFDC_Lambda","dEdx vs p/M, proton cands.",1000,0,10,500,0,25.0);
  
  gDirectory->mkdir("purity")->cd(); 

  KpDenom_Lambda=new TH2F("KpDenom_Lambda","P vs #theta",40,0,40,100,0.05,10.05); 
  KpNumer_Kp_Lambda=new TH2F("KpNumer_Kp_Lambda","P vs #theta",40,0,40,100,0.05,10.05);
  KpNumer_Pip_Lambda=new TH2F("KpNumer_Pip_Lambda","P vs #theta",40,0,40,100,0.05,10.05);
  KpNumer_P_Lambda=new TH2F("KpNumer_P_Lambda","P vs #theta",40,0,40,100,0.05,10.05);
  KpNumer_Ep_Lambda=new TH2F("KpNumer_Ep_Lambda","P vs #theta",40,0,40,100,0.05,10.05);
  
  PimDenom_Lambda=new TH2F("PimDenom_Lambda","P vs #theta",100,0,100,100,0.05,10.05); 
  PimNumer_Km_Lambda=new TH2F("PimNumer_Km_Lambda","P vs #theta",100,0,100,100,0.05,10.05);
  PimNumer_Pim_Lambda=new TH2F("PimNumer_Pim_Lambda","P vs #theta",100,0,100,100,0.05,10.05);
  PimNumer_aP_Lambda=new TH2F("PimNumer_aP_Lambda","P vs #theta",100,0,100,100,0.05,10.05);
  PimNumer_Em_Lambda=new TH2F("PimNumer_Em_Lambda","P vs #theta",100,0,100,100,0.05,10.05);
  
  PDenom_Lambda=new TH2F("PDenom_Lambda","P vs #theta",80,0,80,100,0.05,10.05); 
  PNumer_Kp_Lambda=new TH2F("PNumer_Kp_Lambda","P vs #theta",80,0,80,100,0.05,10.05);
  PNumer_Pip_Lambda=new TH2F("PNumer_Pip_Lambda","P vs #theta",80,0,80,100,0.05,10.05);
  PNumer_P_Lambda=new TH2F("PNumer_P_Lambda","P vs #theta",80,0,80,100,0.05,10.05);
  PNumer_Ep_Lambda=new TH2F("PNumer_Ep_Lambda","P vs #theta",80,0,80,100,0.05,10.05);
  
  gDirectory->cd("../../../");

  CL_CUT=0.01;
  gPARMS->SetDefaultParameter("LAMBDASTUDIES:CL_CUT",CL_CUT);
 
  return NOERROR;
}

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

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

  vector<const DBeamPhoton*>beamphotons;
  loop->Get(beamphotons); 
  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];
  
  japp->WriteLock("myLambdalock");

  unsigned int neg_index=0;
  unsigned int num_neg=0;
  for (int i=0;i<3;i++){
    if (ctracks[i]->Get_BestTrackingFOM()->charge()<0){
      num_neg++;
      neg_index=i;
    }
  }
  if (num_neg==1){
    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[2][2];
    if (neg_index==0){
      posindexes[0][0]=1;
      posindexes[0][1]=2;
      posindexes[1][0]=2;
      posindexes[1][1]=1;
    } 
    if (neg_index==1){
      posindexes[0][0]=0;
      posindexes[0][1]=2;
      posindexes[1][0]=2;
      posindexes[1][1]=0;
    }
    if (neg_index==2){
      posindexes[0][0]=1;
      posindexes[0][1]=0;
      posindexes[1][0]=0;
      posindexes[1][1]=1;
    }
      
    const DChargedTrackHypothesis *piminus_hyp=ctracks[neg_index]->Get_Hypothesis(PiMinus);
    if (piminus_hyp!=NULL){
      const DTrackTimeBased *piminus_track=piminus_hyp->Get_TrackTimeBased();
      if (piminus_track!=NULL){
	for (unsigned int i=0;i<beamphotons.size();i++){
	  for (int k=0;k<2;k++){
	    const DChargedTrackHypothesis *kplus_hyp=ctracks[posindexes[k][0]]->Get_Hypothesis(KPlus);
	    if (kplus_hyp==NULL) continue;
	    const DChargedTrackHypothesis *proton_hyp=ctracks[posindexes[k][1]]->Get_Hypothesis(Proton);
	    if (proton_hyp==NULL) continue;
	    
	    const DTrackTimeBased *kplus_track=kplus_hyp->Get_TrackTimeBased();
	    if (kplus_track==NULL) continue;
	    
	    const DTrackTimeBased *proton_track=proton_hyp->Get_TrackTimeBased();
	    if (proton_track==NULL) continue;

	    double t0_rf=kplus_hyp->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>myProton=dKinFitUtils->Make_DetectedParticle(proton_track);
	    shared_ptr<DKinFitParticle>myKplus=dKinFitUtils->Make_DetectedParticle(kplus_track);
	    shared_ptr<DKinFitParticle>myPiminus=dKinFitUtils->Make_DetectedParticle(piminus_track);  
	    
	    FinalParticles.insert(myProton);
	    FinalParticles.insert(myKplus);
	    FinalParticles.insert(myPiminus);
	    
	    // make constraint
	    shared_ptr<DKinFitConstraint_P4>locP4Constraint = dKinFitUtils->Make_P4Constraint(InitialParticles, FinalParticles);
	    // set constraint
	    dKinFitter->Add_Constraint(locP4Constraint);
	      	
	    //MAKE ppi- pair
	    set<shared_ptr<DKinFitParticle>> locFromInitialState, locFromFinalState;
	    locFromFinalState.insert(myPiminus);  
	    locFromFinalState.insert(myProton);
	    shared_ptr<DKinFitParticle> myLambda = dKinFitUtils->Make_DecayingParticle(Lambda, locFromInitialState, locFromFinalState);
	    
	    // Decay Vertex constraint
	    set<shared_ptr<DKinFitParticle>> locFullConstrainParticles, locNoConstrainParticles;
	    locFullConstrainParticles.insert(myPiminus);  
	    locFullConstrainParticles.insert(myProton);
	    locNoConstrainParticles.insert(myLambda);
		
	    DVector3 second_vertex;
	    dAnalysisUtilities->Calc_DOCAVertex(myPiminus.get(),myProton.get(),second_vertex);
		
	    shared_ptr<DKinFitConstraint_Vertex> locSecondaryVertexConstraint = dKinFitUtils->Make_VertexConstraint(locFullConstrainParticles, locNoConstrainParticles,second_vertex);
	    dKinFitter->Add_Constraint(locSecondaryVertexConstraint);
	      
	    // PERFORM THE KINEMATIC FIT
	    dKinFitter->Fit_Reaction();
	    
	    // GET THE FIT RESULTS
	    double locConfidenceLevel = dKinFitter->Get_ConfidenceLevel();
	    LambdaConfidenceLevel->Fill(locConfidenceLevel);
	    LambdaChiSq->Fill(dKinFitter->Get_ChiSq()/double(dKinFitter->Get_NDF()));
	    if (locConfidenceLevel>CL_CUT){
	      if (LambdaAnalysis(dKinFitter,
				 ctracks[posindexes[k][1]],// proton
				 ctracks[posindexes[k][0]], //kplus
				 ctracks[neg_index], // piminus
				 weight)){
		FillPIDHistos(t0_rf,proton_hyp,proton_track,kplus_hyp,
			      kplus_track,piminus_hyp,piminus_track,
			      pid_algorithm);
		
	      }
	      
	    }
	    
	  } // loop over 2 possible positive track combinations
	} // loop over beam photons
      } // check for piminus track
    }
    
    delete dAnalysisUtilities;
    delete dKinFitter;
    delete dKinFitUtils;
  }
  
  japp->Unlock("myLambdalock");

  return NOERROR;
}

//------------------
// erun
//------------------
jerror_t JEventProcessor_LambdaStudies::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_LambdaStudies::fini(void)
{
	// Called before program exit after event processing is finished.
	return NOERROR;
}

void JEventProcessor_LambdaStudies::FillPIDHistos(double t0_rf,
						  const DChargedTrackHypothesis *proton_hyp,
						  const DTrackTimeBased *proton_track,
						  const DChargedTrackHypothesis *kplus_hyp,
						  const DTrackTimeBased *kplus_track,
						  const DChargedTrackHypothesis *piminus_hyp,
						  const DTrackTimeBased *piminus_track,
						  const DParticleID *pid_algorithm){
  double p=kplus_track->momentum().Mag();
  double beta=p/kplus_track->energy();
  double gammabeta=p/ParticleMass(KPlus);
  double dEdxAmp=1e6*kplus_track->ddEdx_CDC_amp;	 
  double dEdx_fdc=1e6*kplus_track->ddEdx_FDC;
  if (dEdxAmp>0){
    double dEdxMean=0.,dEdxSigma=0.;		      
    pid_algorithm->GetdEdxMean_CDC(beta,0,dEdxMean,KPlus);
    pid_algorithm->GetdEdxSigma_CDC(beta,0,dEdxSigma,KPlus);
    double dEdxDiff=kplus_track->ddEdx_CDC_amp-dEdxMean;
    double dEdxPull=dEdxDiff/dEdxSigma;
    double dEdxChiSq=dEdxPull*dEdxPull;
    
    KpdEdxCDC_Lambda->Fill(gammabeta,dEdxAmp);
    KpdEdxCDCPull_Lambda->Fill(gammabeta,dEdxPull);
    KpdEdxCDCCL_Lambda->Fill(TMath::Prob(dEdxChiSq,1.));
  }
  if (dEdx_fdc>0){
    double dEdxMean=0.,dEdxSigma=0.;		      
    pid_algorithm->GetdEdxMean_FDC(beta,0,dEdxMean,KPlus);
    pid_algorithm->GetdEdxSigma_FDC(beta,0,dEdxSigma,KPlus);
    double dEdxDiff=kplus_track->ddEdx_FDC-dEdxMean;
    double dEdxPull=dEdxDiff/dEdxSigma;
    double dEdxChiSq=dEdxPull*dEdxPull;
	  
    KpdEdxFDCPull_Lambda->Fill(gammabeta,dEdxPull);
    KpdEdxFDCCL_Lambda->Fill(TMath::Prob(dEdxChiSq,1.));
    KpdEdxFDC_Lambda->Fill(gammabeta,dEdx_fdc);
  }
  
  shared_ptr<const DBCALShowerMatchParams>bcalparms=kplus_hyp->Get_BCALShowerMatchParams(); 
  shared_ptr<const DFCALShowerMatchParams>fcalparms=kplus_hyp->Get_FCALShowerMatchParams();	  
  shared_ptr<const DTOFHitMatchParams>tofparms=kplus_hyp->Get_TOFHitMatchParams();
  if (bcalparms!=NULL){
    double dt_bcal=bcalparms->dBCALShower->t-bcalparms->dFlightTime-t0_rf;
    KpBCALdt_Lambda->Fill(p,dt_bcal);  
	  
    double var=pid_algorithm->GetTimeVariance(SYS_BCAL,KPlus,p);
    double pull=dt_bcal/sqrt(var);
    double chisq=dt_bcal*dt_bcal/var;
    KpBCALdtPull->Fill(p,pull);
    KpBCALdtCL->Fill(TMath::Prob(chisq,1.));
  }
  if (tofparms!=NULL){ 
    double dt=tofparms->dHitTime-tofparms->dFlightTime-t0_rf;
    KpTOFdt_Lambda->Fill(p,dt);
    
    double var=pid_algorithm->GetTimeVariance(SYS_TOF,KPlus,p);
    double pull=dt/sqrt(var);
    double chisq=dt*dt/var;
    KpTOFdtPull->Fill(p,pull);
    KpTOFdtCL->Fill(TMath::Prob(chisq,1.));
	} 
  if (fcalparms!=NULL){
    double dt_fcal=fcalparms->dFCALShower->getTime()-fcalparms->dFlightTime-t0_rf;
    KpFCALdt_Lambda->Fill(p,dt_fcal);
    
    double var=pid_algorithm->GetTimeVariance(SYS_FCAL,KPlus,p);
    double pull=dt_fcal/sqrt(var);
    double chisq=dt_fcal*dt_fcal/var;
    KpFCALdtPull->Fill(p,pull);
    KpFCALdtCL->Fill(TMath::Prob(chisq,1.));
  }
  
  p=proton_track->momentum().Mag();
  dEdxAmp=1e6*proton_track->ddEdx_CDC_amp;	 
  dEdx_fdc=1e6*proton_track->ddEdx_FDC;
  if (dEdxAmp>0){
    PdEdxCDC_Lambda->Fill(p/ParticleMass(Proton),dEdxAmp);
  }
  if (dEdx_fdc>0){
    PdEdxFDC_Lambda->Fill(p/ParticleMass(Proton),dEdx_fdc);
  }
  
  shared_ptr<const DBCALShowerMatchParams>bcalparms2=kplus_hyp->Get_BCALShowerMatchParams(); 
  shared_ptr<const DFCALShowerMatchParams>fcalparms2=kplus_hyp->Get_FCALShowerMatchParams();	  
  shared_ptr<const DTOFHitMatchParams>tofparms2=kplus_hyp->Get_TOFHitMatchParams();
	
  if (bcalparms2!=NULL){
    double dt_bcal=bcalparms2->dBCALShower->t-bcalparms2->dFlightTime-t0_rf;
    PBCALdt_Lambda->Fill(p,dt_bcal);  
  }
  if (tofparms2!=NULL){ 
    double dt=tofparms2->dHitTime-tofparms2->dFlightTime-t0_rf;
    PTOFdt_Lambda->Fill(p,dt);
  } 
  if (fcalparms2!=NULL){
    double dt_fcal=fcalparms2->dFCALShower->getTime()
      -fcalparms2->dFlightTime-t0_rf;
    PFCALdt_Lambda->Fill(p,dt_fcal);
  }
}

bool JEventProcessor_LambdaStudies::LambdaAnalysis(DKinFitter *dKinFitter,
						   const DChargedTrack *proton_hyp,
						   const DChargedTrack *kplus_hyp,
						   const DChargedTrack *piminus_hyp,
						   double weight){
  DLorentzVector beam_kf,gamma_kf;
  DLorentzVector proton_kf,kp_kf,pim_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+
	kp_kf=(*locParticleIterator)->Get_P4();	
	break;
      case -211: // pi-
	pim_kf=(*locParticleIterator)->Get_P4();
	break;	  
      default:
	break;
      }
    }
  }

  DLorentzVector pPim_kf=proton_kf+pim_kf;
  PimPMass->Fill(pPim_kf.M(),weight);
  
  if (weight>0. && fabs(pPim_kf.M()-ParticleMass(Lambda))<LAMBDA_CUT){
    double kp_theta=kp_kf.Theta()*180./M_PI;
    double kp_momentum=kp_kf.P();
    KpDenom_Lambda->Fill(kp_theta,kp_momentum);

    const DTrackTimeBased *best_fom_track=kplus_hyp->Get_BestFOM()->Get_TrackTimeBased();
    switch(best_fom_track->PID()){
    case PiPlus:
      KpNumer_Pip_Lambda->Fill(kp_theta,kp_momentum);
      break;
    case KPlus:
      KpNumer_Kp_Lambda->Fill(kp_theta,kp_momentum);
      break;
    case Proton:
      KpNumer_P_Lambda->Fill(kp_theta,kp_momentum);
      break;
    case Positron:
      KpNumer_Ep_Lambda->Fill(kp_theta,kp_momentum);
      break;
    default:
      break;
    }

    double pim_theta=pim_kf.Theta()*180./M_PI;
    double pim_momentum=pim_kf.P();
    PimDenom_Lambda->Fill(pim_theta,pim_momentum);
    best_fom_track=piminus_hyp->Get_BestFOM()->Get_TrackTimeBased();
    switch(best_fom_track->PID()){
    case PiMinus:
      PimNumer_Pim_Lambda->Fill(pim_theta,pim_momentum);
      break;
    case KMinus:
      PimNumer_Km_Lambda->Fill(pim_theta,pim_momentum);
      break;
    case AntiProton:
      PimNumer_aP_Lambda->Fill(pim_theta,pim_momentum);
      break;
    case Electron:
      PimNumer_Em_Lambda->Fill(pim_theta,pim_momentum);
      break;
    default:
      break;
    }
    
    double proton_theta=proton_kf.Theta()*180./M_PI;
    double proton_momentum=proton_kf.P();
    PDenom_Lambda->Fill(proton_theta,proton_momentum);
    best_fom_track=proton_hyp->Get_BestFOM()->Get_TrackTimeBased();
    switch(best_fom_track->PID()){
    case PiPlus:
      PNumer_Pip_Lambda->Fill(proton_theta,proton_momentum);
      break;
    case KPlus:
      PNumer_Kp_Lambda->Fill(proton_theta,proton_momentum);
      break;
    case Proton:
      PNumer_P_Lambda->Fill(proton_theta,proton_momentum);
      break;
    case Positron:
      PNumer_Ep_Lambda->Fill(proton_theta,proton_momentum);
      break;
    default:
      break;
    }


    return true;
  }
  return false;
}