#include "HCAmplitude_Chain.h"

HCAmplitude_Chain::HCAmplitude_Chain(const vector<string>& locArguments) : UserAmplitude<HCAmplitude_Chain>(locArguments)
{
	//args: external_helicities pids (_-separated), # steps, steps (JXx2 L12 S12x2 ContractionIndex ParticleX__Particle1__Particle2__...etc)
		//external_helicities & pids: by p4 array index, EXCEPT target inserted at index = 1 (with indices starting at 0)
		//particles: _-separated p4 array index locations

	string locHelicityString = locArguments[0].substr(4); //remove Hx2_ prefix
	ExtractInts_SingleUnderscore(locHelicityString, dExternalHelicitiesTimes2);

	vector<int> locPIDInts;
	ExtractInts_SingleUnderscore(locArguments[1], locPIDInts);
	for(size_t loc_i = 0; loc_i < locPIDInts.size(); ++loc_i)
		dPIDs.push_back(Particle_t(locPIDInts[loc_i]));

	//Extract Step Info
	dNumSteps = atoi(locArguments[2].c_str());
	for(int loc_i = 0; loc_i < dNumSteps; ++loc_i)
	{
		dJXTimes2Vector.push_back(atoi(locArguments[3 + 5*loc_i].c_str()));
		dL12Vector.push_back(atoi(locArguments[4 + 5*loc_i].c_str()));
		dS12Times2Vector.push_back(atoi(locArguments[5 + 5*loc_i].c_str()));
		dContractionIndexVector.push_back(atoi(locArguments[6 + 5*loc_i].c_str()));

		vector<vector<int> > locStepParticleIndices;
		ExtractInts_DoubleUnderscore(locArguments[7 + 5*loc_i], locStepParticleIndices);
		dChainParticleIndices.push_back(locStepParticleIndices);
	}

	//Now, loop over steps and resonance helicities and create amplitude objects
	vector<vector<HCAmplitude_Step*> > locChainAmplitudes;
	for(int loc_i = 0; loc_i < dNumSteps; ++loc_i)
	{
		vector<HCAmplitude_Step*> locStepAmplitudes;
		for(int locXHelicityTimes2 = -1*dJXTimes2Vector[loc_i]; locXHelicityTimes2 <= dJXTimes2Vector[loc_i]; locXHelicityTimes2 += 2)
		{
			vector<int> locCurrentHelicitiesTimes2(1, locXHelicityTimes2);
			vector<int> locCurrentJsTimes2(1, dJXTimes2Vector[loc_i]);
			vector<HCAmplitude_Step*> locTempStepAmplitudes;
			Build_DecayAmplitudes(loc_i, 1, locCurrentHelicitiesTimes2, locCurrentJsTimes2, locStepAmplitudes);
			locStepAmplitudes.insert(locStepAmplitudes.end(), locTempStepAmplitudes.begin(), locTempStepAmplitudes.end());
		}
		locChainAmplitudes.push_back(locStepAmplitudes);
	}

	//ok, amplitudes are in vectors, but can't operate with vectors: convert to arrays
	dDecayAmplitudes = new HCAmplitude_Step**[dNumSteps];
	dAmplitudeValues = new complex<GDouble>*[dNumSteps];
	dAmplitudeArrayPermutationSizes = new int[dNumSteps];
	for(int loc_i = 0; loc_i < dNumSteps; ++loc_i)
	{
		dDecayAmplitudes[loc_i] = new HCAmplitude_Step*[locChainAmplitudes[loc_i].size()];
		dAmplitudeValues[loc_i] = new complex<GDouble>[locChainAmplitudes[loc_i].size()];
		dAmplitudeArrayPermutationSizes[loc_i] = locChainAmplitudes[loc_i].size();
		for(size_t loc_j = 0; loc_j < locChainAmplitudes[loc_i].size(); ++loc_j)
			dDecayAmplitudes[loc_i][loc_j] = locChainAmplitudes[loc_i][loc_j];
	}

	//build all possible permutations of resonance helicities
	deque<int> locResumeAtIndices;
	for(size_t loc_i = 0; loc_i < dJXTimes2Vector.size(); ++loc_i)
		locResumeAtIndices.push_back(0);

	int locStepIndex = 0;
	deque<int> locResonanceHelicityIndices;
	deque<deque<int> > locResonanceHelicityIndicesPermutations;
	while(true)
	{
		if(locStepIndex == dNumSteps)
		{
			//permutation finished, save results
			locResonanceHelicityIndicesPermutations.push_back(locResonanceHelicityIndices);
			//go through previous indices, resetting them, until reach an index where we can resume
			if(!Handle_HelicityDecursion(locStepIndex, locResonanceHelicityIndices, locResumeAtIndices))
				break;
			continue; //continue on previous index
		}

		int locFinalJTimes2 = dJXTimes2Vector[locStepIndex];
		int locResonanceHelicityIndex = locResumeAtIndices[locStepIndex];
		if(locResonanceHelicityIndex >= (locFinalJTimes2 + 1))
		{
			if(!Handle_HelicityDecursion(locStepIndex, locResonanceHelicityIndices, locResumeAtIndices))
				break;
			continue;
		}
		locResonanceHelicityIndices.push_back(locResonanceHelicityIndex);
		locResumeAtIndices[locStepIndex] = locResonanceHelicityIndex + 1;
		++locStepIndex;

		//check if s-channel resonance appears twice
		if((locStepIndex == 1) && (locStepIndex < int(dChainParticleIndices.size())))
		{
			if(dChainParticleIndices[0][0] == dChainParticleIndices[0][1])
			{
				locResonanceHelicityIndices.push_back(locResonanceHelicityIndex);
				locResumeAtIndices[locStepIndex] = locResonanceHelicityIndex + 1;
				++locStepIndex;
			}
		}
	}

	//loop over helicity permutations and find corresponding amplitude indices
	for(size_t loc_i = 0; loc_i < locResonanceHelicityIndicesPermutations.size(); ++loc_i)
		Find_AmplitudeIndices(locResonanceHelicityIndicesPermutations[loc_i]);

	//amplitude indices are in vectors, but can't operate with vectors: convert to arrays
	dNumAmplitudeCombos = dChainAmplitudeIndices.size();
	dAmplitudeHelicityComboIndices = new int*[dNumAmplitudeCombos];
	for(int loc_i = 0; loc_i < dNumAmplitudeCombos; ++loc_i)
	{
		dAmplitudeHelicityComboIndices[loc_i] = new int[dNumSteps];
		for(int loc_j = 0; loc_j < dNumSteps; ++loc_j)
			dAmplitudeHelicityComboIndices[loc_i][loc_j] = dChainAmplitudeIndices[loc_i][loc_j];
	}

	//finally, create object for other amplitude component (if present)
	size_t locOtherClassNameIndex = 3 + 5*dNumSteps;
	dOtherAmplitudeComponent = NULL;
	if(locArguments.size() == locOtherClassNameIndex)
		return; //no other amplitude class

	vector<string> locClassArguments(locArguments.begin() + locOtherClassNameIndex + 1, locArguments.end());
	dOtherAmplitudeComponent = (locArguments[locOtherClassNameIndex] == "TChannel_Basic") ? new TChannel_Basic(locClassArguments) : NULL;
}

