// $Id$
//
//    File: JEventProcessor_Primakoff.cc
// Created: Fri Sep  5 04:31:44 PM EDT 2025
// Creator: staylor (on Linux ifarm2401.jlab.org 5.14.0-570.33.2.el9_6.x86_64 x86_64)
//

/// For more information on the syntax changes between JANA1 and JANA2, visit: https://jeffersonlab.github.io/JANA2/#/jana1to2/jana1-to-jana2

#include "JEventProcessor_Primakoff.h"
#include <vector>
#include <map>

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


//------------------
// JEventProcessor_Primakoff (Constructor)
//------------------
JEventProcessor_Primakoff::JEventProcessor_Primakoff()
{
    // Parameters and Services should be accessed from Init() instead of here!
}

//------------------
// ~JEventProcessor_Primakoff (Destructor)
//------------------
JEventProcessor_Primakoff::~JEventProcessor_Primakoff()
{
	SetTypeName(NAME_OF_THIS); // Provide JANA with this class's name
}

//------------------
// Init
//------------------
void JEventProcessor_Primakoff::Init()
{
  auto app = GetApplication();
  lockService = app->GetService<JLockService>();
  lockService->CreateLock("mylock");
  
  lockService->RootWriteLock();
  gDirectory->mkdir("Primakoff")->cd();
  {
    BeamPhotonTimeDiff=new TH1F("BeamPhotonTimeDiff","t(tag)-t(rf) at vertex",
			      1000,-50,50);
    KinFitCL=new TH1F("KinFitCL","fit probability",100,0,1);
    TwoGammaMass_vs_theta=new TH2F("TwoGammaMass_vs_theta",
				   ";#theta(degrees);M(2#gamma) [GeV]",
				   400,0,10.,400,0,1.); 
    TwoGammaMass_vs_t=new TH2F("TwoGammaMass_vs_t",
			       ";-t [GeV^{2}];M(2#gamma) [GeV]",
			       2000,0,2.,400,0,1.);
  }
  gDirectory->cd("../");
  lockService->RootUnLock();
    
  FCAL_THRESHOLD=0.1;
  app->SetDefaultParameter("PRIMAKOFF:FCAL_THRESHOLD",FCAL_THRESHOLD);
  BCAL_THRESHOLD=0.05;
  app->SetDefaultParameter("PRIMAKOFF:BCAL_THRESHOLD",BCAL_THRESHOLD);
  FCAL_POS_CUT=4.0;
  app->SetDefaultParameter("PRIMAKOFF:FCAL_POS_CUT",FCAL_POS_CUT); 
  FCAL_RADIAL_CUT=104.0;
  app->SetDefaultParameter("PRIMAKOFF:FCAL_RADIAL_CUT",FCAL_RADIAL_CUT);
  BCAL_Z_CUT=384.0;
  app->SetDefaultParameter("PRIMAKOFF:BCAL_Z_CUT",BCAL_Z_CUT); 
  BCAL_R_CUT=89.0;
  app->SetDefaultParameter("PRIMAKOFF:BCAL_R_CUT",BCAL_R_CUT);

}

//------------------
// BeginRun
//------------------
void JEventProcessor_Primakoff::BeginRun(const std::shared_ptr<const JEvent> &event)
{
  const DGeometry *geom=DEvent::GetDGeometry(event);
  double x0,y0,z0;
  geom->GetFCALPosition(x0,y0,z0);
  mFCALCenter.SetXYZ(x0,y0,z0);
}

