// $Id$
//
//    File: JEventProcessor_CPPMVA.cc
// Created: Sat Jul  2 10:49:58 EDT 2016
// Creator: davidl (on Darwin harriet 15.5.0 Darwin Kernel Version 15.5.0)
//


#include "JEventProcessor_CPPMVA.h"
using namespace jana;

#include <TMath.h>
#include <particleType.h>
#include <TRACKING/DMCThrown.h>
#include <TRACKING/DTrackTimeBased.h>
#include <PID/DNeutralShower.h>
#include <PID/DChargedTrack.h>
#include <PID/DBeamPhoton.h>
#include <PID/DEventRFBunch.h>
#include <TOF/DTOFHit.h>
#include <FCAL/DFCALHit.h>
#include <FCAL/DFCALCluster.h>
#include <FMWPC/DFMWPCHit.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_CPPMVA());
}
} // "C"


//------------------
// JEventProcessor_CPPMVA (Constructor)
//------------------
JEventProcessor_CPPMVA::JEventProcessor_CPPMVA()
{

}

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

}

//------------------
// init
//------------------
jerror_t JEventProcessor_CPPMVA::init(void)
{

	cppmva = new TTree("cppmva", "Tree for MVA to classify pi/mu events for CPP experiment");

	cppmva->Branch("Ntracks",              &mva.Ntracks,              "Ntracks/I");
	cppmva->Branch("Ntof",                 &mva.Ntof,                 "Ntof/I");
	cppmva->Branch("Nfcal_hits",           &mva.Nfcal_hits,           "Nfcal_hits/I");
	cppmva->Branch("Nfcal_clusters",       &mva.Nfcal_clusters,       "Nfcal_clusters/I");
	cppmva->Branch("Efcal_clusters",       &mva.Efcal_clusters,       "Efcal_clusters/F");
	cppmva->Branch("Efcal_hits",           &mva.Efcal_hits,           "Efcal_hits/F"); 
	cppmva->Branch("Nfmwpc",               &mva.Nfmwpc,               "Nfmwpc/I");
	cppmva->Branch("Nfmwpc1",              &mva.Nfmwpc1,              "Nfmwpc1/I");
	cppmva->Branch("Nfmwpc2",              &mva.Nfmwpc2,              "Nfmwpc2/I");
	cppmva->Branch("Nfmwpc3",              &mva.Nfmwpc3,              "Nfmwpc3/I");
	cppmva->Branch("Nfmwpc4",              &mva.Nfmwpc4,              "Nfmwpc4/I");
	cppmva->Branch("Nfmwpc5",      	      &mva.Nfmwpc5,              "Nfmwpc5/I");
	cppmva->Branch("Nfmwpc6",              &mva.Nfmwpc6,              "Nfmwpc6/I");
	cppmva->Branch("Nfmwpc7",     	      &mva.Nfmwpc7,              "Nfmwpc7/I");
	cppmva->Branch("Nfmwpc8",              &mva.Nfmwpc8,              "Nfmwpc8/I");
	cppmva->Branch("Nfmwpc9",      	      &mva.Nfmwpc9,              "Nfmwpc9/I");
	cppmva->Branch("Nfmwpc10",             &mva.Nfmwpc10,             "Nfmwpc10/I");
	cppmva->Branch("Nfmwpc11",     	      &mva.Nfmwpc11,             "Nfmwpc11/I");
	cppmva->Branch("Nfmwpc12",             &mva.Nfmwpc12,             "Nfmwpc12/I");
	cppmva->Branch("Nfmwpc_with_2hits",    &mva.Nfmwpc_with_2hits,    "Nfmwpc_with_2hits/I");
	cppmva->Branch("Nfmwpc_with_lt_2hits", &mva.Nfmwpc_with_lt_2hits, "Nfmwpc_with_lt_2hits/I");
	cppmva->Branch("Nfmwpc_with_gt_2hits", &mva.Nfmwpc_with_gt_2hits, "Nfmwpc_with_gt_2hits/I");

	cppmva->Branch("pip_theta",   	      &mva.pip_theta,   	      "pip_theta/F");
	cppmva->Branch("pim_theta",            &mva.pim_theta,   	      "pim_theta/F");
	cppmva->Branch("pip_E",                &mva.pip_E,   	            "pip_E/F");
	cppmva->Branch("pim_E",   	            &mva.pim_E,   	            "pim_E/F");
	cppmva->Branch("pip_theta_Thrown",     &mva.pip_theta_Thrown,     "pip_theta_Thrown/F");
	cppmva->Branch("pim_theta_Thrown",     &mva.pim_theta_Thrown,     "pim_theta_Thrown/F");
	cppmva->Branch("pip_E_Thrown",         &mva.pip_E_Thrown,         "pip_E_Thrown/F");
	cppmva->Branch("pim_E_Thrown",         &mva.pim_E_Thrown,         "pim_E_Thrown/F");

	cppmva->Branch("is_pion",              &mva.is_pion,              "is_pion/I");

	return NOERROR;
}

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

