// $Id$
//
//    File: JEventProcessor_CalCal.cc
// Created: Mon Feb  2 08:42:18 AM EST 2026
// Creator: staylor (on Linux ifarm2402.jlab.org 5.14.0-611.20.1.el9_7.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_CalCal.h"
#include <PID/DChargedTrack.h>
#include <TRACKING/DTrackCandidate.h>
#include <ECAL/DECALShower.h>
#include <ECAL/DECALCluster.h>
#include <ECAL/DECALHit.h>
#include <TRIGGER/DTrigger.h>
#include "DANA/DEvent.h"
#include <TMatrixD.h>
#include <fstream>

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


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

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

//------------------
// Init
//------------------
void JEventProcessor_CalCal::Init()
{
  auto app = GetApplication();  
  // lockService should be initialized here like this
  lockService = app->GetService<JLockService>();
  lockService->RootWriteLock();

  hDeltaT_Electron=new TH1F("hDeltaT_Electron",
			    ";t_{shower}-t_{flight}-t_{rf} [ns]",100,-10,10);
  hDeltaT_Positron=new TH1F("hDeltaT_Positron",
			    ";t_{shower}-t_{flight}-t_{rf} [ns]",100,-10,10);
  hEoverP_ECAL_Electron=new TH2F("hEoverP_ECAL_Electron",";p [GeV/c];E/p",100,0,10,100,0,1.5);
  hEoverP_ECAL_Positron=new TH2F("hEoverP_ECAL_Positron",";p [GeV/c];E/p",100,0,10,100,0,1.5);
  hEfitOverP_Electron=new TH2F("hEfitOverP_Electron",";p [GeV/c];E/p",100,0,10,100,0,1.5);
  hEfitOverP_Positron=new TH2F("hEfitOverP_Positron",";p [GeV/c];E/p",100,0,10,100,0,1.5);
  hFitProb1=new TH1F("hFitProb1",";E/p fit CL",1000,0,1); 
  hdEdx_CDC=new TH2F("hdEdx_CDC",";p [GeV/c];dE/dx [keV/cm]",100,0,2.5,200,0,25.);
  hdEdx_CDC_cut=new TH2F("hdEdx_CDC_cut",";p [GeV/c];dE/dx [keV/cm]",100,0,2.5,200,0,25.);
  hProtonGammaTimeDiff=new TH1F("hProtonGammaTimeDiff",";t_{0}-t(#gamma) [ns]",100,-25,25);
  h2GammaMass=new TH1F("h2GammaMass",";m(2#gamma) [GeV]",1000,0,1.);
  h2GammaMassImproved=new TH1F("h2GammaMassImproved",";m(2#gamma) [GeV]",1000,0,1.);
  h2GammaMass_vs_E=new TH2F("h2GammaMass_vs_E",";E [GeV];m(2#gamma) [GeV]",70,0,7,800,0,0.8);
  hGainScaleFactor=new TH2F("hGainScaleFactor",";channel;g/g_{0}",1600,-0.5,1599.5,200,0.8,1.2);
  hChannelDeltaT=new TH2F("hChannelDeltaT",";channel;t_{hit}-t_{flight}-t_{rf} [ns]",1600,-0.5,1599.5,200,-2.,2.);

  lockService->RootUnLock();

  MIN_NUM_ECAL_BLOCKS=3;
  app->SetDefaultParameter("CalCal:MIN_NUM_ECAL_BLOCKS",MIN_NUM_ECAL_BLOCKS);
  MAX_NUM_ECAL_BLOCKS=25;
  app->SetDefaultParameter("CalCal:MAX_NUM_ECAL_BLOCKS",MAX_NUM_ECAL_BLOCKS);
  TRACK_CL_CUT=0.001;
  app->SetDefaultParameter("CalCal:TRACK_CL_CUT",TRACK_CL_CUT);

  // ECAL hit resolution parameters
  VAR_E_PAR0=7.7e-5;
  VAR_E_PAR1=-6.9e-5;
  VAR_E_PAR2=2.3e-5;
  app->SetDefaultParameter("CalCal:VAR_E_PAR0",VAR_E_PAR0);
  app->SetDefaultParameter("CalCal:VAR_E_PAR1",VAR_E_PAR1);
  app->SetDefaultParameter("CalCal:VAR_E_PAR2",VAR_E_PAR2);
}