//------------------
// Process
//------------------
void JEventProcessor_Primakoff::Process(const std::shared_ptr<const JEvent> &event)
{
  vector<const DBeamPhoton*>beamphotons;
  event->Get(beamphotons);
  if (beamphotons.size()==0) return;
  
  vector<const DNeutralParticle*>neutrals;
  event->Get(neutrals);
  if (neutrals.size()==0) return;

  // Find t0 (at "vertex") for event and assign list of neutral particles
  double t0_rf=-1000.;
  DVector3 vertex(0.,0.,65.);
  vector<DNeutralParticleHypothesis>gammaParticles;
  FillGammaParticleVector(neutrals,gammaParticles,t0_rf);

  if (gammaParticles.size()==2){
    lockService->WriteLock("mylock");

    // Fill vector of beam photons to pass on the analysis along with weights 
    // for in-time/ out-of-time
    vector<pair<const DBeamPhoton*,double>>beamPhotons;
    for (unsigned int i=0;i<beamphotons.size();i++){
      double dt_st=beamphotons[i]->time()-t0_rf;
      BeamPhotonTimeDiff->Fill(dt_st);
      
      double weight=1.;
      bool got_beam_photon=false;
      if (fabs(dt_st)>6.012&&fabs(dt_st)<18.036){
	weight=-1./6.;
	got_beam_photon=true;
      }
      if (fabs(dt_st)<2.004){
	got_beam_photon=true;     
      }
      if (got_beam_photon){
	beamPhotons.push_back(make_pair(beamphotons[i],weight));
      }
    }
 
    // Setup for performing kinematic fits
    DKinFitUtils_GlueX *dKinFitUtils = new DKinFitUtils_GlueX(event);
    DKinFitter *dKinFitter = new DKinFitter(dKinFitUtils);   
    dKinFitUtils->Reset_NewEvent();
    dKinFitter->Reset_NewEvent();
    
    for (unsigned int i=0;i<beamPhotons.size();i++){
      double weight=beamPhotons[i].second;
      if (DoKinematicFit(t0_rf,vertex,beamPhotons[i].first,gammaParticles,
			 dKinFitUtils,dKinFitter)){
	double CL=dKinFitter->Get_ConfidenceLevel();
	KinFitCL->Fill(CL,weight);
	
	if (CL>0.01){
	  DLorentzVector beam_kf,missing_kf;
	  map<Particle_t,vector<DLorentzVector>>final_kf;
	  GetKF4vectors(dKinFitter,beam_kf,missing_kf,final_kf);
	  MissingProtonAnalysis(beam_kf,missing_kf,final_kf,weight);
	}
      }
    }

    delete dKinFitter;
    delete dKinFitUtils;

    lockService->Unlock("mylock");
  }
}


//------------------
// EndRun
//------------------
void JEventProcessor_Primakoff::EndRun()
{
    // 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.
}

//------------------
// Finish
//------------------
void JEventProcessor_Primakoff::Finish()
{
    // Called before program exit after event processing is finished.
}

void
JEventProcessor_Primakoff::FillGammaParticleVector(vector<const DNeutralParticle*>&neutrals,
						   vector<DNeutralParticleHypothesis>&gammaParticles,
						   double &t0_rf
						   ) const{
  for (unsigned int i=0;i<neutrals.size();i++){
    const DNeutralParticleHypothesis *gamHyp=neutrals[i]->Get_Hypothesis(Gamma);
    const DNeutralShower *shower=gamHyp->Get_NeutralShower();
    DLorentzVector gamma_v4=gamHyp->lorentzMomentum(); 

    // Look for good photons
    bool got_photon=true;
    if (shower->dDetectorSystem==SYS_FCAL||shower->dDetectorSystem==SYS_ECAL
	|| shower->dDetectorSystem==SYS_ECAL_FCAL){
      DVector3 pos=shower->dSpacetimeVertex.Vect();
      DVector3 pos_rel=pos-mFCALCenter;
      
      if (pos_rel.Perp()>FCAL_RADIAL_CUT) got_photon=false;
      if (fabs(pos_rel.X())<FCAL_POS_CUT && fabs(pos_rel.Y())<FCAL_POS_CUT){
	got_photon=false;
      }
      if (gamHyp->lorentzMomentum().E()<FCAL_THRESHOLD){
	got_photon=false;
      }
    }
    else if (shower->dDetectorSystem==SYS_BCAL){
      const DBCALShower *bcalshower=dynamic_cast<const DBCALShower*>(shower->dBCALFCALShower);
      DVector3 bcal_pos(bcalshower->x,bcalshower->y,bcalshower->z);
      if (bcal_pos.Perp()>BCAL_R_CUT){
	got_photon=false;
      }
      if (gamma_v4.E()<BCAL_THRESHOLD){
	got_photon=false;
      }
      if (bcalshower->z>BCAL_Z_CUT) got_photon=false;
    }
    if (got_photon){
      t0_rf=gamHyp->t0();
      gammaParticles.push_back(*gamHyp);
    }
  }
}

