// $Id$
//
//    File: DCustomAction_InvariantMassDifference.cc
// Created: Fri Jul  7 15:51:45 EDT 2017
// Creator: aaustreg (on Linux ifarm1402.jlab.org 3.10.0-327.el7.x86_64 x86_64)
//

#include "DCustomAction_InvariantMassDifference.h"

void DCustomAction_InvariantMassDifference::Initialize(JEventLoop* locEventLoop)
{
	string locHistName, locHistTitle;
	double locMassPerBin = 1000.0*(dMaxMass - dMinMass)/((double)dNumMassBins);
	locEventLoop->GetSingle(dAnalysisUtilities);

	string locMotherParticleNamesForHist = "";
	if(dMother != Unknown)
	{
		auto locChainPIDs = DAnalysis::Get_ChainPIDs(Get_Reaction(), dMother, !Get_UseKinFitResultsFlag());
		locMotherParticleNamesForHist = DAnalysis::Convert_PIDsToROOTName(locChainPIDs);
	}
	else
	{
	    for(size_t loc_i = 0; loc_i < dToIncludePIDs_Mother.size(); ++loc_i)
	      locMotherParticleNamesForHist += ParticleName_ROOT(dToIncludePIDs_Mother[loc_i]);
	}
	
	string locDaughterParticleNamesForHist = "";
	if(dDaughter != Unknown)
	{
		auto locChainPIDs = DAnalysis::Get_ChainPIDs(Get_Reaction(), dDaughter, !Get_UseKinFitResultsFlag());
		locDaughterParticleNamesForHist = DAnalysis::Convert_PIDsToROOTName(locChainPIDs);
	}
	else
	{
	    for(size_t loc_i = 0; loc_i < dToIncludePIDs_Daughter.size(); ++loc_i)
	      locDaughterParticleNamesForHist += ParticleName_ROOT(dToIncludePIDs_Daughter[loc_i]);
	}
	
	bool locBeamPresent = DAnalysis::Get_IsFirstStepBeam(Get_Reaction());

	//CREATE THE HISTOGRAMS
	//Since we are creating histograms, the contents of gDirectory will be modified: must use JANA-wide ROOT lock
	japp->RootWriteLock(); //ACQUIRE ROOT LOCK!!
	{
		CreateAndChangeTo_ActionDirectory();

		locHistName = "InvariantMassDifference";
		ostringstream locStream;
		locStream << locMassPerBin;
		locHistTitle = string(";") + locMotherParticleNamesForHist + string(" - ") + locDaughterParticleNamesForHist + string(" Invariant Mass Difference (GeV/c^{2});# Combos / ") + locStream.str() + string(" MeV/c^{2}");
		dHist_InvariantMassDifference = GetOrCreate_Histogram<TH1I>(locHistName, locHistTitle, dNumMassBins, dMinMass, dMaxMass);

		if(locBeamPresent)
		{
			locHistName = "InvariantMassDifferenceVsBeamE";
			locHistTitle = string(";Beam Energy (GeV);") + locMotherParticleNamesForHist + string(" - ") + locDaughterParticleNamesForHist + string(" Invariant Mass Difference (GeV/c^{2})");
			dHist_InvariantMassDifferenceVsBeamE = GetOrCreate_Histogram<TH2D>(locHistName, locHistTitle, dNum2DBeamEBins, dMinBeamE, dMaxBeamE, dNum2DMassBins, dMinMass, dMaxMass);
		}

		//Return to the base directory
		ChangeTo_BaseDirectory();
	}
	japp->RootUnLock(); //RELEASE ROOT LOCK!!
}

