#ifndef _ResonanceNodeGenerator_
#define _ResonanceNodeGenerator_

#include <deque>
#include <map>
#include <set>
#include <vector>
#include <string>
#include <vector>
#include <sstream>
#include <iostream>
#include <stdlib.h>
#include <unistd.h>

#include "particleType.h"
#include "HCA_Amplitudes/resonanceinfo.h"

#include "Particle.h"
#include "ResonanceDecay.h"
#include "Transition.h"
#include "HCA_TensorContraction/HCA_TensorContractor.h"

using namespace std;

class ResonanceNodeGenerator
{
	public:
		ResonanceNodeGenerator(Particle_t locBeamPID, Particle_t locTargetPID, deque<Particle_t> locFinalStatePIDs) : 
		dBeamPID(locBeamPID), dTargetPID(locTargetPID), dFinalStatePIDs(locFinalStatePIDs)
		{
			dMaxParticleCharge = 2;
			dMinParticleCharge = -1;

			dDefaultBWMassOffset = 0.05;
			dDefaultBWWidth = 0.1;

			dLimitToOneSFlag = true;
			dFixParticlePermutationFlag = false;

			dDebugFlag = false;
		}

		//PIDs
		Particle_t dBeamPID;
		Particle_t dTargetPID;
		deque<Particle_t> dFinalStatePIDs;

		//Control Params
		int dMaxParticleCharge;
		int dMinParticleCharge;

		//Limit Flags
		bool dLimitToOneSFlag; //if true, choose only one S12 (the smallest) for each possible L12
		bool dFixParticlePermutationFlag; //if true, don't permute final state particles amongst decay chain locations: keep fixed

		//Breit-Wigner params
		double dDefaultBWMassOffset;
		double dDefaultBWWidth;

		bool dDebugFlag;

		void Create_Particles(Particle_t locBeamPID, Particle_t locTargetPID, deque<Particle_t>& locFinalStatePIDs, Particle*& locPsiBeam, Particle*& locPsiTarget, deque<Particle*>& locFinalStateParticles);
		void Register_Particles(Particle* locPsiBeam, Particle* locPsiTarget, deque<Particle*>& locFinalStateParticles);

		//if locMaxSChannelJTimes2 < 0, no s-channel (generate for other, input channels)
		deque<NodeBase*>& Generate(int locMaxSChannelJTimes2, int locMaxDecayJTimes2, deque<ExchangeNode*>& locExchangeNodes);
		void Print_Tree(NodeBase* locTreeBase, bool locPrintChainFlag = true, int locNumTabs = 0);
		void Print_ResonanceNode(ResonanceNode* locCurrentNode, bool locPrintChainFlag = true, int locNumTabs = 0);
		void Get_FitParams(set<FitParamInfo*>& locFitParams) const;

		deque<ResonanceDecay*> dResonanceDecays;

	private:

		//Decay Trees
		deque<NodeBase*> dSavedTreeNodes; //decay tree nodes
		NodeBase* gProductionNode; //top node

		//Particles
		Particle* dPsiBeam;
		Particle* dPsiTarget;
		set<Particle*> dFinalStateParticleSet;

		//Control Params
		int dMaxSChannelJTimes2;
		int dMaxDecayJTimes2;

		//other
		map<Particle*, int> dParticleIndexMap; //int is prodcut index //same order as data
		map<Particle_t, deque<Particle*> > dFinalStateParticleMap;
		set<FitParamInfo*> dFitParams;

		ExchangeNode* Clone_DecayTree(ExchangeNode* locInputExchangeNode);
		ResonanceNode* Clone_DecayTree(ResonanceNode* locInputResonanceNode, NodeBase* locPreviousDecayNode = NULL);
		double CalcAndSet_ResonanceMasses_Init(NodeBase* locInputTransitionTree);
		double CalcAndSet_ResonanceMasses(ResonanceNode* locCurrentNode);

		string Get_ProductIndices(ResonanceNode* locIsParentNode);
		void Get_ProductIndices(ResonanceNode* locCurrentNode, vector<int>& locProductIndices);

		void Generate_DecayTrees(ResonanceNode* locCurrentNode, set<Particle*> locRemainingFinalStateParticles, deque<ResonanceNode*> locOtherUndecayedNodesToProcess);
		void LoopOver_DecayAmplitudes_PickPsi1(ResonanceNode* locCurrentNode, set<Particle*> locRemainingFinalStateParticles, deque<ResonanceNode*> locOtherUndecayedNodesToProcess);
		void LoopOver_DecayAmplitudes_PickPsi2(ResonanceNode* locCurrentNode, set<Particle*> locRemainingFinalStateParticles, deque<ResonanceNode*> locOtherUndecayedNodesToProcess);
		void LoopOver_DecayAmplitudes(ResonanceNode* locCurrentNode, set<Particle*> locRemainingFinalStateParticles, deque<ResonanceNode*> locOtherUndecayedNodesToProcess);

