#include <vector>
#include <cassert>

#include "AMPTOOLS_DATAIO/ROOTDataWriter.h"
#include "CLHEP/Vector/LorentzVector.h"

#include "TFile.h"
#include "TTree.h"
#include "TH1.h"

ROOTDataWriter::ROOTDataWriter( const string& outFile )
{
	
	TH1::AddDirectory( kFALSE );
	
	m_outFile = new TFile( outFile.c_str(), "recreate" );
	m_outTree = new TTree( "kin", "Kinematics" );
	
	m_outTree->Branch( "nPart", &m_nPart, "nPart/I" );
	m_outTree->Branch( "e", m_e, "e[nPart]/F" );
	m_outTree->Branch( "px", m_px, "px[nPart]/F" );
	m_outTree->Branch( "py", m_py, "py[nPart]/F" );
	m_outTree->Branch( "pz", m_pz, "pz[nPart]/F" );
	m_outTree->Branch( "eBeam", &m_eBeam, "eBeam/F" );
	m_outTree->Branch( "pxBeam", &m_pxBeam, "pxBeam/F" );
	m_outTree->Branch( "pyBeam", &m_pyBeam, "pyBeam/F" );
	m_outTree->Branch( "pzBeam", &m_pzBeam, "pzBeam/F" );
	m_outTree->Branch( "eRecoil", &m_eRecoil, "eRecoil/F" );
	m_outTree->Branch( "pxRecoil", &m_pxRecoil, "pxRecoil/F" );
	m_outTree->Branch( "pyRecoil", &m_pyRecoil, "pyRecoil/F" );
	m_outTree->Branch( "pzRecoil", &m_pzRecoil, "pzRecoil/F" );
	
	m_eventCounter = 0;
}

ROOTDataWriter::~ROOTDataWriter()
{
	m_outFile->cd();
	m_outTree->Write();
	m_outFile->Close();
}

void
ROOTDataWriter::writeEvent( const Kinematics& kin )
{
	vector< HepLorentzVector > particleList = kin.particleList();
	
	m_nPart = particleList.size() - 2;
	
	assert( particleList.size() <= Kinematics::kMaxParticles );

  m_eBeam = particleList[0].t();
	m_pxBeam = particleList[0].px();
	m_pyBeam = particleList[0].py();
	m_pzBeam = particleList[0].pz();
    
	m_eRecoil = particleList[1].t();
	m_pxRecoil = particleList[1].px();
	m_pyRecoil = particleList[1].py();
	m_pzRecoil = particleList[1].pz();
    
	for( int i = 0; i < m_nPart; ++i ){
		
		m_e[i] = particleList[i+2].t();
		m_px[i] = particleList[i+2].x();
		m_py[i] = particleList[i+2].y();
		m_pz[i] = particleList[i+2].z();
	}
	
	
	m_outTree->Fill();
    
  m_eventCounter++;
}