// $Id$
//
//    File: JEventProcessor_cpp_pythia.cc
// Created: Wed Aug 22 10:25:39 EDT 2012
// Creator: davidl (on Linux ifarm1102 2.6.18-274.3.1.el5 x86_64)
//

#include "JEventProcessor_cpp_pythia.h"
using namespace jana;

#include <TRACKING/DMCThrown.h>


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

#include <particleType.h>
double CHARGED_PION_MASS = ParticleMass(PiPlus);

typedef struct{
	Int_t event;
	Int_t isPiPi_exclusive;
	Int_t isPiPi_inclusive;
	Int_t isKK;
	Float_t W;     // using perfect pid
	Float_t Wpipi; // using pion assumption
	Float_t t;
	Float_t Egamma;
	Float_t Epm;   // total energy of +/- particle pair with perfect PID
	Float_t Epipi; // same as above but with pion assumption
	Float_t Ebaryon; // proton mass if positive is a proton. Zero otherwise
}KININFO_t;
KININFO_t *kinfo = NULL;

string kinfo_description = "event/I:isPiPi_exclusive:isPiPi_inclusive:isKK:W/F:Wpipi:t:Egamma:Epm:Epipi:Ebaryon";

//------------------
// LorentzVectorDebugString
//------------------
string LorentzVectorDebugString(const DLorentzVector &v)
{
	// For debugging
	char str[256];
	sprintf(str, "m=%f E=%f theta=%f phi=%f", v.M(), v.E(), v.Theta(), v.Phi());
	return string(str);
}

//------------------
// JEventProcessor_cpp_pythia (Constructor)
//------------------
JEventProcessor_cpp_pythia::JEventProcessor_cpp_pythia()
{

}

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

}

//------------------
// init
//------------------
jerror_t JEventProcessor_cpp_pythia::init(void)
{
	Wplus_minus = new TH1D("Wplus_minus", "W_{+-}", 2000, 0.0, 4.0);
	Eplus_minus = new TH1D("Eplus_minus", "E_{+-}", 500, 0.0, 12.0);
	Wplus_minus_coherent = new TH1D("Wplus_minus_coherent", "W_{+-} in Coherent Brem. Peak", 2000, 0.0, 4.0);

	Wplus_minus->SetStats(0);
	Eplus_minus->SetStats(0);
	Wplus_minus_coherent->SetStats(0);
	
	Wplus_minus->SetXTitle("Invariant mass W of + and - (GeV)");
	Eplus_minus->SetXTitle("Total Energy E of + and - (GeV)");
	Wplus_minus_coherent->SetXTitle("Invariant mass W of + and - (GeV)");

	Wpipi = (TH1D*)Wplus_minus->Clone("Wpipi");
	Epipi = (TH1D*)Eplus_minus->Clone("Epipi");
	Wpipi_coherent = (TH1D*)Wplus_minus_coherent->Clone("Wpipi_coherent");

	Wplus_minus->SetLineColor(kRed);
	Eplus_minus->SetLineColor(kRed);
	Wplus_minus_coherent->SetLineColor(kRed);

	Wpipi->SetLineColor(kBlue);
	Epipi->SetLineColor(kBlue);
	Wpipi_coherent->SetLineColor(kBlue);
	
	Wplus_minusKK = (TH1D*)Wplus_minus->Clone("Wplus_minusKK");
	WpipiKK = (TH1D*)Wpipi->Clone("WpipiKK");
	Wplus_minusKK_coherent = (TH1D*)Wplus_minus_coherent->Clone("Wplus_minusKK_coherent");
	WpipiKK_coherent = (TH1D*)Wpipi_coherent->Clone("WpipiKK_coherent");

	Wplus_minusPiPi_exclusive = (TH1D*)Wpipi_coherent->Clone("Wplus_minusPiPi_exclusive");
	Wplus_minusPiPi_inclusive = (TH1D*)Wpipi_coherent->Clone("Wplus_minusPiPi_inclusive");

	t_all = new TH1D("t_all", "Momentum transfer^2 t", 1000, -1.0, 0.0);
	t_KK = (TH1D*)t_all->Clone("t_KK");
	t_PiPi = (TH1D*)t_all->Clone("t_PiPi");
	t_coherent = (TH1D*)t_all->Clone("t_coherent");
	t_KK_coherent = (TH1D*)t_all->Clone("t_KK_coherent");
	t_PiPi_coherent = (TH1D*)t_all->Clone("t_PiPi_coherent");
	
	kinfo = new KININFO_t;
	kininfo_tree = new TTree("kinfo", "");
	kininfo_tree->Branch("K", kinfo, kinfo_description.c_str());

	return NOERROR;
}

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