//------------------
// BeginRun
//------------------
void JEventProcessor_CalCal::BeginRun(const std::shared_ptr<const JEvent> &event)
{
  auto runnumber = event->GetRunNumber();
  auto app = event->GetJApplication();
  auto geo_manager = app->GetService<DGeometryManager>();
  auto dgeom = geo_manager->GetDGeometry(runnumber);
  if (dgeom->HaveInsert()){
    event->GetSingle(dECALGeom);
    ifstream infile("ecal_gains.dat");
    scale_factors.clear();
    for (size_t i=0;i<1600;i++){
      double dummy;
      infile >> dummy;
      scale_factors.push_back(dummy);
    }
    infile.close();
    DEvent::GetCalib(event, "/ECAL/gains", old_gains);
  }
}

//------------------
// Process
//------------------
void JEventProcessor_CalCal::Process(const std::shared_ptr<const JEvent> &event)
{
  // Look for events with energy in the forward calorimeters and one track
  auto triggers=event->Get<DTrigger>();
  if (triggers.size()==0) return;
  if (!(triggers[0]->Get_L1TriggerBits()&0x1)) return;
  auto candidates=event->Get<DTrackCandidate>();
  if (candidates.size()>2) return;
  auto tracks=event->Get<DChargedTrack>();
  if (tracks.size()!=1) return;
  
  auto ecalshowers=event->Get<DECALShower>();
  auto fcalshowers=event->Get<DFCALShower>();

  // Acquire root fill lock
  lockService->RootFillLock(this); 

  for (auto &track:tracks){
    auto hyp=(*track).Get_Hypothesis(Electron);
    if (hyp!=nullptr){
      LeptonAnalysis(hyp,hEoverP_ECAL_Electron,hEfitOverP_Electron,
		     hDeltaT_Electron);
    }
    hyp=(*track).Get_Hypothesis(Positron);
    if (hyp!=nullptr){
      LeptonAnalysis(hyp,hEoverP_ECAL_Positron,hEfitOverP_Positron,
		     hDeltaT_Positron);
    }
    if (tracks.size()==1){
      hyp=(*track).Get_Hypothesis(Proton);
      if (hyp!=nullptr){
	GammaAnalysis(hyp,ecalshowers);
      }
    }
  }

  // Release root fill lock
  lockService->RootFillUnLock(this);
}

//------------------
// EndRun
//------------------
void JEventProcessor_CalCal::EndRun()
{

}

//------------------
// Finish
//------------------
void JEventProcessor_CalCal::Finish()
{
  ofstream gainfile("ecal_gains.dat");
  for (size_t i=0;i<scale_factors.size();i++){
    gainfile << scale_factors[i] << endl;
  }
}

// E/p analysis for electron/positron tracks
void JEventProcessor_CalCal::LeptonAnalysis(const DChargedTrackHypothesis *hyp,
					    TH2F *histo,TH2F *hEfitOverP,TH1F *thisto){
  auto tb_track=hyp->Get_TrackTimeBased();
  if (tb_track->FOM>TRACK_CL_CUT){
    auto ecal_match=hyp->Get_ECALShowerMatchParams();
    if (ecal_match!=nullptr
	&& ecal_match->dECALShower->nBlocks>=MIN_NUM_ECAL_BLOCKS
	&& ecal_match->dECALShower->nBlocks<=MAX_NUM_ECAL_BLOCKS
	){
      double dt_flight_rf=ecal_match->dFlightTime+hyp->t0();
      double dt=ecal_match->dECALShower->t-dt_flight_rf;
      thisto->Fill(dt);
      if (fabs(dt)<2.0){
	auto ecal_extraps=tb_track->extrapolations.at(SYS_ECAL);
	double p=ecal_extraps[0].momentum.Mag();
	double Es=ecal_match->dECALShower->E;
	histo->Fill(p,Es/p);

	// Fitted energy
	auto cluster=ecal_match->dECALShower->GetSingle<DECALCluster>();
	hEfitOverP->Fill(p,cluster->Efit/p);
	
	if (Es/p>0.8 && p>0.95 && p<1.05 && cluster->status==1){
	  auto hits=cluster->Get<DECALHit>();
	  vector<double>E(hits.size());
	  vector<double>V(hits.size());
	  double sumE=0.,sumV=0.0025*p*p;
	  for (size_t i=0;i<hits.size();i++){
	    E[i]=hits[i]->E;
	    V[i]=VAR_E_PAR0+VAR_E_PAR1*E[i]+VAR_E_PAR2*E[i]*E[i];
	    sumE+=E[i];
	    sumV+=V[i];
	  }
	  double Pdiff=sumE-p;
	  // Lagrange multiplier
	  double lagrange=Pdiff/sumV;
	  double new_sumE=0.;
	  for (size_t i=0;i<hits.size();i++){
	    E[i]-=lagrange*V[i];
	    new_sumE+=E[i];
	    if (E[i]>0.1){
	      int channel=dECALGeom->channel(hits[i]->row,hits[i]->column);
	      //hGainScaleFactor->Fill(channel,E[i]/hits[i]->E);
	    }
	  }
	  double new_p=p+lagrange*0.0025*p*p;	      
	  double chisq=lagrange*(Pdiff+2.*(new_sumE-new_p));
	  hFitProb1->Fill(TMath::Prob(chisq,1));
	}
      }
    }
  }
}

