(* ----- AMPLITUDE FORMULAS FROM THE LITERATURE, FOR COMPARISON ----- *)


(* Only 3 independent NMHV helicity amplitudes need to be specified *)


AMP[{-1, -1, -1, 1, 1, 1}] = -
	braket[4,{5,6},1]^3/(bra[2,3] bra[3,4] braket[2,{3,4},5] ket[5,6] ket[6,1] S[2,3,4]) -
	braket[6,{1,2},3]^3/(bra[1,2] bra[6,1] braket[2,{3,4},5] ket[3,4] ket[4,5] S[1,2,6]);


AMP[{1, -1, 1, -1, 1, -1}] =
        ket[6,4]^4 bra[1,3]^4/(S[4,5,6] braket[1,{5,6},4] braket[3,{4,5},6] ket[4,5] ket[5,6] bra[1,2] bra[2,3]) +
        ket[6,2]^4 bra[3,5]^4/(S[6,1,2] braket[3,{1,2},6] braket[5,{6,1},2] ket[6,1] ket[1,2] bra[3,4] bra[4,5]) +
        ket[2,4]^4 bra[1,5]^4/(S[2,3,4] braket[5,{3,4},2] braket[1,{2,3},4] ket[2,3] ket[3,4] bra[5,6] bra[6,1]);


AMP[{1, 1, -1, 1, -1, -1}] =
	braket[4,{1,2},3]^4/(bra[4,5] bra[5,6] braket[4,{3,2},1] braket[6,{1,2},3] ket[1,2] ket[2,3] S[1,2,3]) +
	bra[2,4]^4 ket[5,6]^3/(bra[2,3] bra[3,4] braket[2,{3,4},5] braket[4,{3,2},1] ket[6,1] S[2,3,4]) +
	bra[1, 2]^3 ket[3,5]^4/(bra[6,1] braket[2,{3,4},5] braket[6,{5,4},3] ket[3,4] ket[4,5] S[3,4,5]);


