// $Id$
//
//    File: JEventProcessor_PipPimPi0Studies.cc
// Created: Fri Jan 17 16:34:34 EST 2020
// Creator: staylor (on Linux ifarm1802.jlab.org 3.10.0-1062.4.1.el7.x86_64 x86_64)
//

#include "JEventProcessor_PipPimPi0Studies.h"
using namespace jana;


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


//------------------
// JEventProcessor_PipPimPi0Studies (Constructor)
//------------------
JEventProcessor_PipPimPi0Studies::JEventProcessor_PipPimPi0Studies()
{

}

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

}

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

  gDirectory->mkdir("ThreePiQAnalysis")->cd();
  
  AcceptedP=new TH2F("AcceptedP","Accepted proton",200,0,2,90,0,90);
  AcceptedP->Sumw2(); 
  ThrownP=new TH2F("ThrownP","Thrown proton",200,0,2,90,0,90);
  ThrownP->Sumw2(); 

  string leaves = "Eg/F:t:M2pi:M2g:M2pi2g:prob:weight";
  threePiQTree=new TTree("threePiQTree","Pi+ Pi- 2Gamma Analysis");
  threePiQTree->Branch("T",treeval,leaves.c_str());

  string mcleaves = "Eg/F:t:M3PiQ";  
  threePiQMCTree=new TTree("threePiQMCTree","Pi+ Pi- 2Gamma MC Analysis");
  threePiQMCTree->Branch("T",mctreeval,mcleaves.c_str());


  gDirectory->cd("../");

  TIME_CUT=5.0;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:TIME_CUT",TIME_CUT);

  FCAL_RADIAL_CUT=104.0;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:FCAL_RADIAL_CUT",FCAL_RADIAL_CUT);
  BCAL_Z_CUT=384.0;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:BCAL_Z_CUT",BCAL_Z_CUT);

  FCAL_THRESHOLD=0.1;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:FCAL_THRESHOLD",FCAL_THRESHOLD);
  BCAL_THRESHOLD=0.05;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:BCAL_THRESHOLD",BCAL_THRESHOLD);
  
  SPLIT_CUT=0.5;
  gPARMS->SetDefaultParameter("PipPimPi0Studies:SPLIT_CUT",SPLIT_CUT); 


  return NOERROR;
}

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