// Final state photon analysis
void JEventProcessor_CalCal::GammaAnalysis(const DChargedTrackHypothesis *hyp,
					  vector<const DECALShower *>&ecalshowers){
  auto tb_track=hyp->Get_TrackTimeBased();
  if (tb_track->FOM>TRACK_CL_CUT){
    double dEdx=1e6*hyp->Get_dEdx_CDC_amp();
    if (dEdx>0.){
      double p=tb_track->momentum().Mag();
      hdEdx_CDC->Fill(p,dEdx);
      if (hyp->Get_FOM()>0.01){
	hdEdx_CDC_cut->Fill(p,dEdx);
	auto vertex=tb_track->position();
	double t0=hyp->t0();
	if (ecalshowers.size()==2){
	  for (size_t i=0;i<ecalshowers.size();i++){   
	    // Check the timing with respect to the "vertex" time at the target
	    auto diff1=ecalshowers[i]->pos-vertex;
	    double dt=t0-(ecalshowers[i]->t-diff1.Mag()/29.98);
	    hProtonGammaTimeDiff->Fill(dt);
	    if (fabs(dt)<2.){
	      // Get the cluster object associated with this shower
	      auto cluster1=ecalshowers[i]->GetSingle<DECALCluster>();
	      if (cluster1->status!=1) continue; // Single fitted peak?

	      // Skip if block with maximum energy is at the ECAL/FCAL
	      // border
	      int channel_Emax=cluster1->channel_Emax;
	      int row=dECALGeom->row(channel_Emax);
	      int col=dECALGeom->column(channel_Emax);
	      if (row==0 || col==0 || row==39 || col==39) continue;

	      // Channel-by-channel timing
	      int channel=cluster1->channel_Emax;
	      hChannelDeltaT->Fill(channel,dt);
	      
	      // Shower energy
	      double E1S=ecalshowers[i]->E;
	      
	      // Get the hits associated with this cluster
	      vector<const DECALHit*>cluster1_hits;
	      vector<double>E1,E1Improved,V1;
	      double V1sum=0;
	      if (cluster1!=NULL){
		cluster1_hits=cluster1->Get<DECALHit>();
		// Store energy and variance for each hit
		for (size_t k=0;k<cluster1_hits.size();k++){
		  const DECALHit *hit=cluster1_hits[k];
		  int channel=dECALGeom->channel(hit->row,hit->column);
		  double E=scale_factors[channel]*cluster1_hits[k]->E
		    /old_gains[channel];
		  E1.push_back(E);
		  E1Improved.push_back(E);
		  double Vtemp=VAR_E_PAR0+VAR_E_PAR1*E+VAR_E_PAR2*E*E;
		  V1.push_back(Vtemp);
		  V1sum+=Vtemp;
		}
	      }
	      
	      for (size_t j=i+1;j<ecalshowers.size();j++){
		// Check timing with respect to the "vertex" in the target
		auto diff2=ecalshowers[j]->pos-vertex;
		dt=t0-(ecalshowers[j]->t-diff2.Mag()/29.98);
		hProtonGammaTimeDiff->Fill(dt);
		if (fabs(dt)<2.){
		  auto cluster2=ecalshowers[j]->GetSingle<DECALCluster>();
		  if (cluster2->status!=1) continue; // Single fitted peak?

		  // Skip if block with maximum energy is at the ECAL/FCAL
		  // border
		  int channel_Emax=cluster2->channel_Emax;
		  int row=dECALGeom->row(channel_Emax);
		  int col=dECALGeom->column(channel_Emax);
		  if (row==0 || col==0 || row==39 || col==39) continue;
				  
		  // Shower energy
		  double E2S=ecalshowers[j]->E;

		  // Two-photon invariant mass
		  auto dir1=diff1.Unit();
		  auto dir2=diff2.Unit();
		  double one_minus_costheta12=1.-dir1.Dot(dir2);
		  double Msq=2.*E1S*E2S*one_minus_costheta12;
		  double M=sqrt(Msq);
		  h2GammaMass->Fill(M);
		  if (fabs(E1S-E2S)<0.1){
		    h2GammaMass_vs_E->Fill(0.5*(E1S+E2S),M);
		  }
		  
		  if (M>0.1 && M<0.15 && E1S>1.0 && E2S>1.0){
		    // Get the hits associated with this second cluster
		    vector<const DECALHit*>cluster2_hits;
		    vector<double>E2,E2Improved,V2;
		    double V2sum=0;
		    if (cluster2!=NULL){
		      cluster2_hits=cluster2->Get<DECALHit>();
		      // Store energy and variance for each hit
		      for (size_t k=0;k<cluster2_hits.size();k++){
			const DECALHit *hit=cluster2_hits[k];
			int channel=dECALGeom->channel(hit->row,hit->column);
			double E=scale_factors[channel]*cluster2_hits[k]->E
			  /old_gains[channel];
			E2.push_back(E);
			E2Improved.push_back(E);
			double Vtemp=VAR_E_PAR0+VAR_E_PAR1*E+VAR_E_PAR2*E*E;
			V2.push_back(Vtemp);
			V2sum+=Vtemp;
		      }
		    }
		    // Apply pi0 mass constraint
		    double E1sum=E1S;
		    double E2sum=E2S;
		    double Mpi0_sq=pow(ParticleMass(Pi0),2.);
		    double MsqDiff=Msq-Mpi0_sq;
		    int iter=0;
		    double chisq=1e10,chisq_old;
		    do {
		      chisq_old=chisq;
		      double r=2.*one_minus_costheta12
			*(E2sum*E1S+E1sum*E2S-E1sum*E2sum)-Mpi0_sq;
		      double S=4*one_minus_costheta12*one_minus_costheta12
			*(E1sum*E1sum*V2sum+E2sum*E2sum*V1sum);
		      // Lagrange multiplier
		      double lagrange=r/S;
		      
		      for (size_t k=0;k<cluster1_hits.size();k++){
			E1Improved[k]=E1[k]-2.*one_minus_costheta12*E2sum*lagrange*V1[k];
		      }
		      for (size_t k=0;k<cluster2_hits.size();k++){
			E2Improved[k]=E2[k]-2.*one_minus_costheta12*E1sum*lagrange*V2[k];
		      }
		      E1sum=0.;
		      for (size_t k=0;k<cluster1_hits.size();k++){
			E1sum+=E1Improved[k];
		      }
		      E2sum=0.;
		      for (size_t k=0;k<cluster2_hits.size();k++){
			E2sum+=E2Improved[k];
		      }
		      chisq=lagrange*(S*lagrange+2.*r);
		      iter++;
		    }
		    while (fabs(chisq-chisq_old)>0.01 && iter<20);
		    
		    // "Improved" mass
		    double M_improved=sqrt(2.*E1sum*E2sum*one_minus_costheta12);
		    h2GammaMassImproved->Fill(M_improved);
		    
		    // Update gain factor histograms
		    for (size_t k=0;k<cluster1_hits.size();k++){
		      const DECALHit *hit=cluster1_hits[k];
		      if (hit->E>0.25){
			int channel=dECALGeom->channel(hit->row,hit->column);
			double gain_ratio=E2Improved[k]/hit->E;
			int channel_Emax=cluster1->channel_Emax;
			if (gain_ratio>0.8 && gain_ratio<1.2){
			  scale_factors[channel]=gain_ratio*old_gains[channel];
			  hGainScaleFactor->Fill(channel,gain_ratio);
			}
		      }
		    }
		    for (size_t k=0;k<cluster2_hits.size();k++){
		      const DECALHit *hit=cluster2_hits[k];
		      if (hit->E>0.25){
			int channel=dECALGeom->channel(hit->row,hit->column);
			double gain_ratio=E2Improved[k]/hit->E;
			if (gain_ratio>0.8 && gain_ratio<1.2){
			  scale_factors[channel]=gain_ratio*old_gains[channel];
			  hGainScaleFactor->Fill(channel,gain_ratio);
			}
		      }
		    }
		  } // rough pi0 mass cut
		}
	      }
	    }
	  }
	}
      }
    }
  }
}