#include "HCAmplitude_Step.h"

void HCAmplitude_Step::Calculate_TensorContraction(void)
{
	// This MUST be called prior to calcAmplitude()!!!
	deque<HCA_ContractionResult> locAllFinalAmplitudes;
	dHCABuilder.Calc_Invariants_2BodyDecay(dJ1Times2/2, dJ2Times2/2, dHelicity12Times2/2, dS12Times2/2, dL12, dJXTimes2/2, dProduct1IsPhotonFlag, dProduct2IsPhotonFlag, locAllFinalAmplitudes);
	for(size_t loc_i = 0; loc_i < locAllFinalAmplitudes.size(); ++loc_i)
	{
		if(locAllFinalAmplitudes[loc_i].dContractionIndex != dContractionIndex)
			continue;
		dFinalAmplitude = locAllFinalAmplitudes[loc_i];
		return;
	}
}

complex<GDouble> HCAmplitude_Step::calcAmplitude(GDouble** locKinematics) const
{
	if(dStepDebugFlag)
		cout << "H12x2, JXx2, S12x2 = " << dHelicity12Times2 << ", " << dJXTimes2 << ", " << dS12Times2 << endl;

	if((abs(dHelicity12Times2) > dS12Times2) || (abs(dHelicity12Times2) > dJXTimes2))
		return 0.0; //invalid combo
	if(dFinalAmplitude.dGammaTerms.empty())
		return 0.0; //tensor contraction result is always zero

	// Get/Build Product 1
	TLorentzVector locP4_Product1;
	for(int loc_i = 0; loc_i < dNumParticleIndices_Product1; ++loc_i)
	{
		int locIndex = dParticleIndices_Product1[loc_i];
		if(locIndex == -1) //target
			locP4_Product1 += TLorentzVector(0.0, 0.0, 0.0, dTargetMass);
		else
			locP4_Product1 += TLorentzVector(locKinematics[locIndex][1], locKinematics[locIndex][2], locKinematics[locIndex][3], locKinematics[locIndex][0]);
	}
	if(dStepDebugFlag)
	{
		cout << "Product1: " << dNumParticleIndices_Product1 << ": ";
		for(int loc_i = 0; loc_i < dNumParticleIndices_Product1; ++loc_i)
			cout << dParticleIndices_Product1[loc_i] << ", ";
		cout << endl;
	}

	// Get/Build Product 2
	TLorentzVector locP4_Product2;
	for(int loc_i = 0; loc_i < dNumParticleIndices_Product2; ++loc_i)
	{
		int locIndex = dParticleIndices_Product2[loc_i];
		if(locIndex == -1) //target
			locP4_Product2 += TLorentzVector(0.0, 0.0, 0.0, dTargetMass);
		else
			locP4_Product2 += TLorentzVector(locKinematics[locIndex][1], locKinematics[locIndex][2], locKinematics[locIndex][3], locKinematics[locIndex][0]);
	}
	if(dStepDebugFlag)
	{
		cout << "Product2: " << dNumParticleIndices_Product2 << ": ";
		for(int loc_i = 0; loc_i < dNumParticleIndices_Product2; ++loc_i)
			cout << dParticleIndices_Product2[loc_i] << ", ";
		cout << endl;
	}

	// Build Decaying Particle's Parent's P4
	TLorentzVector locP4_XsParentP4;
	for(int loc_i = 0; loc_i < dNumParticleIndices_XsParent; ++loc_i)
	{
		int locIndex = dParticleIndices_XsParent[loc_i];
		if(locIndex == -1) //target
			locP4_XsParentP4 += TLorentzVector(0.0, 0.0, 0.0, dTargetMass);
		else
			locP4_XsParentP4 += TLorentzVector(locKinematics[locIndex][1], locKinematics[locIndex][2], locKinematics[locIndex][3], locKinematics[locIndex][0]);
	}
	if(dStepDebugFlag)
	{
		cout << "XsParent: " << dNumParticleIndices_XsParent << ": ";
		for(int loc_i = 0; loc_i < dNumParticleIndices_XsParent; ++loc_i)
			cout << dParticleIndices_XsParent[loc_i] << ", ";
		cout << endl;
	}

	TLorentzVector locP4_YsParentP4;
	for(int loc_i = 0; loc_i < dNumParticleIndices_YsParent; ++loc_i)
	{
		int locIndex = dParticleIndices_YsParent[loc_i];
		if(locIndex == -1) //target
			locP4_YsParentP4 += TLorentzVector(0.0, 0.0, 0.0, dTargetMass);
		else
			locP4_YsParentP4 += TLorentzVector(locKinematics[locIndex][1], locKinematics[locIndex][2], locKinematics[locIndex][3], locKinematics[locIndex][0]);
	}
	if(dStepDebugFlag)
	{
		cout << "YsParent: " << dNumParticleIndices_YsParent << ": ";
		for(int loc_i = 0; loc_i < dNumParticleIndices_YsParent; ++loc_i)
			cout << dParticleIndices_YsParent[loc_i] << ", ";
		cout << endl;
	}

	// Boost to CM Frame
	TLorentzVector locTotalFinalStateP4 = locP4_Product1 + locP4_Product2;
	TLorentzRotation locLorentzBoost(-1.0*locTotalFinalStateP4.BoostVector());
	TLorentzVector locP4_1_CM = locLorentzBoost*locP4_Product1;
	TLorentzVector locP4_2_CM = locLorentzBoost*locP4_Product2;

	// Define coordinate system
	TLorentzRotation locLorentzBoost_XsParent(-1.0*locP4_XsParentP4.BoostVector());
	TLorentzVector locP4_X_XsParentCM = locLorentzBoost_XsParent*locTotalFinalStateP4;

	TLorentzRotation locLorentzBoost_YsParent(-1.0*locP4_YsParentP4.BoostVector());
	TLorentzVector locP4_Y_YsParentCM = locLorentzBoost_YsParent*locP4_XsParentP4;

	TVector3 locXHat, locYHat, locZHat;
	if(dNumParticleIndices_XsParent == 0)
	{
		//B, T -> X
		//        X -> 1, 2
		//OR
		//B, T -> X
		locZHat = TVector3(0.0, 0.0, 1.0);
		locYHat = TVector3(0.0, 1.0, 0.0);
		locXHat = TVector3(1.0, 0.0, 0.0);
	}
	else if(dNumParticleIndices_YsParent == 0)
	{
		//B, T -> Y
		//        Y -> X, C
		//             X -> 1, 2
		locZHat = locP4_X_XsParentCM.Vect().Unit();
		locYHat = TVector3(0.0, 0.0, 1.0).Cross(locZHat);
		locXHat = locYHat.Cross(locZHat);
	}
	else
	{
		//B, T -> Z
		//        Z -> Y, A
		//             Y -> X, C
		//                  X -> 1, 2
		locZHat = locP4_X_XsParentCM.Vect().Unit();
		locYHat = locP4_Y_YsParentCM.Vect().Unit().Cross(locZHat);
		locXHat = locYHat.Cross(locZHat);
	}

	TVector3 locP3_1_CM_Rotated(locP4_1_CM.Vect().Dot(locXHat), locP4_1_CM.Vect().Dot(locYHat), locP4_1_CM.Vect().Dot(locZHat));

	GDouble locDecayCosTheta = locP3_1_CM_Rotated.CosTheta();
	if((fabs(locDecayCosTheta - 1.0) < 1.0E-20) && (dHelicityXTimes2 != dHelicity12Times2))
		return 0.0; //particles are defined along z-axis (theta = 0) and helicities don't match: wigner-D-function term is zero

	GDouble locDecayP = locP3_1_CM_Rotated.Mag();
	GDouble locDecayPhi = locP3_1_CM_Rotated.Phi();
	GDouble locInvariantMass = locTotalFinalStateP4.M();

	//wignerD_1: 1st term is spin, 2nd & 3rd terms are the 1st & 2nd z-components, cos(theta), phi (rad)
	complex<GDouble> locWignerD = wignerD(double(dJXTimes2)/2.0, double(dHelicityXTimes2)/2.0, double(dHelicity12Times2)/2.0, locDecayCosTheta, locDecayPhi);
	complex<GDouble> locWignerDStar = std::conj(locWignerD);

//	GDouble locHankelX = locDecayP/0.197327; //0.197327 -> 1 fm converted to GeV/c
//	GDouble locBarrierFactor = Calc_BarrierFactor(locHankelX, dL12);
	GDouble locBarrierFactor = barrierFactor(locDecayP, dL12); //not technically accurate, but faster (and good enough)

	complex<GDouble> locCovariantDecayAmplitude = Calc_CovariantDecayAmplitude(locP4_1_CM, locP4_2_CM);
	complex<GDouble> locHelicityDecayAmplitude = sqrt(double(2*dL12 + 1))*dClebschGordanCoefficient*locBarrierFactor*locCovariantDecayAmplitude;
	if(dStepDebugFlag)
		cout << "Amp Step Terms = " << locInvariantMass << ", " << locDecayP << ", " << locWignerDStar << ", " << dClebschGordanCoefficient << ", " << locBarrierFactor << ", " << locCovariantDecayAmplitude << endl;
	return sqrt(4.0*3.14159265359*locInvariantMass/locDecayP)*locWignerDStar*locHelicityDecayAmplitude;
}

