from .process import *

def _reconstruct_in_channel_S(proc, inters, table1, table2):
	m_A, m_B, m_C, m_D = m_for_part(*proc.abcd)

	S_part, P_part = sp.S.Zero, sp.S.Zero

	for inter in inters:
		lh_S, lh_P = lookup_SP_wave(scatter(proc.ab, inter), table1)
		rh_S, rh_P = lookup_SP_wave(scatter(proc.cd, inter), table2)

		sym_factor = 2 if (inter == inter[::-1]) else 1 
		heavi = sp.Heaviside(s - (m_for_part(inter[0]) + m_for_part(inter[1]))**2)

		S_part += lambda_sqrt(s, m_for_part(inter[0]), m_for_part(inter[1]))/s/sym_factor \
				  * heavi * lh_S*rh_S

		P_part += lambda_sqrt(s, m_for_part(inter[0]), m_for_part(inter[1]))/s/sym_factor \
				  * heavi * lh_P*rh_P \
					 / lambda_sqrt(s, m_A, m_B) / lambda_sqrt(s, m_C, m_D)

	W_0 = proc.phase * disp_integral(S_part, s)
	W_1 = proc.phase * disp_integral(s*P_part, s)

	return W_0, W_1

def _reconstruct_in_channel(channel, proc, inters, table1, table2):
	if channel == "S":
		W_0, W_1 = _reconstruct_in_channel_S(proc, inters, table1, table2)
		subs = {}
	elif channel == "T":
		W_0, W_1 = _reconstruct_in_channel_S(proc.T, inters, table1, table2)
		subs = {s: t, t: u, u: s}
	elif channel == "U":
		W_0, W_1 = _reconstruct_in_channel_S(proc.U.tu, inters, table1, table2)
		subs = {s: u, u: s}
	else:
		raise ValueError("invalid channel: %r (must be one of 'S', 'T', 'U')" % channel)

	return W_0.subs(subs, simultaneous=True), W_1.subs(subs, simultaneous=True)

def reconstruct_in_channel(channel, target, inters, table):
	return _reconstruct_in_channel(channel, target, inters, table, table)

def reconstruct(target, ptable, S=[], T=[], U=[]):
	W_0_S, W_1_S = reconstruct_in_channel("S", target, S, ptable)
	W_0_T, W_1_T = reconstruct_in_channel("T", target, T, ptable)
	W_0_U, W_1_U = reconstruct_in_channel("U", target, U, ptable)

	return 16*sp.S.Pi*(
		W_0_S + 3*(t - u)*W_1_S
		+ W_0_T + 3*(u - s)*W_1_T
		+ W_0_U + 3*(t - s)*W_1_U
	)

def reconstruct_in_channel_ext(channel, target, inters, table1, table2):
	W_0_a, W_1_a = _reconstruct_in_channel(channel, target, inters, table1, table2)
	W_0_b, W_1_b = _reconstruct_in_channel(channel, target, inters, table2, table1)
	return W_0_a + W_0_b, W_1_a + W_1_b

def reconstruct_ext(target, ptable1, ptable2, S=[], T=[], U=[]):
	W_0_S, W_1_S = reconstruct_in_channel_ext("S", target, S, ptable1, ptable2)
	W_0_T, W_1_T = reconstruct_in_channel_ext("T", target, T, ptable1, ptable2)
	W_0_U, W_1_U = reconstruct_in_channel_ext("U", target, U, ptable1, ptable2)

	return 16*sp.S.Pi*(
		W_0_S + 3*(t - u)*W_1_S
		+ W_0_T + 3*(u - s)*W_1_T
		+ W_0_U + 3*(t - s)*W_1_U
	)