bool DCustomAction_InvariantMassDifference::Perform_Action(JEventLoop* locEventLoop, const DParticleCombo* locParticleCombo)
{
	auto locFirstStep = locParticleCombo->Get_ParticleComboStep(0);
	bool locBeamPresent = DAnalysis::Get_IsFirstStepBeam(Get_Reaction());
	auto locBeam = Get_UseKinFitResultsFlag() ? locFirstStep->Get_InitialParticle() : locFirstStep->Get_InitialParticle_Measured();

	vector<double> locMassDifferencesToFill;
	vector<double> loc2DMassDifferencesToFill;	
	for(size_t loc_i = 0; loc_i < locParticleCombo->Get_NumParticleComboSteps(); ++loc_i)
	{
	        //const DParticleComboStep* locParticleComboStep = locParticleCombo->Get_ParticleComboStep(loc_i);
		const DReactionStep* locReactionStep = Get_Reaction()->Get_ReactionStep(loc_i);
		// if((dInitialPID != Unknown) && (locParticleComboStep->Get_InitialParticleID() != dInitialPID))
		//   continue;
		if((dStepIndex != -1) && (int(loc_i) != dStepIndex))
		  continue;

		//build all possible combinations of the included pids
		set<set<size_t> > locIndexCombos_Mother = dAnalysisUtilities->Build_IndexCombos(locReactionStep, dToIncludePIDs_Mother);

		const DReactionStep* locReactionStepPlusOne = Get_Reaction()->Get_ReactionStep(loc_i+1);
		
		//build all possible combinations of the included pids
		set<set<size_t> > locIndexCombos_Daughter = dAnalysisUtilities->Build_IndexCombos(locReactionStepPlusOne, dToIncludePIDs_Daughter);
		
		//loop over mother
		for(auto& locMothers : locIndexCombos_Mother)
		  {
		        set<pair<const JObject*, unsigned int> > locSourceObjects_Mother;
			DLorentzVector locMotherP4 = dAnalysisUtilities->Calc_FinalStateP4(Get_Reaction(), locParticleCombo, loc_i, locMothers, locSourceObjects_Mother, Get_UseKinFitResultsFlag());

			//loop over daughter
			for(auto& locDaughters : locIndexCombos_Daughter)
			  {
			    //check if mother includes daughter
			    if(!std::includes(locMothers.begin(), locMothers.end(), locDaughters.begin(), locDaughters.end()))
			      continue;
			    
			    set<pair<const JObject*, unsigned int> > locSourceObjects_Daughter;
			    DLorentzVector locDaughterP4 = dAnalysisUtilities->Calc_FinalStateP4(Get_Reaction(), locParticleCombo, loc_i+1, locDaughters, locSourceObjects_Daughter, Get_UseKinFitResultsFlag());
				
			    auto locSourceObjects_Couple = std::make_pair(locSourceObjects_Mother, locSourceObjects_Daughter);
			    if(dPreviousSourceObjects_Couple.find(locSourceObjects_Couple) == dPreviousSourceObjects_Couple.end())
			      {
				dPreviousSourceObjects_Couple.insert(locSourceObjects_Couple);
				locMassDifferencesToFill.push_back(locMotherP4.M() - locDaughterP4.M());
			      }
			    
			    auto locBeamSourceObjects = std::make_tuple(locSourceObjects_Mother, locSourceObjects_Daughter, locBeam);
			    if(locBeamPresent && (dPreviousSourceObjects_Beam.find(locBeamSourceObjects) == dPreviousSourceObjects_Beam.end()))
			      {
				dPreviousSourceObjects_Beam.insert(locBeamSourceObjects);
				loc2DMassDifferencesToFill.push_back(locMotherP4.M() - locDaughterP4.M());
			      }
			  }
		  }
		//don't break: e.g. if multiple pi0's, histogram invariant mass of each one
	}

	//FILL HISTOGRAMS
	//Since we are filling histograms local to this action, it will not interfere with other ROOT operations: can use action-wide ROOT lock
	//Note, the mutex is unique to this DReaction + action_string combo: actions of same class with different hists will have a different mutex
	Lock_Action();
	{
		for(size_t loc_i = 0; loc_i < locMassDifferencesToFill.size(); ++loc_i)
			dHist_InvariantMassDifference->Fill(locMassDifferencesToFill[loc_i]);
		for(size_t loc_i = 0; loc_i < loc2DMassDifferencesToFill.size(); ++loc_i)
		  dHist_InvariantMassDifferenceVsBeamE->Fill(locBeam->energy(), loc2DMassDifferencesToFill[loc_i]);
	}
	Unlock_Action();

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