// $Id$
//
//    File: DCustomAction_MCSelector3pip.cc
// Created: Tue Sep 16 09:42:07 EDT 2014
// Creator: mstaib (on Linux maria 2.6.32-431.20.3.el6.x86_64 x86_64)
//

#include "DCustomAction_MCSelector3pip.h"

void DCustomAction_MCSelector3pip::Initialize(JEventLoop* locEventLoop)
{
	
	//Optional: Create histograms and/or modify member variables.
	//Create any histograms/trees/etc. within a ROOT lock. 
		//This is so that when running multithreaded, only one thread is writing to the ROOT file at a time. 

	japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
	{
		// Optional: Useful utility functions.
		// locEventLoop->GetSingle(dAnalysisUtilities);

		//Required: Create a folder in the ROOT output file that will contain all of the output ROOT objects (if any) for this action.
			//If another thread has already created the folder, it just changes to it. 
		CreateAndChangeTo_ActionDirectory();
		// Some constants for the histograms
		Int_t nBinInvMass = 1000; Float_t xMinInvMass = 0.0, xMaxInvMass = 8.0;
		Int_t nBinMissingMass = 1000; Float_t xMinMissingMass = -0.4, xMaxMissingMass = 0.4;
		Int_t nBinInvMassPi0 = 100; Float_t xMinInvMassPi0 = 0.105, xMaxInvMassPi0 = 0.165;
		Int_t nBinKinFitFOM = 500;
		Int_t nBinExtraPhotons = 10;
		Int_t nBinExtraCharged = 10;

		// Optional: Create a ROOT subfolder.
			//If another thread has already created the folder, it just changes to it. 
		CreateAndChangeTo_Directory("MCCorrectThrown", "MCCorrectThrown");
			//make sub-directory content here
		if(gDirectory->Get("hInvMassFullCombo_Correct") == NULL)
		hInvMassFullCombo_Correct = new TH1I("hInvMassFullCombo_Correct","Total Meson Invariant Mass; Inv. Mass [GeV/c^{2}];", nBinInvMass, xMinInvMass, xMaxInvMass);
		if(gDirectory->Get("hMissingMassSquared_Correct") == NULL)
                hMissingMassSquared_Correct = new TH1I("hMissingMassSquared_Correct","Missing Mass Squared; MM^{2} [[GeV^2/c^{4}];",nBinMissingMass, xMinMissingMass, xMaxMissingMass);
		if(gDirectory->Get("hInvMassPi0_Correct") == NULL)
                hInvMassPi0_Correct = new TH1I ("hInvMassPi0_Correct","#pi^{0} Invariant Mass; Inv. Mass [GeV/c^{2}];",nBinInvMassPi0, xMinInvMassPi0,xMaxInvMassPi0);
		if(gDirectory->Get("hKinFitFOM_Correct") == NULL)
                hKinFitFOM_Correct = new TH1I ("hKinFitFOM_Correct", "Kinematic Fit FOM", nBinKinFitFOM, 0.0, 1.0);
		if(gDirectory->Get("hNumExtraPhotons_Correct") == NULL)
                hNumExtraPhotons_Correct = new TH1I ( "hNumExtraPhotons_Correct", "Number of Extra #gamma", nBinExtraPhotons + 1, -0.5, (Float_t) nBinExtraPhotons + 0.5);
		if(gDirectory->Get("hNumExtraChargedTracks_Correct") == NULL)
                hNumExtraChargedTracks_Correct = new TH1I("hNumExtraChargedTracks_Correct", "Number of Extra Charged Tracks", nBinExtraCharged + 1, -0.5, (Float_t) nBinExtraCharged+0.5);

		gDirectory->cd(".."); //return to the action directory
		CreateAndChangeTo_Directory("MCIncorrectThrown", "MCIncorrectThrown");

		if(gDirectory->Get("hInvMassFullCombo_Incorrect") == NULL)
		hInvMassFullCombo_Incorrect = new TH1I("hInvMassFullCombo_Incorrect","Total Meson Invariant Mass; Inv. Mass [GeV/c^{2}];", nBinInvMass, xMinInvMass, xMaxInvMass);
		if(gDirectory->Get("hMissingMassSquared_Incorrect") == NULL)
                hMissingMassSquared_Incorrect = new TH1I("hMissingMassSquared_Incorrect","Missing Mass Squared; MM^{2} [[GeV^2/c^{4}];",nBinMissingMass, xMinMissingMass, xMaxMissingMass);
		if(gDirectory->Get("hInvMassPi0_Incorrect") == NULL)
                hInvMassPi0_Incorrect = new TH1I ("hInvMassPi0_Incorrect","#pi^{0} Invariant Mass; Inv. Mass [GeV/c^{2}];",nBinInvMassPi0, xMinInvMassPi0,xMaxInvMassPi0);
		if(gDirectory->Get("hKinFitFOM_Incorrect") == NULL)
                hKinFitFOM_Incorrect = new TH1I ("hKinFitFOM_Incorrect", "Kinematic Fit FOM", nBinKinFitFOM, 0.0, 1.0);
		if(gDirectory->Get("hNumExtraPhotons_Incorrect") == NULL)
                hNumExtraPhotons_Incorrect = new TH1I ( "hNumExtraPhotons_Incorrect", "Number of Extra #gamma", nBinExtraPhotons + 1, -0.5 , (Float_t) nBinExtraPhotons + 0.5);
		if(gDirectory->Get("hNumExtraChargedTracks_Incorrect") == NULL)
                hNumExtraChargedTracks_Incorrect = new TH1I("hNumExtraChargedTracks_Incorrect", "Number of Extra Charged Tracks", nBinExtraCharged + 1, -0.5 , (Float_t) nBinExtraCharged + 0.5);

		if(gDirectory->Get("nTupleIncorrectThrown") == NULL)
		hNTupleIncorrect = new TNtuple("nTupleIncorrectThrown", "nTupleIncorrectThrown", "numPiPlus:numPiMinus:numPi0:numKPlus:numKMinus:numK0:numGamma:numProton:numNeutron");
		

		//	(Optional) Example: Create a histogram. 
		//	if(gDirectory->Get("MyHistName") == NULL) //check to see if already created by another thread
		//		dMyHist = new TH1D("MyHistName", "MyHistTitle", 100, 0.0, 1.0);
		//	else //already created by another thread
		//		dMyHist = static_cast<TH1D*>(gDirectory->Get("MyHistName"));
	}
	japp->RootUnLock(); //RELEASE ROOT LOCK!!
	
}