//------------------
// evnt
//------------------
jerror_t JEventProcessor_cpp_pythia::evnt(JEventLoop *loop, int eventnumber)
{
	vector<const DMCThrown*> throwns;
	loop->Get(throwns);
	
	vector<const DMCThrown*> positives;
	vector<const DMCThrown*> negatives;
	topology_t topology;
	topology.Clear();
	
	// Get incident photon energy by adding up energy
	// of all thrown particles and subtracting off the
	// proton mass.
	double Egamma = 0.0;
	
	// Calculate momentum transfer "t" by looking for the
	// baryon (proton or neutron) in the final state 
	// particle list. Declare and initialize the variable
	// here.
	double t = +1000.0;
	
	// Sort into lists of positive and negative particles
	for(unsigned int i=0; i<throwns.size(); i++){
		const DMCThrown *thrown = throwns[i];
		
		// Add to total energy of event
		if(thrown->parentid==0)Egamma += thrown->energy();
		
		// Fill in event topology
		int type = thrown->type;
		if(type>=0 && type<50)topology.Nparts[type]++;	

		// If this is baryon, calculate "t"
		if(type==14 || type==13){
			DLorentzVector target(0.0, 0.0, 0.0, 0.938272);
			DLorentzVector pdiff = target - thrown->lorentzMomentum();
			t = pdiff.Mag2();
		}
				
		// Only interested in forward going particles
		double theta_deg = thrown->momentum().Theta()*57.3;
		if(theta_deg<0.5 || theta_deg>10.0)continue;
		
		// Exclude intermediaries that PYTHIA has already decayed
		if(type<1 || type>50)continue;
		
		if(thrown->charge() == +1.0)positives.push_back(thrown);
		if(thrown->charge() == -1.0)negatives.push_back(thrown);
	}
	
	// Subtract proton mass from total energy of final state
	// particles to get incident photon energy
	Egamma -= 0.938272;
	
	// Fill in non-topology-dependent t histos
	t_all->Fill(t);
	if(Egamma>=8.5 && Egamma<=9.5) t_coherent->Fill(t);
	
	// Check if this is a K+K- event
	topology_t topologyKK;
	topologyKK.Clear();
	topologyKK.Nparts[11]++; // K+
	topologyKK.Nparts[12]++; // K-
	topologyKK.Nparts[14]++; // proton
	bool isKK = (topology==topologyKK);

	// Check if this is an exclusive pi+pi- event
	topology_t topologyPiPi;
	topologyPiPi.Clear();
	topologyPiPi.Nparts[8]++; // pi+
	topologyPiPi.Nparts[9]++; // pi-
	topologyPiPi.Nparts[14]++; // proton
	bool isPiPi_exclusive = (topology==topologyPiPi);

	// Check if this is an inclusive pi+pi- event
	bool isPiPi_inclusive = true;
	isPiPi_inclusive &= (topology.Nparts[8]>0);
	isPiPi_inclusive &= (topology.Nparts[9]>0);
	isPiPi_inclusive &= (topology.Nparts[14]>0);
	
	// Loop over all pairs of positive and negative particles
	bool record_toplogy = false;
	for(unsigned int i=0; i<positives.size(); i++){
		DLorentzVector p = positives[i]->lorentzMomentum();
		DLorentzVector pip = p;
		pip.SetE(sqrt(CHARGED_PION_MASS*CHARGED_PION_MASS + p.P()*p.P()));

		for(unsigned int j=0; j<negatives.size(); j++){
			DLorentzVector n = negatives[j]->lorentzMomentum();
			DLorentzVector pim = n;
			pim.SetE(sqrt(CHARGED_PION_MASS*CHARGED_PION_MASS + n.P()*n.P()));

			DLorentzVector tot = p+n;
			double W = tot.M();
			double E = tot.E();
			Wplus_minus->Fill(W);
			Eplus_minus->Fill(E);
		
			// Assumption that both particles are pions
			DLorentzVector tot_pipi = pip+pim;
			double Wpi_pi = tot_pipi.M();
			double Epi_pi = tot_pipi.E();
			Wpipi->Fill(Wpi_pi);
			Epipi->Fill(Epi_pi);
			
			// Tree
			kinfo->event = eventnumber;
			kinfo->isPiPi_exclusive = isPiPi_exclusive;
			kinfo->isPiPi_inclusive = isPiPi_inclusive;
			kinfo->isKK = isKK;
			kinfo->W = W;
			kinfo->Wpipi = Wpi_pi;
			kinfo->t = t;
			kinfo->Egamma = Egamma;
			kinfo->Epm = E;
			kinfo->Epipi = Epi_pi;
			kinfo->Ebaryon = positives[i]->type==14 ? 0.938272:0.0;
			kininfo_tree->Fill();

#if 0
if(E>Egamma){
_DBG_<<eventnumber<<": "<<MakeTopologyString(topology)<<endl;
_DBG_<<"Egamma="<<Egamma<<"  E="<<E<<"  diff="<<E-Egamma<<endl;
_DBG_<<LorentzVectorDebugString(p)<<endl;
_DBG_<<LorentzVectorDebugString(n)<<endl;
_DBG_<<LorentzVectorDebugString(tot)<<endl;
for(unsigned int i=0; i<throwns.size(); i++){
	const DMCThrown *thrown = throwns[i];
	_DBG_<<" thrown "<<i<<" type="<<thrown->type<<" myid="<<thrown->myid<<" parentid="<<thrown->parentid<<" E="<<thrown->energy()<<endl;
}
}
#endif
			// Inside coherent peak
			// (n.b. on 9/14/2012 this was changed to use "Egamma"
			// instead of "E". Plots generated earlier used "E".)
			if(Egamma>=8.5 && Egamma<=9.5){
				Wplus_minus_coherent->Fill(W);
				Wpipi_coherent->Fill(Wpi_pi);

				if(Wpi_pi<0.4)record_toplogy = true;

				if(isKK){
					Wplus_minusKK_coherent->Fill(W);
					WpipiKK_coherent->Fill(Wpi_pi);
					t_KK_coherent->Fill(t);
				}
				
				if(isPiPi_inclusive){
					t_PiPi_coherent->Fill(t);
				}
			}

			// Topology-dependent histograms
			if(isKK){
				Wplus_minusKK->Fill(W);
				WpipiKK->Fill(Wpi_pi);
				t_KK->Fill(t);
			}
			if(isPiPi_exclusive){
				Wplus_minusPiPi_exclusive->Fill(W);
			}
			if(isPiPi_inclusive){
				Wplus_minusPiPi_inclusive->Fill(W);
			}
		}
	}

	if(record_toplogy){
		string str = MakeTopologyString(topology);
		topologies[str]++;
	}

	return NOERROR;
}