(* We'll automatically generate everything else from these *)


(* First rewrite them in terms of spinor helicity bras [] and kets <> only *)

rewrite = {
	braket[a_, {b_, c_}, d_] -> bra[a, b] ket[b, d] + bra[a, c] ket[c, d],
	S[a__] :> (bra @@@ Subsets[{a}, {2}]).(ket @@@ Subsets[{a}, {2}])};

AMP[{-1, -1, -1, 1, 1, 1}] = AMP[{-1, -1, -1, 1, 1, 1}]/.rewrite;
AMP[{1, -1, 1, -1, 1, -1}] = AMP[{1, -1, 1, -1, 1, -1}]/.rewrite;
AMP[{1, 1, -1, 1, -1, -1}] = AMP[{1, 1, -1, 1, -1, -1}]/.rewrite;


(* Now generate all cyclic rotations of the above *)

AMP[{-1, 1, -1, 1, -1, 1}] =
	AMP[{1, -1, 1, -1, 1, -1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+1, 6, 1], Mod[b+1, 6, 1]];

For[i = 1, i < 6, i++,

	AMP[RotateRight[{-1, -1, -1, 1, 1, 1}, i]] =
		AMP[{-1, -1, -1, 1, 1, 1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+i, 6, 1], Mod[b+i, 6, 1]];

	AMP[RotateRight[{1, 1, -1, 1, -1, -1}, i]] =
		AMP[{1, 1, -1, 1, -1, -1}]/.f_[a_Integer, b_Integer] :> f[Mod[a+i, 6, 1], Mod[b+i, 6, 1]];

];


(* Finally the last six are obtained by parity conjugation *)

For[i = 0, i < 6, i++,

	AMP[RotateRight[{-1, -1, 1, -1, 1, 1}, i]] =
		AMP[RotateRight[{1, 1, -1, 1, -1, -1}, i]]/.{bra -> ket, ket -> bra};

];


(* Now we're done with NMHV amplitudes; here are the rest *)


(* MHV amplitudes *)

AMP[H_] := (ket @@ Flatten[Position[H, -1]])^4/Product[ket[i, Mod[i+1, Length[H], 1]], {i, Length[H]}] /; (Length[H] > 2 && Count[H, -1] == 2);


(* MHV-bar amplitudes *)

AMP[H_] := (bra @@ Flatten[Position[H, +1]])^4/Product[bra[i, Mod[i+1, Length[H], 1]], {i, Length[H]}] /; (Length[H] > 2 && Count[H, +1] == 2);

(* All others vanish *)


AMP[H_] := 0 /; (Length[H] < 7);



(* ----- SPINOR HELICITY AND MOMENTUM SPACE DEFINITIONS ----- *)


(* Note: the next two imply a hard-coded infinity twistor {{0,0,1,0},{0,0,0,1}} *)

(* The spinor helicity lambda variables *)

L[i_] := Take[Z[[Mod[i, Length[Z], 1]]], 2];


(* The spinor helicity variables *)

M[i_] := Take[Z[[Mod[i, Length[Z], 1]]], -2];


(* The < , > spinor bracket *)

ket[i_, j_] := L[i][[1]] L[j][[2]] - L[i][[2]] L[j][[1]];


(* The lambda-tilde variables *)

T[i_] := (ket[i-1,i] M[i+1] + ket[i+1,i-1] M[i] + ket[i,i+1] M[i-1])/(ket[i-1,i] ket[i,i+1]);


(* The [ , ] spinor bracket *)

bra[i_, j_] := T[i][[1]] T[j][[2]] - T[i][[2]] T[j][[1]];


(* The four-brakcet < , , , > on momentum twistors *)

det[x__] := Det[Z[[{x}]]];


(* Convert from 2x2 matrix in alpha, alpha-dot space to 4-vector *)

unpack[m_] :=
	{m[[1, 1]] + m[[2, 2]], m[[1, 2]] + m[[2, 1]],
	I m[[2, 1]] - I m[[1, 2]], m[[2, 2]] - m[[1, 1]]}/2;

unpack[a_, b_] := unpack[Outer[Times, a, b]];


(* Momentum 4-vector in terms of lambda and lambda-tilde *)

P[i_] := unpack[L[i], T[i]];


(* Minkowksi dot-product on 4-vectors *)

dot[a_, b_] := a.DiagonalMatrix[{1, -1, -1, -1}].b;



(* ---- EXTRACT NMHV HELICITY AMPLITUDE FROM MOMENTUM SUPERTWISTORS ----- *)


(* Definition of the 5-bracket, or R-invariant *)

R[a_, b_, c_, d_, e_] := (det[b,c,d,e] chi[a] + det[c,d,e,a] chi[b] + det[d,e,a,b] chi[c] + det[e,a,b,c] chi[d] + det[a,b,c,d] chi[e])^4/(det[a,b,c,d] det[b,c,d,e] det[c,d,e,a] det[d,e,a,b] det[e,a,b,c]);


(* Here is the code to extract a desired helicity component amplitude *)

AMT[H_] := Module[{ANMHV, m1, m2, m3, eqn, sol},


	(* Here is the NMHV superamplitude, including the MHV prefactor *)

	ANMHV = (R[6, 1, 2, 3, 4] + R[6, 1, 2, 4, 5] + R[6, 2, 3, 4, 5])/(ket[1, 2] ket[2, 3] ket[3, 4] ket[4, 5] ket[5, 6] ket[6, 1]);

	(* Identify the locations of the three negative helicities *)

	m1 = Position[H, -1][[1,1]];
	m2 = Position[H, -1][[2,1]];
	m3 = Position[H, -1][[3,1]];

	(* Solve for the first two using the supermomentum conservation equations *)

	eqn = Array[eta, 6].Array[L, 6];
	sol = Solve[Thread[eqn == 0], {eta[m1], eta[m2]}];

	(* Now we can solve for the complementary set of chi's *)

	eqn = Table[eta[i] - (ket[Mod[i-1,6,1],i] chi[Mod[i+1,6,1]] + ket[Mod[i+1,6,1],Mod[i-1,6,1]] chi[i] + ket[i,Mod[i+1,6,1]] chi[Mod[i-1,6,1]])/(ket[Mod[i-1,6,1], i] ket[i, Mod[i+1,6,1]]), {i, 1, 6}]/.sol[[1]];

	sol = Solve[Thread[eqn == 0], Delete[Array[chi, 6], {{m1}, {m2}}]];

	Return[Coefficient[Expand[ket[m1, m2]^4 ANMHV/.sol[[1]]], eta[m3]^4]];
];



(* ----- COMPUTE AMPLITUDES FROM FEYNMAN DIAGRAMS ----- *)


(* Propagator *)

prop[x_] := -I/dot[x, x];


(* Cubic vertex *)

V3[ei_, ej_, ek_, pi_, pj_, pk_] := I/Sqrt[2] * (
        dot[ei, ej] dot[pi - pj, ek] +
        dot[ej, ek] dot[pj - pk, ei] +
        dot[ek, ei] dot[pk - pi, ej]);


(* Quartic vertex *)

V4[a_, b_, c_, d_] := I dot[eps[a], eps[c]] dot[eps[b], eps[d]] -
	I/2 (dot[eps[a], eps[b]] dot[eps[c], eps[d]] +
	dot[eps[a], eps[d]] dot[eps[b], eps[c]]);


(* Five-point diagram type 1 *)

D1[a_,b_,c_,d_,e_] :=
        V3[eps[a], eps[b], eps[-1], P[a], P[b], -P[a] - P[b]] *
        V3[eps[d], eps[e], eps[-2], P[d], P[e], -P[d] - P[e]] *
        V3[eps[c], eps[-2], eps[-1], P[c], P[d]+P[e], P[a]+P[b]] *
	prop[P[a]+P[b]] prop[P[d]+P[e]];


(* Five-point diagram type 2 *)

D2[a_, b_, c_, d_, e_] :=
        V4[a, b, c, -1] prop[P[d]+P[e]] *
        V3[eps[d], eps[e], eps[-1], P[d], P[e], -P[d] - P[e]];


(* Six-point diagram type 3 *)

D3[a_, b_, c_, d_, e_, f_] :=
	 V3[eps[a], eps[b], eps[-1], P[a], P[b], - P[a] - P[b]] *
	V3[eps[-1], eps[c], eps[-2], P[a] + P[b], P[c], - P[a] - P[b] - P[c]] *
	V3[eps[-2], eps[d], eps[-3], P[a] + P[b] + P[c], P[d], - P[a] - P[b] - P[c] - P[d]] *
	V3[eps[-3], eps[e], eps[f], - P[e] - P[f], P[e], P[f]] *
	prop[P[a]+P[b]] prop[P[a]+P[b]+P[c]] prop[P[e]+P[f]];


(* Six-point diagram type 4 *)

D4[a_, b_, c_, d_, e_, f_] :=
	 V3[eps[a], eps[b], eps[-1], P[a], P[b], - P[a]-P[b]] *
	V4[-1,c,-2,f] *
	V3[eps[-2], eps[d], eps[e], -P[d]-P[e], P[d], P[e]] *
	prop[P[a]+P[b]] prop[P[d]+P[e]];


(* Six-point diagram type 5 *)

D5[a_, b_, c_, d_, e_, f_] :=
	 V4[a,b,c,-1] V4[-1,d,e,f] prop[P[a]+P[b]+P[c]];


(* Six-point diagram type 6 *)

D6[a_, b_, c_, d_, e_, f_] :=
	 V4[a, b, c, -1] V3[eps[-1], eps[d], eps[-2], P[a]+P[b]+P[c], P[d], P[e]+P[f]] *
	V3[eps[-2], eps[e], eps[f], -P[e]-P[f],P[e],P[f]] prop[P[a]+P[b]+P[c]] prop[P[e]+P[f]];


(* Six-point diagram type 7 *)

D7[a_, b_, c_, d_, e_, f_] :=
	 V4[-1,-2,e,f] V3[eps[a], eps[b], eps[-1], P[a], P[b], -P[a]-P[b]] *
	V3[eps[c],eps[d],eps[-2],P[c],P[d],-P[c]-P[d]] *
	prop[P[c]+P[d]] prop[P[a]+P[b]];


(* Six-point diagram type 8 *)

D8[a_, b_, c_, d_, e_, f_] :=
	 V3[eps[a], eps[b], eps[-1], P[a], P[b], -P[a]-P[b]] *
	V3[eps[c], eps[-2], eps[-1], P[c], -P[c]-P[a]-P[b], P[a]+P[b]] *
	V3[eps[f],eps[-2],eps[-3],P[f],P[a]+P[b]+P[c],-P[a]-P[b]-P[c]-P[f]] *
	V3[eps[d],eps[e],eps[-3],P[d],P[e],-P[d]-P[e]] *
	prop[P[a]+P[b]] prop[P[a]+P[b]+P[c]] prop[P[d]+P[e]];


(* Six-point diagram type 9 *)

D9[a_, b_, c_, d_, e_, f_] :=
	 V3[eps[a], eps[b], eps[-1], P[a],P[b],-P[a]-P[b]] *
	V3[eps[-1],eps[-2],eps[f],P[a]+P[b],-P[a]-P[b]-P[f],P[f]] *
	V3[eps[-2],eps[c],eps[-3],P[a]+P[b]+P[f],P[c],-P[a]-P[b]-P[c]-P[f]] *
	V3[eps[-3],eps[d],eps[e],-P[d]-P[e],P[d],P[e]] *
	prop[P[a]+P[b]] prop[P[a]+P[b]+P[f]] prop[P[d]+P[e]];


(* Six-point diagram type 10 *)

D10[a_, b_, c_, d_, e_, f_] :=
	 V4[a,b,c,-1] V3[eps[-1],eps[-2],eps[f],P[a]+P[b]+P[c],-P[a]-P[b]-P[c]-P[f],P[f]] *
	V3[eps[d],eps[e],eps[-2],P[d],P[e],-P[d]-P[e]] *
	prop[P[a]+P[b]+P[c]] prop[P[d]+P[e]];


(* Six-point diagram type 11 *)

D11[a_, b_, c_, d_, e_, f_] :=
	V3[eps[a],eps[b],eps[-1],P[a],P[b],-P[a]-P[b]] *
	V3[eps[c],eps[d],eps[-2],P[c],P[d],-P[c]-P[d]] *
	V3[eps[e],eps[f],eps[-3],P[e],P[f],-P[e]-P[f]] *
	V3[eps[-1],eps[-2],eps[-3],P[a]+P[b],P[c]+P[d],P[e]+P[f]] *
	prop[P[a]+P[b]] prop[P[c]+P[d]] prop[P[e]+P[f]];


(* Dummy variables for internal polarizations; these get summed over *)

eps[-1] = {aa, bb, cc, dd};
eps[-2] = {ee, ff, gg, hh};
eps[-3] = {ii, jj, kk, ll};


(* Here is the main function *)

AFD[H_] := Module[{n},

n = Length[Z];

(* Make the polarization vectors appropriate for the specified helicities *)

(* Note: {R1, R2} and {R3, R4} are the reference spinors; they will drop out *)

For[i = 1, i <= n, i++, If[H[[i]] == 1,
	eps[i] = - Sqrt[2] unpack[{R1, R2}, T[i]]/(L[i][[1]] R2 - L[i][[2]] R1),
	eps[i] = - Sqrt[2] unpack[L[i], {R3, R4}]/(T[i][[1]] R4 - T[i][[2]] R3)
]];

If[n == 3,

	A[H] = Factor[V3[eps[1], eps[2], eps[3], P[1], P[2], P[3]]];

];

If[n == 4,

	A[H] =	V3[eps[1], eps[2], eps[-1], P[1], P[2], -P[1] - P[2]] *
		V3[eps[3], eps[4], eps[-1], P[3], P[4], -P[3] - P[4]] prop[P[1]+P[2]] +
		V3[eps[2], eps[3], eps[-1], P[2], P[3], -P[2] - P[3]] *
		V3[eps[4], eps[1], eps[-1], P[4], P[1], -P[4] - P[1]] prop[P[2]+P[3]];

	A[H] = (A[H]/.{aa->1,bb->0,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->1,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->1,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->0,dd->1});

	A[H] = Factor[A[H] + V4[1, 2, 3, 4]];

];

If[n == 5,

	A[H] = Total[D1[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5},i],{i,5}]];

	A[H] =
	(A[H]/.{ee->1,ff->0,gg->0,hh->0}) -
	(A[H]/.{ee->0,ff->1,gg->0,hh->0}) -
	(A[H]/.{ee->0,ff->0,gg->1,hh->0}) -
	(A[H]/.{ee->0,ff->0,gg->0,hh->1});

	A[H] += Total[D2[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5},i],{i,5}]];

	A[H] = (A[H]/.{aa->1,bb->0,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->1,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->1,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->0,dd->1});

	A[H] = Factor[A[H]];

];

