from .process import *

def pp_subthreshold_poly(degree, subscript=None):
    '''
    Generates a polynomial of given degree for use as subtraction polynomial
    in reconstruction of pi-pi scattering in isospin limit.

    The function returns an expression in terms of the global symbols 'alpha',
    'beta', 'lambda1', 'lambda2' and 'M_pi'. If a subscript is given as argument,
    it instead uses indexed symbols 'alpha_sub', 'beta_sub', 'lambda1_sub',
    'lambda2_sub' and attaches the given subscript to them.
    '''
    ret = sp.S.Zero
    if degree >= 1:
        ret += alpha*M_pi**2/(3*F_pi**2) + beta/F_pi**2*(s-4*M_pi**2/3)
    if degree >= 2:
        ret += lambda1/F_pi**4*(s-2*M_pi**2) \
                + lambda2/F_pi**4*(t-2*M_pi**2) \
                + lambda2/F_pi**4*(u-2*M_pi**2)
    if degree < 1 or degree >= 3:
        raise NotImplementedError("degree %d subthreshold parametrization "
                                  "not implemented" % degree)
    if subscript is not None:
        ret = ret.subs({
            alpha: alpha_sub[subscript],
            beta: beta_sub[subscript],
            lambda1: lambda1_sub[subscript],
            lambda2: lambda2_sub[subscript],
        })
    return ret

def pp_ptable_from_form_factor(A):
    '''Returns a minimal process table of isospin-symmetric pi-pi scattering
    expressed in terms of form factor A'''

    B = A.subs({s: t, t: s}, simultaneous=True)
    C = B.subs({t: u, u: t}, simultaneous=True)

    return {
        scatter("pm", "00"): -A,
        scatter("pm", "pm"): A + B,
        scatter("00", "00"): A + B + C,
    }