void JEventProcessor_Primakoff::GetKF4vectors(DKinFitter *dKinFitter,
					      DLorentzVector &beam_kf,
					      DLorentzVector &missing_kf,
					      map<Particle_t,vector<DLorentzVector>>&final_kf) const{
  
  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_MissingParticle){
      missing_kf=(*locParticleIterator)->Get_P4();
    }
    else if ((*locParticleIterator)->Get_KinFitParticleType()==d_DetectedParticle){
      Particle_t myPID=PDGtoPType((*locParticleIterator)->Get_PID());
      final_kf[myPID].push_back((*locParticleIterator)->Get_P4());
    }
  }
}


// Run the kinematic fitter requiring energy and momentum conservation
bool JEventProcessor_Primakoff::DoKinematicFit(double t0_rf,const DVector3 &vertex,

					       const DBeamPhoton *beamphoton,
					       vector<DNeutralParticleHypothesis>&gammaParticles,
					       DKinFitUtils_GlueX *dKinFitUtils,
					       DKinFitter *dKinFitter) const {
  set<shared_ptr<DKinFitParticle>> InitialParticles, FinalParticles;
  
  dKinFitter->Reset_NewFit();
  shared_ptr<DKinFitParticle>myBeam=dKinFitUtils->Make_BeamParticle(beamphoton);
  shared_ptr<DKinFitParticle>myTarget=dKinFitUtils->Make_TargetParticle(Proton);
  
  InitialParticles.insert(myBeam);  
  InitialParticles.insert(myTarget);

   // Vertex info
  DLorentzVector vert4;
  vert4.SetVect(vertex);
  vert4.SetT(t0_rf);

  // Missing proton
  shared_ptr<DKinFitParticle>myRecoil=dKinFitUtils->Make_MissingParticle(Proton);
  FinalParticles.insert(myRecoil);    

  // Final state photons
  for (unsigned int k=0;k<gammaParticles.size();k++){
    DNeutralParticleHypothesis gammaHyp=gammaParticles[k];
    shared_ptr<DKinFitParticle>myGamma=dKinFitUtils->Make_DetectedParticle(22,0,0.,vert4,gammaHyp.momentum(),0.,gammaHyp.errorMatrix());
    FinalParticles.insert(myGamma);
  }

  // make constraint
  shared_ptr<DKinFitConstraint_P4>locP4Constraint = dKinFitUtils->Make_P4Constraint(InitialParticles, FinalParticles);
  // set constraint
  dKinFitter->Add_Constraint(locP4Constraint);
  
  // PERFORM THE KINEMATIC FIT
  return dKinFitter->Fit_Reaction();

}

void JEventProcessor_Primakoff::MissingProtonAnalysis(const DLorentzVector &beam_kf,
						      
						      const DLorentzVector &missing_kf,
						      map<Particle_t,vector<DLorentzVector>>&final_kf,
						      double weight) const{
  DLorentzVector twogam_kf=final_kf[Gamma][0]+final_kf[Gamma][1];
  double twogam_mass=twogam_kf.M(); 
  double t=(beam_kf-twogam_kf).M2();

  TwoGammaMass_vs_theta->Fill(twogam_kf.Theta()*180./M_PI,
			      twogam_mass,weight); 
  TwoGammaMass_vs_t->Fill(-t, twogam_mass,weight);
}