GDouble HCAmplitude_Step::Calc_BarrierFactor(GDouble locHankelX, int locL) const
{
	complex<GDouble> locI(0.0, 1.0);
	complex<GDouble> locSum = 0.0;
	for(int loc_k = 0; loc_k <= locL; ++loc_k)
	{
		complex<GDouble> locTerm = Factorial(locL + loc_k, locL - loc_k + 1)/(pow(2.0*locHankelX, loc_k)*Factorial(loc_k));
		if(((locL + loc_k) % 2 == 1) != ((loc_k % 4 == 2) || (loc_k % 4 == 3)))
			locTerm *= -1.0; //each term contributes -1, so don't multiply if both true (or false)
		if((loc_k % 4 == 1) || (loc_k % 4 == 3))
			locTerm *= locI;
		locSum += locTerm;
	}
	complex<GDouble> locBarrierFactor = (sin(locHankelX) - locI*cos(locHankelX))*locSum;
	return 1.0/(std::abs(locBarrierFactor));
}

complex<GDouble> HCAmplitude_Step::Calc_CovariantDecayAmplitude(TLorentzVector locP4_1_CM, TLorentzVector locP4_2_CM) const
{
	TLorentzVector locTP4_1_CM(locP4_1_CM.Px(), locP4_1_CM.Py(), locP4_1_CM.Pz(), locP4_1_CM.E());
	TLorentzVector locTP4_2_CM(locP4_2_CM.Px(), locP4_2_CM.Py(), locP4_2_CM.Pz(), locP4_2_CM.E());
	complex<double> locResult = dFinalAmplitude.Calculate_Result(locTP4_1_CM, locTP4_2_CM);
	return complex<GDouble>(locResult.real(), locResult.imag());
}