//------------------
// evnt
//------------------
jerror_t JEventProcessor_CPPMVA::evnt(JEventLoop *loop, uint64_t eventnumber)
{
	vector<const DMCThrown*>       mcthrowns;
	vector<const DChargedTrack*>   cts;
	vector<const DTrackTimeBased*> tbts;
	vector<const DNeutralShower*>  nss;
	vector<const DTOFHit*>         tofhits;
	vector<const DFCALHit*>        fcalhits;
	vector<const DFCALCluster*>    fcalclusters;
	vector<const DFMWPCHit*>       fmwpchits;

	loop->Get(mcthrowns);
	loop->Get(cts);
	loop->Get(tbts);
	loop->Get(nss);
	loop->Get(tofhits);
	loop->Get(fcalhits);
	loop->Get(fcalclusters);
	loop->Get(fmwpchits);
	
	double Pi = 3.141592626358979; //pi

	// Pion or muon event?
	Int_t is_pion = false;
	for(auto mcthrown : mcthrowns) if(mcthrown->PID()==PiPlus) is_pion = true;

	// Visible energy
	double Evisible = 0.0;
	for(auto ct : cts){
		const DChargedTrackHypothesis *cth = ct->Get_BestFOM();
		double confidence_level = TMath::Prob((double)cth->dChiSq, (double)cth->dNDF);
		if(confidence_level > 0.0001){
			Evisible += cth->energy();
			
			// Do not include proton mass in visible energy
			if(fabs(cth->mass()-0.938) < 0.010) Evisible -= cth->mass();
		}
	}
	for(auto ns : nss) Evisible += ns->dEnergy;
	
	// Calorimeter energy
	double Efcal_clusters = 0.0;
	for(auto fc : fcalclusters) Efcal_clusters += fc->getEnergy();



	// Calorimeter energy - not clustered
	double Efcal_hits = 0.0;
	for(auto efc : fcalhits) Efcal_hits += efc->E; //MeV

	// Only count TOF hits with dE>1MeV
	Int_t Ntof = 0;
	for(auto tofhit : tofhits) if(tofhit->dE >= 0.001) Ntof++;

	// FMWPC
	vector<Int_t> Nfmwpc(8, 0);
	for(auto h : fmwpchits){
		if(h->layer>0 && h->layer<=(int)Nfmwpc.size());
		Nfmwpc[h->layer-1]++;
	}

	Int_t   Nfmwpc_with_2hits = 0;
	Int_t   Nfmwpc_with_lt_2hits = 0;
	Int_t   Nfmwpc_with_gt_2hits = 0;
	
	for(auto N : Nfmwpc){
		if( N<2  ) Nfmwpc_with_lt_2hits++;
		if( N==2 ) Nfmwpc_with_2hits++;
		if( N>2  ) Nfmwpc_with_gt_2hits++;
	}

	const DMCThrown *pip_Thrown = NULL;
	const DMCThrown *pim_Thrown = NULL;
	for(auto s : mcthrowns){
		if(s->PID()==PiPlus || s->PID()==MuonPlus) pip_Thrown =s;
		if(s->PID()==PiMinus || s->PID()==MuonMinus) pim_Thrown =s;	
	}
	
	if(pip_Thrown==NULL || pim_Thrown==NULL) return NOERROR;

	double pip_theta_Thrown = pip_Thrown->momentum().Theta()*180/Pi;
	double pim_theta_Thrown = pim_Thrown->momentum().Theta()*180/Pi;
	double pip_E_Thrown = pip_Thrown->energy();
	double pim_E_Thrown = pim_Thrown->energy();

	const DTrackTimeBased *pip = NULL;
	const DTrackTimeBased *pim = NULL;
	for(auto t : tbts){
		if( fabs(t->mass()-0.139) > 0.01 ) continue;
		if(t->charge() == +1.0) pip = t;
		if(t->charge() == -1.0) pim = t;
	}
	
	if(pip==NULL || pim==NULL) return NOERROR;

	double pip_theta = pip->momentum().Theta()*180/Pi;
	double pim_theta = pim->momentum().Theta()*180/Pi;
	double pip_E = pip->energy();
	double pim_E = pim->energy();

	
	
	
	japp->RootWriteLock();

	mva.Ntracks        	= cts.size();
	mva.Ntof           	= Ntof;

	mva.Nfcal_hits     	= fcalhits.size();
	mva.Nfcal_clusters 	= fcalclusters.size();
	mva.Efcal_clusters 	= Efcal_clusters;
	mva.Efcal_hits		= Efcal_hits;
	mva.pip_theta	   	= pip_theta;
	mva.pim_theta	   	= pim_theta; 		
	mva.pip_E	   	= pip_E;	
	mva.pim_E	   	= pim_E;

	mva.pip_theta_Thrown	= pip_theta_Thrown;
	mva.pim_theta_Thrown	= pim_theta_Thrown; 		
	mva.pip_E_Thrown	= pip_E_Thrown;	
	mva.pim_E_Thrown	= pim_E_Thrown;

	mva.Nfmwpc         	= fmwpchits.size();
	mva.Nfmwpc1        	= Nfmwpc[ 0];
	mva.Nfmwpc2        	= Nfmwpc[ 1];
	mva.Nfmwpc3        	= Nfmwpc[ 2];
	mva.Nfmwpc4        	= Nfmwpc[ 3];
	mva.Nfmwpc5        	= Nfmwpc[ 4];
	mva.Nfmwpc6        	= Nfmwpc[ 5];
	mva.Nfmwpc7        	= Nfmwpc[ 6];
	mva.Nfmwpc8        	= Nfmwpc[ 7];
// 	mva.Nfmwpc9        	= Nfmwpc[ 8];
// 	mva.Nfmwpc10        	= Nfmwpc[ 9];
// 	mva.Nfmwpc11        	= Nfmwpc[10];
// 	mva.Nfmwpc12        	= Nfmwpc[11];
	mva.Nfmwpc_with_2hits    = Nfmwpc_with_2hits;
	mva.Nfmwpc_with_lt_2hits = Nfmwpc_with_lt_2hits;
	mva.Nfmwpc_with_gt_2hits = Nfmwpc_with_gt_2hits;
	mva.is_pion = is_pion ? 1:0;

	cppmva->Fill();

	japp->RootUnLock();


	return NOERROR;
}

//------------------
// erun
//------------------
jerror_t JEventProcessor_CPPMVA::erun(void)
{
	japp->RootWriteLock();
	cppmva->FlushBaskets();
	japp->RootUnLock();

	return NOERROR;
}

//------------------
// fini
//------------------
jerror_t JEventProcessor_CPPMVA::fini(void)
{
	japp->RootWriteLock();
	cppmva->FlushBaskets();
	japp->RootUnLock();
		
	return NOERROR;
}