void HCAmplitude_Chain::Build_DecayAmplitudes(int locStepIndex, int locKidIndex, vector<int> locCurrentHelicitiesTimes2, vector<int> locCurrentJsTimes2, vector<HCAmplitude_Step*>& locStepAmplitudes)
{
	if(locKidIndex == int(dChainParticleIndices[locStepIndex].size()))
	{
		//end of step: create amplitude object
		Create_DecayAmplitude(locStepIndex, locCurrentHelicitiesTimes2, locCurrentJsTimes2, locStepAmplitudes);
		return;
	}

	vector<int> locParticleIndices = dChainParticleIndices[locStepIndex][locKidIndex];
	if(locParticleIndices.size() == 1) //final-state particle, helicity is set externally
	{
		int locArrayIndex = locParticleIndices[0];
		if(locParticleIndices[0] == -1)
			locArrayIndex = 1; //target
		if(locParticleIndices[0] > 0)
			++locArrayIndex; //because of target
		locCurrentHelicitiesTimes2.push_back(dExternalHelicitiesTimes2[locArrayIndex]);
		locCurrentJsTimes2.push_back(ParticleJTimes2(dPIDs[locArrayIndex]));
		Build_DecayAmplitudes(locStepIndex, locKidIndex + 1, locCurrentHelicitiesTimes2, locCurrentJsTimes2, locStepAmplitudes);
		return;
	}

	//find resonance J: look at future steps for when it is a parent
	int locResonanceJTimes2 = -1;
	for(size_t locFutureStepIndex = locStepIndex + 1; locFutureStepIndex < dChainParticleIndices.size(); ++locFutureStepIndex)
	{
		if(locParticleIndices != dChainParticleIndices[locFutureStepIndex][0])
			continue;
		locResonanceJTimes2 = dJXTimes2Vector[locFutureStepIndex];
		break;
	}
	if(locResonanceJTimes2 < 0)
	{
		cout << "ERROR: BAD RESONANCE Jx2 = " << locResonanceJTimes2 << endl;
		abort();
	}

	//loop over resonance helicities
	for(int locHelicityTimes2 = -1*locResonanceJTimes2; locHelicityTimes2 <= locResonanceJTimes2; locHelicityTimes2 += 2)
	{
		vector<int> locNewHelicitiesTimes2 = locCurrentHelicitiesTimes2;
		locNewHelicitiesTimes2.push_back(locHelicityTimes2);
		vector<int> locNewJsTimes2 = locCurrentJsTimes2;
		locNewJsTimes2.push_back(locResonanceJTimes2);
		Build_DecayAmplitudes(locStepIndex, locKidIndex + 1, locNewHelicitiesTimes2, locNewJsTimes2, locStepAmplitudes);
	}
}

