// $Id$
//
//    File: JEventProcessor_pip_pim.cc
// Created: Sat Apr 14 13:18:21 EDT 2012
// Creator: davidl (on Darwin genmacbook.local 11.3.0 i386)
//

#include <iostream>
using namespace std;

#include "JEventProcessor_pip_pim.h"
using namespace jana;

#include <PID/DPiPlus.h>
#include <PID/DPiMinus.h>
#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_pip_pim());
}
} // "C"


//------------------
// JEventProcessor_pip_pim (Constructor)
//------------------
JEventProcessor_pip_pim::JEventProcessor_pip_pim()
{
	pthread_rwlock_init(&root_lock, NULL);
}

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

}

//------------------
// init
//------------------
jerror_t JEventProcessor_pip_pim::init(void)
{
	//japp->RootWriteLock();
	pthread_rwlock_wrlock(&root_lock);

	mass_pip_pim = new TH1D("mass_pip_pim", "#pi^{+} #pi^{-} invariant mass (reconstructed)", 1500, 0.0, 1.5);
	mass_pip_pim_gen = new TH1D("mass_pip_pim_gen", "#pi^{+} #pi^{-} invariant mass (generated)", 1500, 0.0, 1.5);

	pip_tree = new TTree("pip_tree", "Reconstruction pi^{+}");
	pim_tree = new TTree("pim_tree", "Reconstruction pi^{-}");
	pip_tree_gen = new TTree("pip_tree_gen", "Generated pi^{+}");
	pim_tree_gen = new TTree("pim_tree_gen", "Generated pi^{-}");
	
	v4_pip=NULL;
	v4_pim=NULL;
	v4_pip_gen=NULL;
	v4_pim_gen=NULL;
	pip_tree->Branch("P", &v4_pip);
	pim_tree->Branch("P", &v4_pim);
	pip_tree_gen->Branch("P", &v4_pip_gen);
	pim_tree_gen->Branch("P", &v4_pim_gen);

	prob_pip = new TH1D("prob_pip", "#pi^{+} probablity from PID #chi^{2} and N_{dof}", 500, 0.0, 1.0);
	prob_pim = new TH1D("prob_pim", "#pi^{-} probablity from PID #chi^{2} and N_{dof}", 500, 0.0, 1.0);

	prob_track_pip = new TH1D("prob_track_pip", "#pi^{+} probablity from tracking #chi^{2} and N_{dof}", 500, 0.0, 1.0);
	prob_track_pim = new TH1D("prob_track_pim", "#pi^{-} probablity from tracking #chi^{2} and N_{dof}", 500, 0.0, 1.0);

	//japp->RootUnLock();
	pthread_rwlock_unlock(&root_lock);
	
	return NOERROR;
}

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

//------------------
// evnt
//------------------
jerror_t JEventProcessor_pip_pim::evnt(JEventLoop *loop, int eventnumber)
{
	// Get reconstructed pi+ and pi- tracks
	vector<const DPiPlus*> pipluses;
	vector<const DPiMinus*> piminuses;
	vector<const DMCThrown*> throwns;
	loop->Get(pipluses);
	loop->Get(piminuses);
	loop->Get(throwns);

	//japp->RootWriteLock();
	pthread_rwlock_wrlock(&root_lock);

	//----------- Reconstructed --------------
	// Convert to DLorentzVector objects
	vector<DLorentzVector> pip;
	vector<DLorentzVector> pim;
	for(unsigned int i=0; i<pipluses.size(); i++){
		double prob = pipluses[i]->GetProb();
		double prob_track = pipluses[i]->GetProbTrack();
		prob_pip->Fill(prob);
		prob_track_pip->Fill(prob_track);
		if(prob_track < 0.02)continue; // Cut out poorly reconstructed tracks
		*v4_pip = pipluses[i]->lorentzMomentum();
		pip.push_back(*v4_pip);
		pip_tree->Fill();
	}
	for(unsigned int i=0; i<piminuses.size(); i++){
		double prob = piminuses[i]->GetProb();
		double prob_track = piminuses[i]->GetProbTrack();
		prob_pim->Fill(prob);
		prob_track_pim->Fill(prob_track);
		if(prob_track < 0.02)continue; // Cut out poorly reconstructed tracks
		*v4_pim = piminuses[i]->lorentzMomentum();
		pim.push_back(*v4_pim);
		pim_tree->Fill();
	}

	// Fill invariant mass histogram
	FillInvariantMassHisto(mass_pip_pim, pip, pim);
	

	//----------- Generated --------------
	// Convert to DLorentzVector objects
	pip.clear();
	pim.clear();
	for(unsigned int i=0; i<throwns.size(); i++){
		if(throwns[i]->type==8){
			*v4_pip_gen = throwns[i]->lorentzMomentum();
			pip.push_back(*v4_pip_gen);
			pip_tree_gen->Fill();
		}else if(throwns[i]->type==9){
			*v4_pim_gen = throwns[i]->lorentzMomentum();
			pim.push_back(*v4_pim_gen);
			pim_tree_gen->Fill();
		}
	}

	// Fill invariant mass histogram
	FillInvariantMassHisto(mass_pip_pim_gen, pip, pim);

	// Unlock ROOT read-write lock
	//japp->RootUnLock();
	pthread_rwlock_unlock(&root_lock);

	return NOERROR;
}

//------------------
// FillInvariantMassHisto
//------------------
void JEventProcessor_pip_pim::FillInvariantMassHisto(TH1D *h, vector<DLorentzVector> &pip, vector<DLorentzVector> &pim)
{
	// Loop over all combinations of pi+ + pi-
	for(unsigned int i=0; i<pip.size(); i++){
		for(unsigned int j=0; j<pim.size(); j++){
			DLorentzVector sum = pip[i] + pim[j]; // add 4-vectors
			double mass = sum.M();
			h->Fill(mass);
		}
	}
}


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