bool DCustomAction_MCSelector3pip::Perform_Action(JEventLoop* locEventLoop, const DParticleCombo* locParticleCombo)
{
	//Write custom code to perform an action on the INPUT DParticleCombo (DParticleCombo)
	//NEVER: Grab DParticleCombo or DAnalysisResults objects (of any tag!) from the JEventLoop within this function
	//NEVER: Grab objects that are created post-kinfit (e.g. DKinFitResults, etc.) from the JEventLoop if Get_UseKinFitResultsFlag() == false: CAN CAUSE INFINITE DEPENDENCY LOOP


	//Optional: check whether the user wanted to use the kinematic fit results when performing this action
	// bool locUseKinFitResultsFlag = Get_UseKinFitResultsFlag();
	vector<const DMCThrown*> locFinalStateParticles;
	vector<const DMCThrown*> locDecayingParticles;

	locEventLoop->Get(locFinalStateParticles, "FinalState");
	locEventLoop->Get(locDecayingParticles, "Decaying");

	Int_t numPiPlus = 0, numPiMinus = 0, numPi0 = 0, numKPlus = 0 , numKMinus = 0, numK0 = 0, numGamma = 0, numProton = 0, numNeutron = 0;

	//if (locFinalStateParticles.size() != 7) return false; //Needs to have the exact number of particles in the final state
	// Now that we have the right number, count them.
	for (unsigned int i = 0; i < locFinalStateParticles.size(); i++){
		const DMCThrown* locMCThrown = locFinalStateParticles[i];
		int pdgType = locMCThrown->pdgtype;
		if (pdgType == 211) numPiPlus++;
		else if (pdgType == -211) numPiMinus++;
		else if (pdgType == 2212) numProton++;
		else if (pdgType == 22) numGamma++;
		else if (pdgType == 321) numKPlus++;
		else if (pdgType == -321) numKMinus++;
		else if (pdgType == 130) numK0++;
		else if (pdgType == 2112) numNeutron++;
	}
	// This still leaves an ambiguity about whether the photons actually came from a pi0 so lets count those
	
	for (unsigned int i = 0; i < locDecayingParticles.size(); i++){
		const DMCThrown* locMCThrown = locDecayingParticles[i];
                int pdgType = locMCThrown->pdgtype;
                if (pdgType == 111) numPi0++;
	}
	
	//Now to check if we have what we want and retun
	if (numPiPlus == 1 && numPiMinus == 1 && numProton == 1 && numGamma == 2 && numPi0 == 1 && locFinalStateParticles.size() == 5){
		FillCorrectMCHistograms(locEventLoop, locParticleCombo);
		return true;
	}
	else{
		//numPiPlus:numPiMinus:numPi0:numKPlus:numKMinus:numK0:numGamma:numProton:numNeutron
		japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
		{
		numGamma = numGamma - 2 * numPi0;
		hNTupleIncorrect->Fill(numPiPlus,numPiMinus,numPi0,numKPlus,numKMinus,numK0,numGamma,numProton,numNeutron);
		}
		japp->RootUnLock(); //RELEASE ROOT LOCK!!
		
		FillIncorrectMCHistograms(locEventLoop, locParticleCombo);		
		return false;
	}
	/*
	//Optional: Fill histograms
	japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
	{
		// Fill any histograms here
	}
	japp->RootUnLock(); //RELEASE ROOT LOCK!!
	*/

	//return true; //return false if you want to use this action to apply a cut (and it fails the cut!)
}