//------------------
// evnt
//------------------
jerror_t JEventProcessor_PipPimPi0Studies::evnt(JEventLoop *loop, uint64_t eventnumber)
{
  vector<const DTAGMGeometry *>tagm_geom;
  loop->Get(tagm_geom); 
  vector<const DTAGHGeometry *>tagh_geom;
  loop->Get(tagh_geom); 
  double tagged_mc_beam_energy=0.;
  
  vector<const DMCThrown*>throwns;
  loop->Get(throwns); 
  
  double mc_mom=0.,mc_theta=0.;
  if (throwns.size()>0){
    const DMCReaction* reaction=NULL;
    loop->GetSingle(reaction);
    
    if (reaction!=NULL){
      double mc_beam_energy=reaction->beam.energy();
	
      if (tagm_geom.size()&&tagh_geom.size()){
	unsigned int counter=0,column=0;
	if ( tagm_geom[0]->E_to_column(mc_beam_energy,column)){
	  tagged_mc_beam_energy=0.5*(tagm_geom[0]->getEhigh(column)
				     +tagm_geom[0]->getElow(column));
	} 
	else if (tagh_geom[0]->E_to_counter(mc_beam_energy,counter)){
	  tagged_mc_beam_energy=0.5*(tagh_geom[0]->getEhigh(counter)
				     +tagh_geom[0]->getElow(counter));
	}
	if (tagged_mc_beam_energy>0.){
	  for (unsigned int i=0;i<throwns.size();i++){
	    if (throwns[i]->mech==0){     
	      switch(throwns[i]->type){
	      case Proton:
		{
		mc_mom=throwns[i]->momentum().Mag();
		mc_theta=throwns[i]->momentum().Theta()*180./M_PI;
		DLorentzVector target(0,0,0,ParticleMass(Proton));
		DLorentzVector beam(0,0,tagged_mc_beam_energy,tagged_mc_beam_energy);
		DLorentzVector missing=beam+target-throwns[i]->lorentzMomentum();
		double t_mc=(throwns[i]->lorentzMomentum()-target).M2();

		japp->RootWriteLock();
			
		mctreeval[0]=tagged_mc_beam_energy;
		mctreeval[1]=-t_mc;
		mctreeval[2]=missing.M();
		threePiQMCTree->Fill();
		ThrownP->Fill(mc_mom,mc_theta);

		japp->RootUnLock();
		}
		break;
	      default:
		break;
	      }
	    }
	  }
	}
      }
    }
  }
 
  vector<const DChargedTrack*>ctracks;
  loop->Get(ctracks);
  if (ctracks.size()!=3) return RESOURCE_UNAVAILABLE;
  
  vector<const DNeutralParticle*>neutrals;
  loop->Get(neutrals);
  if (neutrals.size()<2) return RESOURCE_UNAVAILABLE;
  
  double t0_rf=ctracks[0]->Get_BestTrackingFOM()->t0();
  vector<const DNeutralParticleHypothesis *>gams;
  for (unsigned int i=0;i<neutrals.size();i++){
    const DNeutralParticleHypothesis *gamHyp=neutrals[i]->Get_Hypothesis(Gamma);
    double tdiff=gamHyp->time()-t0_rf; 
    if (fabs(tdiff)<TIME_CUT){ 
      bool got_photon=true;
      // Get shower info corresponding to this hypothesis
      const DNeutralShower *shower=gamHyp->Get_NeutralShower();
      DLorentzVector gamma_v4=gamHyp->lorentzMomentum(); 
      if (shower->dDetectorSystem==SYS_FCAL){
        if (gamma_v4.E()<FCAL_THRESHOLD){
	  got_photon=false;
        }
        if (shower->dQuality<SPLIT_CUT) got_photon=false;
        DVector3 pos=(dynamic_cast<const DFCALShower*>(shower->dBCALFCALShower))
	  ->getPosition();
        if (pos.Perp()>FCAL_RADIAL_CUT) got_photon=false;
      }
      else if (shower->dDetectorSystem==SYS_BCAL){
        if (gamma_v4.E()<BCAL_THRESHOLD){
          got_photon=false;
        }
        const DBCALShower *bcalshower=dynamic_cast<const DBCALShower*>(shower->dBCALFCALShower);
        if (bcalshower->z>BCAL_Z_CUT) got_photon=false;
      }
      if (got_photon) gams.push_back(gamHyp);
    }
  }
  if (gams.size()!=2) return RESOURCE_UNAVAILABLE;

  vector<const DBeamPhoton*>beamphotons;
  loop->Get(beamphotons);  

// Assign charge track by PID

  map<Particle_t,vector<const DTrackTimeBased *> >particles;
  FillParticleVectors(ctracks,particles);

  japp->WriteLock("my3PiQlock");

  if (tagged_mc_beam_energy>0 && particles[Proton].size()==1){
    AcceptedP->Fill(mc_mom,mc_theta);
  }

  if (particles[Proton].size()==1&&particles[PiPlus].size()==1
      && particles[PiMinus].size()==1){
    DKinFitUtils_GlueX *dKinFitUtils = new DKinFitUtils_GlueX(loop);
    DKinFitter *dKinFitter = new DKinFitter(dKinFitUtils);   
    dKinFitUtils->Reset_NewEvent();
    dKinFitter->Reset_NewEvent();
    
    const DTrackTimeBased *piminus_track=particles[PiMinus][0];
    const DTrackTimeBased *piplus_track=particles[PiPlus][0];
    const DTrackTimeBased *proton_track=particles[Proton][0];

    for (unsigned int i=0;i<beamphotons.size();i++){
      double dt_st=beamphotons[i]->time()-t0_rf; 
      
      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==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>myPiPlus=dKinFitUtils->Make_DetectedParticle(piplus_track);
      shared_ptr<DKinFitParticle>myPiMinus=dKinFitUtils->Make_DetectedParticle(piminus_track);  
      
      FinalParticles.insert(myProton);
      FinalParticles.insert(myPiPlus);
      FinalParticles.insert(myPiMinus);
      
      TLorentzVector vert;
      vert.SetVect(proton_track->position());
      vert.SetT(proton_track->t0());
      shared_ptr<DKinFitParticle>myGamma1
	=dKinFitUtils->Make_DetectedParticle(1,0,0.,vert,
					     gams[0]->momentum(),0.,
					     gams[0]->errorMatrix());
      FinalParticles.insert(myGamma1); 
      shared_ptr<DKinFitParticle>myGamma2
	=dKinFitUtils->Make_DetectedParticle(1,0,0.,vert,
					     gams[1]->momentum(),0.,
					     gams[1]->errorMatrix());
      FinalParticles.insert(myGamma2);
      
      // make constraint
      shared_ptr<DKinFitConstraint_P4>locP4Constraint = dKinFitUtils->Make_P4Constraint(InitialParticles, FinalParticles);
      // set constraint
      dKinFitter->Add_Constraint(locP4Constraint);
      
      // Production Vertex constraint
      set<shared_ptr<DKinFitParticle>> locFullConstrainParticles, locNoConstrainParticles;
      locFullConstrainParticles.insert(myPiPlus);  
      locFullConstrainParticles.insert(myPiMinus);
      locFullConstrainParticles.insert(myProton);		
      locNoConstrainParticles.insert(myBeam);  
      locNoConstrainParticles.insert(myTarget);
      shared_ptr<DKinFitConstraint_Vertex> locProductionVertexConstraint = dKinFitUtils->Make_VertexConstraint(locFullConstrainParticles, locNoConstrainParticles, proton_track->position());
      dKinFitter->Add_Constraint(locProductionVertexConstraint);
      
      // PERFORM THE KINEMATIC FIT
      dKinFitter->Fit_Reaction();
      
      // GET THE FIT RESULTS
      double locConfidenceLevel = dKinFitter->Get_ConfidenceLevel();
      
      if (locConfidenceLevel>0){
	DLorentzVector beam_kf;
	vector<DLorentzVector>gammas_kf;
	DLorentzVector proton_kf,pip_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 211: // pi+
	      pip_kf=(*locParticleIterator)->Get_P4();
	      break;
	    case -211: // pi-
	      pim_kf=(*locParticleIterator)->Get_P4();
	      break;
	    case 1:  // Gamma
	      gammas_kf.push_back((*locParticleIterator)->Get_P4());
	      break;
	    default:
	      break;
	    }
	  }
	}
	
	DLorentzVector target(0,0,0,ParticleMass(Proton));
	double t_mandelstam=(proton_kf-target).M2();
	DLorentzVector twogam=gammas_kf[0]+gammas_kf[1];
	DLorentzVector pippim_kf=pip_kf+pim_kf;
	DLorentzVector pippim2g_kf=pippim_kf+twogam;
	
	treeval[0]=beam_kf.E();
	treeval[1]=-t_mandelstam;
	treeval[2]=pippim_kf.M();
	treeval[3]=twogam.M();
	treeval[4]=pippim2g_kf.M();
	treeval[5]=locConfidenceLevel;
	treeval[6]=weight;
	threePiQTree->Fill();
      }
    }

    delete dKinFitter;
    delete dKinFitUtils;    
  }

  japp->Unlock("my3PiQlock");

  return NOERROR;
}

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


void
JEventProcessor_PipPimPi0Studies::FillParticleVectors(vector<const DChargedTrack *>&tracks,
					     map<Particle_t, vector<const DTrackTimeBased *> >&particles
					     ) const {
  for (unsigned int j=0;j<tracks.size();j++){  
    const DChargedTrackHypothesis *hyp=NULL;
    vector<pair<double,Particle_t> >probabilities;
    Particle_t particle_list[3]={Proton,PiPlus,PiMinus};
    for (unsigned int i=0;i<3;i++){
      if ((hyp=tracks[j]->Get_Hypothesis(particle_list[i]))!=NULL){
	double prob=TMath::Prob(hyp->Get_ChiSq_DCdEdx(),hyp->Get_NDF_DCdEdx());
	probabilities.push_back(make_pair(prob,particle_list[i]));
      } 
    }
    if (probabilities.size()>0){
      sort(probabilities.begin(),probabilities.end(),SortParticleProbability);
      Particle_t myParticle=probabilities[0].second;
      hyp=tracks[j]->Get_Hypothesis(myParticle);
      particles[myParticle].push_back(hyp->Get_TrackTimeBased());
    }
  }
}