from displib import *
import re
A_x, B_x = sp.symbols("A_x B_x")
C_x, D_x = sp.symbols("C_x D_x")
F_pi = sp.symbols("F_pi")

A_Op2 = pp_subthreshold_poly(1)

ptable_Op2 = pp_ptable_from_form_factor(A_Op2)


A_Op4 = -reconstruct(
scatter("00", "pm"), ptable_Op2,
S=["00", "pm"], T=["m0"], U=["p0"],
)

A_Op4 = A_Op4 + A_x + B_x/(F_pi**2) *(s - s_0) + C_x/(F_pi**4) *(s - s_0)**2 + D_x/(F_pi**4) *((t - s_0)**2 + (u - s_0)**2)

ptable_Op4 = pp_ptable_from_form_factor(A_Op4)


A_00_S, A_00_P = lookup_SP_wave(scatter("00", "00"), ptable_Op4)
A_x_S, A_x_P = lookup_SP_wave(scatter("00", "pm"), ptable_Op4)
A_pm_S, A_pm_P = lookup_SP_wave(scatter("pm", "pm"), ptable_Op4)

psi_0 = sp.cancel(A_00_S - 2*A_x_S)
psi_1 = sp.cancel(A_pm_S*2)
psi_2 = sp.cancel(A_00_S + A_x_S)
                  
def derive_xi(psi):
    xi_ = sp.cancel(attempt_poly_reduce(psi * sigma(s, M_pi), s))
    return normalize_k4(xi_, s, M_pi)

analyze_table(derive_xi(psi_0), s)
analyze_table(derive_xi(psi_1), s)

A_x, B_x = sp.symbols("A_x B_x")
A_00 = sp.symbols("A_00")
ptable_eta_Op2 = {
scatter("E0", "pm"): A_x + B_x*(s-s_0)/F_pi**2,
scatter("E0", "00"): A_00,
}
joined_ptable = {**ptable_eta_Op2, **ptable_Op2}
E0_pm_Op4 = reconstruct(
scatter("E0", "pm"), joined_ptable,
S=["00", "pm"], T=["m0"], U=["p0"],
)
E0_00_Op4 = reconstruct(
scatter("E0", "00"), joined_ptable,
S=["00", "pm"], T=["00", "pm"], U=["00", "pm"],
)
ptable_eta_Op4 = {
scatter("E0", "pm"): E0_pm_Op4,
scatter("E0", "00"): E0_00_Op4
}
joined_ptable_Op4 = {**ptable_eta_Op4, **ptable_Op4}
E0_pm_Op6 = reconstruct_ext(
scatter("E0", "pm"), joined_ptable, joined_ptable_Op4,
S=["00", "pm"], T=["0m"], U=["0p"],
)