//------------------
// MakeTopologyString
//------------------
string JEventProcessor_cpp_pythia::MakeTopologyString(JEventProcessor_cpp_pythia::topology_t &top)
{
	stringstream ss;
	ss << "#gamma p #rightarrow ";
	for(int type=1; type<50; type++){
		if(top.Nparts[type]==0)continue;
		string symbol="?";
		switch(type){
			case 1: symbol="#gamma"; break;
			case 2: symbol="e^{+}"; break;
			case 3: symbol="e^{-}"; break;
			case 4: symbol="#nu"; break;
			case 5: symbol="#mu^{+}"; break;
			case 6: symbol="#mu^{-}"; break;
			case 7: symbol="#pi^{o}"; break;
			case 8: symbol="#pi^{+}"; break;
			case 9: symbol="#pi^{-}"; break;
			case 10: symbol="K_{L}"; break;
			case 11: symbol="K^{+}"; break;
			case 12: symbol="K^{-}"; break;
			case 13: symbol="n"; break;
			case 14: symbol="p"; break;
			case 15: symbol="!p"; break;
			case 16: symbol="K_{S}"; break;
			case 17: symbol="#eta"; break;
			case 18: symbol="#lambda"; break;
		}
		ss<<" + ";
		if(top.Nparts[type]>1)ss<<top.Nparts[type];
		ss<<symbol;
	}
	
	return ss.str();
}

//------------------
// erun
//------------------
jerror_t JEventProcessor_cpp_pythia::erun(void)
{
	// Sort according to number of events
	multimap<int,string> itopologies;
	map<string,int>::iterator iter = topologies.begin();
	for(; iter!=topologies.end(); iter++){
		itopologies.insert(pair<int,string>(iter->second, iter->first));
	}

	int Nbins = topologies.size();
	if(Nbins>20) Nbins=20;
	
	// Create and fill histogram to hold all topologies
	TH1D *htopologies = new TH1D("topologies", "Event Topologies", Nbins, 0.5, (double)Nbins+0.5);
	multimap<int,string>::reverse_iterator itermm = itopologies.rbegin();
	for(int bin=1; itermm!=itopologies.rend(); itermm++, bin++){
		htopologies->GetXaxis()->SetBinLabel(bin, itermm->second.c_str());
		htopologies->SetBinContent(bin, itermm->first);
	}
	
	htopologies->SetStats(0);
	
	// Make two copies with alterating bins set to zero.
	// This is so one can be overlaid on the other to help
	// make the bins more visually distinct
	TH1D *htopologies_a = (TH1D*)htopologies->Clone("topologies_a");
	TH1D *htopologies_b = (TH1D*)htopologies->Clone("topologies_b");
	for(int bin=2; bin<=htopologies_a->GetNbinsX(); bin+=2)htopologies_a->SetBinContent(bin, 0.0);
	for(int bin=1; bin<=htopologies_b->GetNbinsX(); bin+=2)htopologies_b->SetBinContent(bin, 0.0);

	htopologies_a->SetFillColor(kBlue);
	htopologies_b->SetFillColor(kCyan);

	htopologies_a->SetFillStyle(3000);
	htopologies_b->SetFillStyle(3001);

	return NOERROR;
}

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