		void Generate_ExchangeTrees(ExchangeNode* locExchangeNode);
};

inline void ResonanceNodeGenerator::Create_Particles(Particle_t locBeamPID, Particle_t locTargetPID, deque<Particle_t>& locFinalStatePIDs, Particle*& locPsiBeam, Particle*& locPsiTarget, deque<Particle*>& locFinalStateParticles)
{
	//GENERATE PARTICLES
	locPsiBeam = new Particle(locBeamPID);
	locPsiTarget = new Particle(locTargetPID);
	for(size_t loc_i = 0; loc_i < locFinalStatePIDs.size(); ++loc_i)
		locFinalStateParticles.push_back(new Particle(locFinalStatePIDs[loc_i]));
}

inline void ResonanceNodeGenerator::Register_Particles(Particle* locPsiBeam, Particle* locPsiTarget, deque<Particle*>& locFinalStateParticles)
{
	dPsiBeam = locPsiBeam;
	dPsiTarget = locPsiTarget;

	dParticleIndexMap[locPsiBeam] = 0;
	dParticleIndexMap[locPsiTarget] = -1;

	for(size_t loc_i = 0; loc_i < locFinalStateParticles.size(); ++loc_i)
	{
		Particle_t locPID = locFinalStateParticles[loc_i]->dPID;
		dFinalStateParticleSet.insert(locFinalStateParticles[loc_i]);
		dFinalStateParticleMap[locPID].push_back(locFinalStateParticles[loc_i]);
		dParticleIndexMap[locFinalStateParticles[loc_i]] = loc_i + 1;
	}
}

inline ExchangeNode* ResonanceNodeGenerator::Clone_DecayTree(ExchangeNode* locInputExchangeNode)
{
	if(locInputExchangeNode == NULL)
		return NULL;

	ExchangeNode* locNewNode = new ExchangeNode(*locInputExchangeNode);

	locNewNode->dLeftFinalNode = Clone_DecayTree(locInputExchangeNode->dLeftFinalNode, locNewNode);
	locNewNode->dRightFinalNode = Clone_DecayTree(locInputExchangeNode->dRightFinalNode, locNewNode);
	return locNewNode;
}

inline ResonanceNode* ResonanceNodeGenerator::Clone_DecayTree(ResonanceNode* locInputResonanceNode, NodeBase* locPreviousDecayNode)
{
	if(locInputResonanceNode == NULL)
		return NULL;

	ResonanceNode* locNewNode = new ResonanceNode(*locInputResonanceNode);
	if(locInputResonanceNode->dIsProductionNodeFlag)
		locNewNode->dPreviousDecayNode = Clone_DecayTree((ResonanceNode*)locInputResonanceNode->dPreviousDecayNode, locNewNode);
	else
		locNewNode->dPreviousDecayNode = locPreviousDecayNode;

	locNewNode->dLeftFinalNode = Clone_DecayTree(locInputResonanceNode->dLeftFinalNode, locNewNode);
	locNewNode->dRightFinalNode = Clone_DecayTree(locInputResonanceNode->dRightFinalNode, locNewNode);
	return locNewNode;
}

inline string ResonanceNodeGenerator::Get_ProductIndices(ResonanceNode* locIsParentNode)
{
	vector<int> locProductIndices;
	Get_ProductIndices(locIsParentNode, locProductIndices);
	ostringstream locProductIndexStream;
	for(size_t loc_i = 0; loc_i < locProductIndices.size(); ++loc_i)
	{
		if(loc_i != 0)
			locProductIndexStream << "_";
		locProductIndexStream << locProductIndices[loc_i];
	}
	return locProductIndexStream.str();
}

inline void ResonanceNodeGenerator::Get_ProductIndices(ResonanceNode* locCurrentNode, vector<int>& locProductIndices)
{
	//traverse decays
	map<Particle*, int>::iterator locIterator = dParticleIndexMap.find(locCurrentNode->dNodeParticle);
	if(locIterator != dParticleIndexMap.end())
		locProductIndices.push_back(locIterator->second);
	else
	{
		Get_ProductIndices(locCurrentNode->dLeftFinalNode, locProductIndices);
		Get_ProductIndices(locCurrentNode->dRightFinalNode, locProductIndices);
	}
}

inline void ResonanceNodeGenerator::Get_FitParams(set<FitParamInfo*>& locFitParams) const
{
	locFitParams = dFitParams;
}

#endif //_ResonanceNodeGenerator_