// $Id$
//
//    File: JEventProcessor_cpp_tree.cc
// Created: Tue May 26 09:55:48 EDT 2015
// Creator: davidl (on Darwin harriet.jlab.org 13.4.0 i386)
//

#include "JEventProcessor_cpp_tree.h"
using namespace jana;

#include <TRACKING/DTrackTimeBased.h>
#include <TRACKING/DMCThrown.h>

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


//------------------
// JEventProcessor_cpp_tree (Constructor)
//------------------
JEventProcessor_cpp_tree::JEventProcessor_cpp_tree()
{

	
}

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

}

//------------------
// init
//------------------
jerror_t JEventProcessor_cpp_tree::init(void)
{
	japp->RootWriteLock();

	cppinfo_ptr = &cppinfo_obj;
	cpptree = new TTree("cpptree","");
	cpptree->Branch("C", &cppinfo_ptr);

	japp->RootUnLock();

	return NOERROR;
}

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

//------------------
// evnt
//------------------
jerror_t JEventProcessor_cpp_tree::evnt(JEventLoop *loop, int eventnumber)
{
	vector<const DMCThrown*> throwns;
	vector<const DTrackTimeBased*> tbts;
	loop->Get(throwns);
	loop->Get(tbts);
	
	//-------- Generated ----------
	vector<TLorentzVector> pip_throwns;
	vector<TLorentzVector> pim_throwns;
	for(uint32_t i=0; i<throwns.size(); i++){
		if(throwns[i]->PID() == 8) pip_throwns.push_back(throwns[i]->lorentzMomentum());
		if(throwns[i]->PID() == 9) pim_throwns.push_back(throwns[i]->lorentzMomentum());
	}
	if(!pip_throwns.empty() && !pim_throwns.empty()){
		TLorentzVector &pip = cppinfo_ptr->pip_thrown = pip_throwns[0];
		TLorentzVector &pim = cppinfo_ptr->pim_thrown = pim_throwns[0];
		TLorentzVector pipi = pip + pim;
		
		cppinfo_ptr->phi_pipi_thrown = pipi.Phi();
		cppinfo_ptr->psi_pipi_thrown = CalcPsi(pip, pipi);
		cppinfo_ptr->W_pipi_thrown = pipi.M();
	}

	//-------- Reconstructed ----------
	vector<TLorentzVector> pips;
	vector<TLorentzVector> pims;
	for(uint32_t i=0; i<tbts.size(); i++){
		if(tbts[i]->PID() == 8) pips.push_back(tbts[i]->lorentzMomentum());
		if(tbts[i]->PID() == 9) pims.push_back(tbts[i]->lorentzMomentum());
	}

	japp->RootWriteLock();
	
	cppinfo_ptr->event = eventnumber;
	cppinfo_ptr->Ncombos = pips.size()*pims.size();
	
	for(uint32_t i=0; i<pips.size(); i++){
		
		TLorentzVector &pip = cppinfo_ptr->pip = pips[i];
	
		for(uint32_t j=0; j<pims.size(); j++){

			TLorentzVector &pim = cppinfo_ptr->pim = pims[j];
			TLorentzVector pipi = pip + pim;
			
			cppinfo_ptr->phi_pipi = pipi.Phi();
			cppinfo_ptr->psi_pipi = CalcPsi(pip, pipi);
			cppinfo_ptr->W_pipi = pipi.M();
			
			cpptree->Fill();
		}
	}
	
	japp->RootUnLock();

	return NOERROR;
}

//------------------
// erun
//------------------
jerror_t JEventProcessor_cpp_tree::erun(void)
{
	return NOERROR;
}

//------------------
// fini
//------------------
jerror_t JEventProcessor_cpp_tree::fini(void)
{
	return NOERROR;
}

//------------------
// CalcPsi
//------------------
double JEventProcessor_cpp_tree::CalcPsi(TLorentzVector &pip, TLorentzVector &pipi)
{
	//------ Rory's recipe ------
	TVector3 k_hat = pipi.Vect();
	k_hat.SetMag(1.0);

	// Rotation angle about y to get into y-z plane
	double theta_yz = atan2(-k_hat.X(), k_hat.Z());
	k_hat.RotateY(theta_yz);

	// pi+ direction in rotated lab frame
	TVector3 pip_hat = pip.Vect();
	pip_hat.RotateY(theta_yz);

	// Angle to rotate k_hat into z-direction
	double theta_x = atan2(k_hat.Y(), k_hat.Z());
	k_hat.RotateX(theta_x);
	pip_hat.RotateX(theta_x);
	double psi = pip_hat.Phi();
	//----------------------------
	
	return psi;
}