If[n == 6,

	A[H] = Total[D3[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5,6},i],{i,6}]];

	A[H] += D8[1,2,3,4,5,6] + D8[2,3,4,5,6,1] + D8[3,4,5,6,1,2];
	A[H] += D9[1,2,3,4,5,6] + D9[2,3,4,5,6,1] + D9[3,4,5,6,1,2];
	A[H] += D11[1,2,3,4,5,6] + D11[2,3,4,5,6,1];

	A[H] =
	(A[H]/.{ii->1,jj->0,kk->0,ll->0}) -
	(A[H]/.{ii->0,jj->1,kk->0,ll->0}) -
	(A[H]/.{ii->0,jj->0,kk->1,ll->0}) -
	(A[H]/.{ii->0,jj->0,kk->0,ll->1});

	A[H] += D4[1,2,3,4,5,6] + D4[2,3,4,5,6,1] + D4[3,4,5,6,1,2];
	A[H] += Total[D6[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5,6},i],{i,6}]];
	A[H] += Total[D7[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5,6},i],{i,6}]];
	A[H] += Total[D10[Sequence @@ #] & /@ Table[RotateRight[{1,2,3,4,5,6},i],{i,6}]];

	A[H] =
	(A[H]/.{ee->1,ff->0,gg->0,hh->0}) -
	(A[H]/.{ee->0,ff->1,gg->0,hh->0}) -
	(A[H]/.{ee->0,ff->0,gg->1,hh->0}) -
	(A[H]/.{ee->0,ff->0,gg->0,hh->1});

	A[H] += D5[1,2,3,4,5,6] + D5[2,3,4,5,6,1] + D5[3,4,5,6,1,2];

	A[H] = (A[H]/.{aa->1,bb->0,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->1,cc->0,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->1,dd->0}) -
	(A[H]/.{aa->0,bb->0,cc->0,dd->1});

	A[H] = Factor[A[H]];

];

Return[A[H]/I];

];