void HCAmplitude_Chain::Create_DecayAmplitude(int locStepIndex, vector<int> locCurrentHelicitiesTimes2, vector<int> locCurrentJsTimes2, vector<HCAmplitude_Step*>& locStepAmplitudes)
{
	HCAmplitude_Step* locStepAmplitude = new HCAmplitude_Step();

	locStepAmplitude->dNumParticleIndices_Product1 = BuildArray(dChainParticleIndices[locStepIndex][1], locStepAmplitude->dParticleIndices_Product1);
	locStepAmplitude->dNumParticleIndices_Product2 = BuildArray(dChainParticleIndices[locStepIndex][2], locStepAmplitude->dParticleIndices_Product2);

	//set y indices
	int locYStepIndex = -1;
	for(int locPreviousStepIndex = locStepIndex - 1; locPreviousStepIndex >= 0; --locPreviousStepIndex)
	{
		//kids
		for(size_t locPreviousStepKidIndex = 1; locPreviousStepKidIndex < dChainParticleIndices[locPreviousStepIndex].size(); ++locPreviousStepKidIndex)
		{
			if(dChainParticleIndices[locStepIndex][0] != dChainParticleIndices[locPreviousStepKidIndex][locPreviousStepKidIndex])
				continue;
			locYStepIndex = locPreviousStepIndex;
			break;
		}
		if(locYStepIndex != -1)
			break;
	}

	if(locYStepIndex > 0)
		locStepAmplitude->dNumParticleIndices_XsParent = BuildArray(dChainParticleIndices[locYStepIndex][0], locStepAmplitude->dParticleIndices_XsParent);
	else
	{
		locStepAmplitude->dNumParticleIndices_XsParent = 0;
		locStepAmplitude->dParticleIndices_XsParent = NULL;
	}

	//set z indices
	int locZStepIndex = -1;
	for(int locPreviousStepIndex = locYStepIndex - 1; locPreviousStepIndex >= 0; --locPreviousStepIndex)
	{
		//kids
		for(size_t locPreviousStepKidIndex = 1; locPreviousStepKidIndex < dChainParticleIndices[locPreviousStepIndex].size(); ++locPreviousStepKidIndex)
		{
			if(dChainParticleIndices[locYStepIndex][0] != dChainParticleIndices[locPreviousStepKidIndex][locPreviousStepKidIndex])
				continue;
			locZStepIndex = locPreviousStepIndex;
			break;
		}
		if(locZStepIndex != -1)
			break;
	}

	if(locZStepIndex > 0)
		locStepAmplitude->dNumParticleIndices_YsParent = BuildArray(dChainParticleIndices[locZStepIndex][0], locStepAmplitude->dParticleIndices_YsParent);
	else
	{
		locStepAmplitude->dNumParticleIndices_YsParent = 0;
		locStepAmplitude->dParticleIndices_YsParent = NULL;
	}

	if((dChainParticleIndices[locStepIndex][1].size() == 1) && (dPIDs[dChainParticleIndices[locStepIndex][1][0]] == Gamma))
		locStepAmplitude->dProduct1IsPhotonFlag = true;
	else
		locStepAmplitude->dProduct1IsPhotonFlag = false;

	if((dChainParticleIndices[locStepIndex][2].size() == 1) && (dPIDs[dChainParticleIndices[locStepIndex][2][0]] == Gamma))
		locStepAmplitude->dProduct2IsPhotonFlag = true;
	else
		locStepAmplitude->dProduct2IsPhotonFlag = false;

	locStepAmplitude->dJ1Times2 = locCurrentJsTimes2[1];
	locStepAmplitude->dHelicity1Times2 = locCurrentHelicitiesTimes2[1];

	locStepAmplitude->dJ2Times2 = locCurrentJsTimes2[2];
	locStepAmplitude->dHelicity2Times2 = locCurrentHelicitiesTimes2[2];

	locStepAmplitude->dJXTimes2 = locCurrentJsTimes2[0];
	locStepAmplitude->dHelicityXTimes2 = locCurrentHelicitiesTimes2[0];

	locStepAmplitude->dL12 = dL12Vector[locStepIndex];
	locStepAmplitude->dS12Times2 = dS12Times2Vector[locStepIndex];

	locStepAmplitude->dContractionIndex = dContractionIndexVector[locStepIndex];

	locStepAmplitude->dTargetMass = ParticleMass(dPIDs[1]);

	locStepAmplitude->dHelicity12Times2 = locStepAmplitude->dHelicity1Times2 - locStepAmplitude->dHelicity2Times2;
	locStepAmplitude->dClebschGordanCoefficient = clebschGordan(double(locStepAmplitude->dJ1Times2)/2.0, double(locStepAmplitude->dJ2Times2)/2.0, double(locStepAmplitude->dHelicity1Times2)/2.0, -1.0*double(locStepAmplitude->dHelicity2Times2)/2.0, double(locStepAmplitude->dS12Times2)/2.0, double(locStepAmplitude->dHelicity12Times2)/2.0);

	locStepAmplitude->Calculate_TensorContraction();

	locStepAmplitudes.push_back(locStepAmplitude);
}

