// $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)
{
	trk_prob_cut = 0.02;
	gPARMS->SetDefaultParameter("trk_prob_cut", trk_prob_cut, "Reconstructed track probability cut used in pip_pim plugin");

	//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", "Reconstructed pi^{+}");
	pim_tree = new TTree("pim_tree", "Reconstructed pi^{-}");
	pip_tree_gen = new TTree("pip_tree_gen", "Generated pi^{+}");
	pim_tree_gen = new TTree("pim_tree_gen", "Generated pi^{-}");
	recon_tree = new TTree("recon_tree", "Reconstructed pi^{+} 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);
	recon_tree->Branch("T", &pip_pim_vals, "event/I:Npairs:pip_px/F:pip_py:pip_pz:pim_px:pim_py:pim_pz:gen_pip_px/F:gen_pip_py:gen_pip_pz:gen_pim_px:gen_pim_py:gen_pim_pz");

	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);

	bad_events = new ofstream("bad_events.txt");

	//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);

	//----------- Generated --------------
	// Convert to DLorentzVector objects
	vector<DLorentzVector> pip;
	vector<DLorentzVector> pim;
	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);

	//----------- Reconstructed --------------
	// Convert to DLorentzVector objects
	pip.clear();
	pim.clear();
	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 < trk_prob_cut)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 < trk_prob_cut)continue; // Cut out poorly reconstructed tracks
		*v4_pim = piminuses[i]->lorentzMomentum();
		pim.push_back(*v4_pim);
		pim_tree->Fill();
	}
	
	// Fill recon tree for all combinations of 1 pi+ and 1 pi-
	pip_pim_vals.event = eventnumber;
	pip_pim_vals.Npairs = pip.size()*pim.size();
	pip_pim_vals.gen_pip_px = v4_pip_gen->X();
	pip_pim_vals.gen_pip_py = v4_pip_gen->Y();
	pip_pim_vals.gen_pip_pz = v4_pip_gen->Z();
	pip_pim_vals.gen_pim_px = v4_pim_gen->X();
	pip_pim_vals.gen_pim_py = v4_pim_gen->Y();
	pip_pim_vals.gen_pim_pz = v4_pim_gen->Z();
	for(unsigned int i=0; i<pip.size(); i++){

		pip_pim_vals.pip_px = pip[i].Px();
		pip_pim_vals.pip_py = pip[i].Py();
		pip_pim_vals.pip_pz = pip[i].Pz();

		for(unsigned int j=0; j<pim.size(); j++){
	
			pip_pim_vals.pim_px = pim[j].Px();
			pip_pim_vals.pim_py = pim[j].Py();
			pip_pim_vals.pim_pz = pim[j].Pz();

			recon_tree->Fill();
		}
	}

	// Fill invariant mass histogram
	bool one_pair_in_range = FillInvariantMassHisto(mass_pip_pim, pip, pim);
	if(!one_pair_in_range){
		(*bad_events) << eventnumber << endl;
	}

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

	return NOERROR;
}

//------------------
// FillInvariantMassHisto
//------------------
bool JEventProcessor_pip_pim::FillInvariantMassHisto(TH1D *h, vector<DLorentzVector> &pip, vector<DLorentzVector> &pim)
{
	// Loop over all combinations of pi+ + pi-
	bool one_pair_in_range = false;
	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);
			
			if(mass>0.2 && mass<0.5)one_pair_in_range=true;
		}
	}
	
	return one_pair_in_range;
}


//------------------
// 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.
	
	if(bad_events){
		bad_events->close();
		delete bad_events;
	}
	
	return NOERROR;
}