void DCustomAction_MCSelector3pip::FillCorrectMCHistograms(JEventLoop* locEventLoop, const DParticleCombo* locParticleCombo)
{

	DLorentzVector locMissingP4 = dAnalysisUtilities->Calc_MissingP4(locParticleCombo, false);
	DLorentzVector locTotalFinalP4 = dAnalysisUtilities->Calc_FinalStateP4(locParticleCombo, 0, true);
	DLorentzVector locPi0P4 = dAnalysisUtilities->Calc_FinalStateP4(locParticleCombo, 1, true);
	double locKinFitFOM = locParticleCombo->Get_KinFitResults()->Get_ConfidenceLevel();
	
	//The last two require a loop over the charged and neutrals
	//Define the number of charged actually in the event
	Int_t numCharged = 3, numNeutral = 2;
	
	//Now count up
	Int_t numExtraCharged = 0, numExtraNeutral = 0;
	vector<const DChargedTrack*> locChargedTracksVector;
	vector<const DNeutralParticle*> locNeutralParticleVector;

	locEventLoop->Get(locChargedTracksVector);
	locEventLoop->Get(locNeutralParticleVector);

	for(size_t i = 0; i < locChargedTracksVector.size(); i++){
		double locFOM = locChargedTracksVector[i]->Get_BestFOM()->dFOM;
		if (locFOM > 5.73303E-7) numExtraCharged++;
	}
	numExtraCharged-=numCharged;

	for(size_t j = 0; j < locNeutralParticleVector.size(); j++){
		double locFOM = locNeutralParticleVector[j]->Get_BestFOM()->dFOM;
		if (locFOM > 5.73303E-7) numExtraNeutral++;
	}
	numExtraNeutral-=numNeutral;
	
	japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
	{
		hInvMassFullCombo_Correct->Fill(locTotalFinalP4.M());
		hMissingMassSquared_Correct->Fill(locMissingP4.M2());
		hInvMassPi0_Correct->Fill(locPi0P4.M());
		hKinFitFOM_Correct->Fill(locKinFitFOM);
		hNumExtraPhotons_Correct->Fill(numExtraNeutral);
		hNumExtraChargedTracks_Correct->Fill(numExtraCharged);
	}
	japp->RootUnLock(); //RELEASE ROOT LOCK!!
}