bool HCAmplitude_Chain::Handle_HelicityDecursion(int& locStepIndex, deque<int>& locResonanceHelicityIndices, deque<int>& locResumeAtIndices)
{
	//go through previous indices, resetting them, until reach an index where we can resume
	while(true)
	{
		if(locResonanceHelicityIndices.empty())
			return false; //we're done

		//reset this index
		if(locStepIndex < int(locResumeAtIndices.size()))
			locResumeAtIndices[locStepIndex] = 0;

		//re-do (resume-on) the previous rank
		locResonanceHelicityIndices.pop_back();
		--locStepIndex;
		if(locStepIndex == 1)
			continue; //continue re-cursing (pop the next one). this index was forced //s-channel resonance appeared twice
		if(locStepIndex < 0)
			return false; //we're done

		if(locResumeAtIndices[locStepIndex] < (dJXTimes2Vector[locStepIndex] + 1))
			break;
	}

	return true;
}

void HCAmplitude_Chain::Find_AmplitudeIndices(const deque<int>& locStepHelicityIndices)
{
	deque<int> locAmplitudeIndices;
	for(size_t locStepIndex = 0; locStepIndex < dChainParticleIndices.size(); ++locStepIndex)
	{
		int locAmplitudeIndex = 0;
		int locIndexMultiplier = 1; //1st one has largest multiplier
		for(int locKidIndex = dChainParticleIndices[locStepIndex].size() - 1; locKidIndex >= 0; --locKidIndex)
		{
			if(dChainParticleIndices[locStepIndex][locKidIndex].size() == 1)
				continue; //initial or final state particle: helicity already chosen

			int locHelicityIndex = -1;
			int locJTimes2 = -1;
			if(locKidIndex == 0)
			{
				//is step parent resonance
				locHelicityIndex = locStepHelicityIndices[locStepIndex];
				locJTimes2 = dJXTimes2Vector[locStepIndex];
			}
			else
			{
				//find resonance J: look at future steps for when it is a parent
				for(size_t locFutureStepIndex = locStepIndex + 1; locFutureStepIndex < dChainParticleIndices.size(); ++locFutureStepIndex)
				{
					if(dChainParticleIndices[locStepIndex][locKidIndex] != dChainParticleIndices[locFutureStepIndex][0])
						continue;
					locHelicityIndex = locStepHelicityIndices[locFutureStepIndex];
					locJTimes2 = dJXTimes2Vector[locFutureStepIndex];
					break;
				}
			}
			if(locHelicityIndex < 0)
			{
				cout << "ERROR: BAD HELICITY INDEX = " << locHelicityIndex << endl;
				abort();
			}

			locAmplitudeIndex += locIndexMultiplier*locHelicityIndex;
			locIndexMultiplier *= locJTimes2 + 1;
		}
		locAmplitudeIndices.push_back(locAmplitudeIndex);
	}
	dChainAmplitudeIndices.push_back(locAmplitudeIndices);
}

