#include "DSelector_klambda.h"

void DSelector_klambda::Init(TTree *locTree)
{
	// The Init() function is called when the selector needs to initialize a new tree or chain.
	// Typically here the branch addresses and branch pointers of the tree will be set.
	// Init() will be called many times when running on PROOF (once per file to be processed).

	//SET OUTPUT FILE NAME //can be overriden by user in PROOF
	dOutputFileName = "klambda.root"; //"" for none
	dOutputTreeFileName = ""; //"" for none

	//DO THIS NEXT
	//Because this function gets called for each TTree in the TChain, we must be careful:
		//We need to re-initialize the tree interface & branch wrappers, but don't want to recreate histograms
	bool locInitializedPriorFlag = dInitializedFlag; //save whether have been initialized previously
	DSelector::Init(locTree); //This must be called to initialize wrappers for each new TTree
	//gDirectory now points to the output file with name dOutputFileName (if any)
	if(locInitializedPriorFlag)
		return; //have already created histograms, etc. below: exit

	//THEN THIS
	Get_ComboWrappers();

	/*********************************** EXAMPLE USER INITIALIZATION: ANALYSIS ACTIONS **********************************/

	//ANALYSIS ACTIONS: //Executed in order //false/true: use measured/kinfit data
	//PID
	dAnalysisActions.push_back(new DHistogramAction_ParticleID(dComboWrapper, false));
	dAnalysisActions.push_back(new DCutAction_dEdx(dComboWrapper, false, KPlus, SYS_CDC));
	dAnalysisActions.push_back(new DCutAction_dEdx(dComboWrapper, false, Proton, SYS_CDC));
	dAnalysisActions.push_back(new DCutAction_dEdx(dComboWrapper, false, KMinus, SYS_CDC));
	//below: value: +/- N ns, Unknown: All PIDs, SYS_NULL: all timing systems
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 0.4, KPlus, SYS_BCAL));
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 0.25, KPlus, SYS_TOF));
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 0.5, KPlus, SYS_FCAL));
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 1.0, Proton, SYS_FCAL));
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 0.8, Proton, SYS_BCAL));
	dAnalysisActions.push_back(new DCutAction_PIDDeltaT(dComboWrapper, false, 0.5, Proton, SYS_TOF));
	dAnalysisActions.push_back(new DHistogramAction_ParticleID(dComboWrapper, false, "PostPID"));

	//MASSES
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, false, Lambda, 1000, 1.0, 1.2, "Lambda"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, false, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, true, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_KinFit"));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, false));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, true, "KinFit"));
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, true, Lambda, 1000, 1.0, 1.2, "Lambda_KinFit"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMassSquared(dComboWrapper, false, 1000, -0.1, 0.1));

	//KINFIT RESULTS
	dAnalysisActions.push_back(new DHistogramAction_KinFitResults(dComboWrapper));
	dAnalysisActions.push_back(new DCutAction_KinFitFOM(dComboWrapper, 1.0E-40));

	//MASSES
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, false, Lambda, 1000, 1.0, 1.2, "Lambda_PostKinFitCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, false, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_PostKinFitCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, true, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_KinFit_PostKinFitCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, false, "PostKinFitCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, true, "PostKinFitCut_KinFit"));
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, true, Lambda, 1000, 1.0, 1.2, "Lambda_KinFit_PostKinFitCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMassSquared(dComboWrapper, false, 1000, -0.1, 0.1, "PostKinFitCut"));

	//CUT LAMBDA
	dAnalysisActions.push_back(new DCutAction_InvariantMass(dComboWrapper, false, Lambda, 1.095, 1.136));

	//MASSES
	dAnalysisActions.push_back(new DHistogramAction_MissingMassSquared(dComboWrapper, false, 1000, -0.1, 0.1, "PostLambdaCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, false, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_PostLambdaCut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, true, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_KinFit_PostKinFitCut"));

	//BEAM ENERGY
	dAnalysisActions.push_back(new DHistogramAction_BeamEnergy(dComboWrapper, false));
	dAnalysisActions.push_back(new DCutAction_BeamEnergy(dComboWrapper, false, 8.4, 9.05));

	//MASSES
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, false, Lambda, 1000, 1.0, 1.2, "Lambda_PostBeamECut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, false, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_PostBeamECut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMass(dComboWrapper, true, 0, KPlus, 1000, 0.6, 1.6, "OffKPlus_KinFit_PostBeamECut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, false, "PostBeamECut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingP4(dComboWrapper, true, "PostBeamECut_KinFit"));
	dAnalysisActions.push_back(new DHistogramAction_InvariantMass(dComboWrapper, true, Lambda, 1000, 1.0, 1.2, "Lambda_KinFit_PostBeamECut"));
	dAnalysisActions.push_back(new DHistogramAction_MissingMassSquared(dComboWrapper, false, 1000, -0.1, 0.1, "PostBeamECut"));

	//KINEMATICS
	dAnalysisActions.push_back(new DHistogramAction_ParticleComboKinematics(dComboWrapper, false));
	dAnalysisActions.push_back(new DHistogramAction_ParticleComboKinematics(dComboWrapper, true, "KinFit"));

	//INITIALIZE ACTIONS
	Initialize_Actions();

	/******************************** EXAMPLE USER INITIALIZATION: STAND-ALONE HISTOGRAMS *******************************/

	//STAND-ALONE HISTOGRAMS
	dHist_NumSurvivingBeamParticles = new TH1I("NumSurvivingBeamParticles", "# Surviving Beam Particles", 100, 0.5, 100.5);

	//BEST COMBO HISTOGRAMS
	gDirectory->mkdir("BestCombo", "BestCombo")->cd();
	{
		dHist_LambdaMass_BestCombo = new TH1I("LambdaMass", ";p#pi^{#minus} Invariant Mass (GeV/c^{2})", 1000, 1.0, 1.2);
		dHist_MandelstamT = new TH1I("MandelstamT", ";t (GeV^{2})", 200, -5.0, 0.0);
		dHist_LambdaMassVsT = new TH2I("LambdaMassVsT", ";t (GeV^{2});p#pi^{#minus} Invariant Mass (GeV/c^{2})", 200, -5.0, 0.0, 400, 1.0, 1.2);

		//BEST COMBO: POLARIZATION
		dHist_ProdPlanePhi_PARA = new TH1I("ProdPlanePhi_PARA", ";Production Plane #phi#circ", 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Sideband_PARA = new TH2I("ProdPlanePhiVsT_Sideband_PARA", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Signal_PARA = new TH2I("ProdPlanePhiVsT_Signal_PARA", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);

		dHist_ProdPlanePhi_PERP = new TH1I("ProdPlanePhi_PERP", ";Production Plane #phi#circ", 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Sideband_PERP = new TH2I("ProdPlanePhiVsT_Sideband_PERP", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Signal_PERP = new TH2I("ProdPlanePhiVsT_Signal_PERP", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);

		dHist_ProdPlanePhi_AMO = new TH1I("ProdPlanePhi_AMO", ";Production Plane #phi#circ", 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Sideband_AMO = new TH2I("ProdPlanePhiVsT_Sideband_AMO", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);
		dHist_ProdPlanePhiVsT_Signal_AMO = new TH2I("ProdPlanePhiVsT_Signal_AMO", ";t (GeV^{2});Production Plane #phi#circ", 200, -5.0, 0.0, 360, -180.0, 180.0);
	}
	gDirectory->cd("..");

	/***************************************** ADVANCED: CHOOSE BRANCHES TO READ ****************************************/

	//TO SAVE PROCESSING TIME
		//If you know you don't need all of the branches/data, but just a subset of it, you can speed things up
		//By default, for each event, the data is retrieved for all branches
		//If you know you only need data for some branches, you can skip grabbing data from the branches you don't need
		//Do this by doing something similar to the commented code below

	//dTreeInterface->Clear_GetEntryBranches(); //now get none
	//dTreeInterface->Register_GetEntryBranch("Proton__P4"); //manually set the branches you want
}

Bool_t DSelector_klambda::Process(Long64_t locEntry)
{
	// The Process() function is called for each entry in the tree. The entry argument
	// specifies which entry in the currently loaded tree is to be processed.
	//
	// This function should contain the "body" of the analysis. It can contain
	// simple or elaborate selection criteria, run algorithms on the data
	// of the event and typically fill histograms.
	//
	// The processing can be stopped by calling Abort().
	// Use fStatus to set the return value of TTree::Process().
	// The return value is currently not used.

	//CALL THIS FIRST
	DSelector::Process(locEntry); //Gets the data from the tree for the entry
	//cout << "RUN " << Get_RunNumber() << ", EVENT " << Get_EventNumber() << endl;

	/******************************************** GET POLARIZATION ORIENTATION ******************************************/

	//Only if the run number changes
	UInt_t locRunNumber = Get_RunNumber();
	if(locRunNumber != dPreviousRunNumber)
	{
		dIsPolarizedFlag = dAnalysisUtilities.Get_IsPolarizedBeam(locRunNumber, dIsPARAFlag);
		dPreviousRunNumber = locRunNumber;
	}

	/********************************************* SETUP UNIQUENESS TRACKING ********************************************/

	//ANALYSIS ACTIONS: Reset uniqueness tracking for each action
	Reset_Actions_NewEvent();

	//PREVENT-DOUBLE COUNTING WHEN HISTOGRAMMING
		//Sometimes, some content is the exact same between one combo and the next
			//e.g. maybe two combos have different beam particles, but the same data for the final-state
		//When histogramming, you don't want to double-count when this happens: artificially inflates your signal (or background)
		//So, for each quantity you histogram, keep track of what particles you used (for a given combo)
		//Then for each combo, just compare to what you used before, and make sure it's unique

	//EXAMPLE 2: Combo-specific info:
		//In general: Could have multiple particles with the same PID: Use a set of Int_t's
		//In general: Multiple PIDs, so multiple sets: Contain within a map
		//Multiple combos: Contain maps within a set (easier, faster to search)

	//INSERT USER ANALYSIS UNIQUENESS TRACKING HERE

	/************************************************* LOOP OVER COMBOS *************************************************/

	//Loop over combos
	UInt_t locNumSurvivingCombos = 0;
	set<Int_t> locSurvivingBeamIDs;
	double locBestMissingMassSquared = 9.9E9;
	double locT_BestCombo = 0.0, locLambdaMass_BestCombo = 0.0, locProdPlanePhi_BestCombo = 0.0;
	bool locIsLambdaSignal_BestCombo = false;
	for(UInt_t loc_i = 0; loc_i < Get_NumCombos(); ++loc_i)
	{
		//Set branch array indices for combo and all combo particles
		dComboWrapper->Set_ComboIndex(loc_i);

		// Is used to indicate when combos have been cut
		if(dComboWrapper->Get_IsComboCut()) // Is false when tree originally created
			continue; // Combo has been cut previously

		/********************************************** GET PARTICLE INDICES *********************************************/

		//Used for tracking uniqueness when filling histograms, and for determining unused particles

		//Step 0
		Int_t locBeamID = dComboBeamWrapper->Get_BeamID();
		Int_t locKPlusTrackID = dKPlusWrapper->Get_TrackID();

		//Step 1
		Int_t locProtonTrackID = dProtonWrapper->Get_TrackID();
		Int_t locPiMinusTrackID = dPiMinusWrapper->Get_TrackID();

		/*********************************************** GET FOUR-MOMENTUM **********************************************/

		// Get P4's: //is kinfit if kinfit performed, else is measured
		//dTargetP4 is target p4
		//Step 0
		TLorentzVector locBeamP4 = dComboBeamWrapper->Get_P4();
		TLorentzVector locKPlusP4 = dKPlusWrapper->Get_P4();
		//Step 1
		TLorentzVector locDecayingLambdaP4 = dDecayingLambdaWrapper->Get_P4();
		TLorentzVector locProtonP4 = dProtonWrapper->Get_P4();
		TLorentzVector locPiMinusP4 = dPiMinusWrapper->Get_P4();

		// Get Measured P4's:
		//Step 0
		TLorentzVector locBeamP4_Measured = dComboBeamWrapper->Get_P4_Measured();
		TLorentzVector locKPlusP4_Measured = dKPlusWrapper->Get_P4_Measured();
		//Step 1
		TLorentzVector locProtonP4_Measured = dProtonWrapper->Get_P4_Measured();
		TLorentzVector locPiMinusP4_Measured = dPiMinusWrapper->Get_P4_Measured();

		/********************************************* COMBINE FOUR-MOMENTUM ********************************************/

		// DO YOUR STUFF HERE

		// Combine 4-vectors
		TLorentzVector locMissingP4_Measured = locBeamP4_Measured + dTargetP4;
		locMissingP4_Measured -= locKPlusP4_Measured + locProtonP4_Measured + locPiMinusP4_Measured;

		TLorentzVector locLambdaP4_Measured = locProtonP4_Measured + locPiMinusP4_Measured;

		/******************************************** EXECUTE ANALYSIS ACTIONS *******************************************/

		// Loop through the analysis actions, executing them in order for the active particle combo
		if(!Execute_Actions()) //if the active combo fails a cut, IsComboCutFlag automatically set
			continue;

		/************************************************* MISCELLANEOUS *************************************************/

		double locLambdaMass = locLambdaP4_Measured.M();
		bool locIsLambdaSignalFlag = ((locLambdaMass >= 1.095) && (locLambdaMass <= 1.136));

		++locNumSurvivingCombos;
		locSurvivingBeamIDs.insert(locBeamID);

		/*************************************************** BEST COMBO ***************************************************/

		// Signal selection and background subtraction are performed on the Lambda mass peak.
		// However, a given Lambda (proton, pi- pair) can be matched to multiple beam photons (or kaon candidates)
		// This is OK in a 1D plot (check for duplicates), but we are quoting results in 2D: In bins of Mandelstam-t
		// We cannot put a signal Lambda mass value into two separate t-bins: We would count the signal in each bin.
		// Therefore, we must choose a t-bin. This is, in effect, a choice of "Best" combo.

		// We must choose this "best" combo in a way that doesn't bias the background subtraction.
		// Since the background subtraction is done on the Lambda mass peak, we must choose something else.
		// Ideally we would pick the combo with the best kinematic-fit confidence level (or BDT response when it's ready).
		// However, at the time of this writing, the tracking errors are bad and this quantity cannot be trusted.
		// Therefore, we choose the combo that has a missing mass squared closest to zero: It's calculation involves all particles.

		// As far as the code is concerned, it is easiest to keep track of the relevant quantities for the "Best" combo
			// and update them as new "Best" combos are found. Then histogram the "Best" results at the end.

		double locMissingMassSquared = locMissingP4_Measured.M2();
		if(fabs(locMissingMassSquared) > fabs(locBestMissingMassSquared))
			continue;

		double locProdPlanePhi = dAnalysisUtilities.Calc_ProdPlanePhi_Pseudoscalar(locBeamP4.E(), dTargetPID, locKPlusP4);
		double locT = (locBeamP4 - locKPlusP4).Mag2();

		locBestMissingMassSquared = locMissingMassSquared;
		locT_BestCombo = locT;
		locLambdaMass_BestCombo = locLambdaMass;
		locProdPlanePhi_BestCombo = locProdPlanePhi;
		locIsLambdaSignal_BestCombo = locIsLambdaSignalFlag;
	}

	//FILL HISTOGRAMS: Num combos / events surviving actions
	Fill_NumCombosSurvivedHists();

	dHist_NumSurvivingBeamParticles->Fill(locSurvivingBeamIDs.size());

	//FILL ASYMMETRY INFO FOR BEST COMBOS
	if(locNumSurvivingCombos > 0)
	{
		dHist_LambdaMass_BestCombo->Fill(locLambdaMass_BestCombo);
		dHist_LambdaMassVsT->Fill(locT_BestCombo, locLambdaMass_BestCombo);
		dHist_MandelstamT->Fill(locT_BestCombo);
		if(locIsLambdaSignal_BestCombo)
		{
			if(!dIsPolarizedFlag)
			{
				dHist_ProdPlanePhiVsT_Signal_AMO->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
				dHist_ProdPlanePhi_AMO->Fill(locProdPlanePhi_BestCombo);
			}
			else if(dIsPARAFlag)
			{
				dHist_ProdPlanePhiVsT_Signal_PARA->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
				dHist_ProdPlanePhi_PARA->Fill(locProdPlanePhi_BestCombo);
			}
			else //PERP
			{
				dHist_ProdPlanePhiVsT_Signal_PERP->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
				dHist_ProdPlanePhi_PERP->Fill(locProdPlanePhi_BestCombo);
			}
		}
		else //sideband
		{
			if(!dIsPolarizedFlag)
				dHist_ProdPlanePhiVsT_Sideband_AMO->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
			else if(dIsPARAFlag)
				dHist_ProdPlanePhiVsT_Sideband_PARA->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
			else //PERP
				dHist_ProdPlanePhiVsT_Sideband_PERP->Fill(locT_BestCombo, locProdPlanePhi_BestCombo);
		}
	}

	/************************************ EXAMPLE: FILL CLONE OF TTREE HERE WITH CUTS APPLIED ************************************/
/*
	Bool_t locIsEventCut = true;
	for(UInt_t loc_i = 0; loc_i < Get_NumCombos(); ++loc_i) {
		//Set branch array indices for combo and all combo particles
		dComboWrapper->Set_ComboIndex(loc_i);
		// Is used to indicate when combos have been cut
		if(dComboWrapper->Get_IsComboCut())
			continue;
		locIsEventCut = false; // At least one combo succeeded
		break;
	}
	if(!locIsEventCut && dOutputTreeFileName != "")
		FillOutputTree();
*/
	return kTRUE;
}

void DSelector_klambda::Finalize(void)
{
	//Save anything to output here that you do not want to be in the default DSelector output ROOT file.

	//Otherwise, don't do anything else (especially if you are using PROOF).
		//If you are using PROOF, this function is called on each thread,
		//so anything you do will not have the combined information from the various threads.
		//Besides, it is best-practice to do post-processing (e.g. fitting) separately, in case there is a problem.

	//DO YOUR STUFF HERE

	//CALL THIS LAST
	DSelector::Finalize(); //Saves results to the output file
}