void DCustomAction_MCSelector3pip::FillIncorrectMCHistograms(JEventLoop* locEventLoop, const DParticleCombo* locParticleCombo)
{

	DLorentzVector locMissingP4 = dAnalysisUtilities->Calc_MissingP4(locParticleCombo, false);
	//DLorentzVector locTotalFinalP4 = dAnalysisUtilities->Calc_FinalStateP4(locParticleCombo, 0, true);
	set<pair<const JObject*, Particle_t> > locSourceObjects;
	DLorentzVector locTotalFinalP4 = Calc_TotalMesonP4(locParticleCombo, 0, locSourceObjects, true);
	DLorentzVector locPi0P4 = dAnalysisUtilities->Calc_FinalStateP4(locParticleCombo, 1, true);
	double locKinFitFOM = locParticleCombo->Get_KinFitResults()->Get_ConfidenceLevel();

	//The last two require a loop over the charged and neutrals
	//Define the number of charged actually in the event
	Int_t numCharged = 3, numNeutral = 2;
	
	//Now count up
	Int_t numExtraCharged = 0, numExtraNeutral = 0;
	vector<const DChargedTrack*> locChargedTracksVector;
	vector<const DNeutralParticle*> locNeutralParticleVector;

	locEventLoop->Get(locChargedTracksVector);
	locEventLoop->Get(locNeutralParticleVector);

	for(size_t i = 0; i < locChargedTracksVector.size(); i++){
		double locFOM = locChargedTracksVector[i]->Get_BestFOM()->dFOM;
		if (locFOM > 5.73303E-7) numExtraCharged++;
	}
	numExtraCharged-=numCharged;

	for(size_t j = 0; j < locNeutralParticleVector.size(); j++){
		double locFOM = locNeutralParticleVector[j]->Get_BestFOM()->dFOM;
		if (locFOM > 5.73303E-7) numExtraNeutral++;
	}
	numExtraNeutral-=numNeutral;

	japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
	{
		hInvMassFullCombo_Incorrect->Fill(locTotalFinalP4.M());
		hMissingMassSquared_Incorrect->Fill(locMissingP4.M2());
		hInvMassPi0_Incorrect->Fill(locPi0P4.M());
		hKinFitFOM_Incorrect->Fill(locKinFitFOM);
		hNumExtraPhotons_Incorrect->Fill(numExtraNeutral);
		hNumExtraChargedTracks_Incorrect->Fill(numExtraCharged);
	}
	japp->RootUnLock(); //RELEASE ROOT LOCK!!
}

DLorentzVector DCustomAction_MCSelector3pip::Calc_TotalMesonP4(const DParticleCombo* locParticleCombo, size_t locStepIndex, set<pair<const JObject*, Particle_t> >& locSourceObjects, bool locUseKinFitDataFlag) const
 {
    if(locUseKinFitDataFlag && (locParticleCombo->Get_KinFitResults() == NULL))
       return Calc_TotalMesonP4(locParticleCombo, locStepIndex, locSourceObjects, false); //kinematic fit failed
 
    DLorentzVector locFinalStateP4;
    const DParticleComboStep* locParticleComboStep = locParticleCombo->Get_ParticleComboStep(locStepIndex);
    if(locParticleComboStep == NULL)
       return (DLorentzVector());
 
    deque<const DKinematicData*> locParticles;
    if(!locUseKinFitDataFlag) //measured
       locParticleComboStep->Get_FinalParticles_Measured(locParticles);
    else //kinfit
       locParticleComboStep->Get_FinalParticles(locParticles);
 
    for(size_t loc_i = 0; loc_i < locParticles.size(); ++loc_i)
    {
       if(locParticleComboStep->Is_FinalParticleMissing(loc_i))
          return (DLorentzVector()); //bad!
       else if (locParticleComboStep->Get_FinalParticleID(loc_i) == Proton || locParticleComboStep->Get_FinalParticleID(loc_i) == Neutron)
	  return (DLorentzVector());
       if(locParticleComboStep->Is_FinalParticleDecaying(loc_i))
       {
          //measured results, or not constrained by kinfit (either non-fixed mass or excluded from kinfit)
          if((!locUseKinFitDataFlag) || (!IsFixedMass(locParticleComboStep->Get_FinalParticleID(loc_i))))
             locFinalStateP4 += Calc_TotalMesonP4(locParticleCombo, locParticleComboStep->Get_DecayStepIndex(loc_i), locSourceObjects, locUseKinFitDataFlag);
          else //want kinfit results, and decaying particle p4 is constrained by kinfit
          {
             locFinalStateP4 += locParticles[loc_i]->lorentzMomentum();
             //still need source objects of decay products! dive down anyway, but ignore p4 result
             Calc_TotalMesonP4(locParticleCombo, locParticleComboStep->Get_DecayStepIndex(loc_i), locSourceObjects, locUseKinFitDataFlag);
          }
       }
       else
       {
          locFinalStateP4 += locParticles[loc_i]->lorentzMomentum();
          locSourceObjects.insert(pair<const JObject*, Particle_t>(locParticleComboStep->Get_FinalParticle_SourceObject(loc_i), locParticles[loc_i]->PID()));
       }
    }
    return locFinalStateP4;
 }