complex<GDouble> HCAmplitude_Chain::calcAmplitude(GDouble** locKinematics) const
{
	//Calculate HC-amplitude terms
	for(int loc_i = 0; loc_i < dNumSteps; ++loc_i)
	{
		for(int loc_j = 0; loc_j < dAmplitudeArrayPermutationSizes[loc_i]; ++loc_j)
			dAmplitudeValues[loc_i][loc_j] = (dDecayAmplitudes[loc_i][loc_j])->calcAmplitude(locKinematics);
	}

	//For each permutation of HC resonance helicities, compute the total amplitudes, and sum them together
	complex<GDouble> locAmplitudeSum(0.0, 0.0);
	for(int loc_i = 0; loc_i < dNumAmplitudeCombos; ++loc_i)
	{
		complex<GDouble> locPermutationAmplitude(1.0, 0.0);
		for(int loc_j = 0; loc_j < dNumSteps; ++loc_j)
		{
			int locAmplitudeIndex = dAmplitudeHelicityComboIndices[loc_i][loc_j];
			locPermutationAmplitude *= dAmplitudeValues[loc_j][locAmplitudeIndex];
		}
		locAmplitudeSum += locPermutationAmplitude;
	}

	//Calculate other terms
	if(dOtherAmplitudeComponent != NULL)
		locAmplitudeSum *= dOtherAmplitudeComponent->calcAmplitude(locKinematics);

	if((fabs(std::real(locAmplitudeSum)) < 1.0E-20) && (fabs(std::imag(locAmplitudeSum)) < 1.0E-20))
		cout << "WARNING, AMPLITUDE SUM IS ZERO" << endl;

	return locAmplitudeSum;
}