import unittest
from .sym import *
from .expr import *
from .func import *
from .poly import *
from .kwaves import lookup_5_27_kwave_sum

class ProcessLabel:
    pass

def _part_conj_one(c):
    if c == "p":
        return "m"
    elif c == "m":
        return "p"
    elif c == "0":
        return "0"
    elif c == "A":
        return "a"
    elif c == "B":
        return "b"
    elif c == "a":
        return "A"
    elif c == "b":
        return "B"
    else:
        return c

def _part_conj(s):
    conj = "".join([_part_conj_one(c) for c in s])
    invs = sum([c in "pm" for c in s])

    return conj, (-sp.S.One) ** invs

class ScatterProcess(ProcessLabel):
    '''
    Class ScatterProcess is an abstract representation of a scattering
    process, specifying the ingoing and outgoing kind of particles,
    that is, it specifies the A, B, C, D in

        AB->CD
    .

    The particle kinds are represented by a character in the following
    correspondence:

        p  positive pion
        n  negative pion
        0  neutral pion
        E  eta meson

    To consider two ScatterProcess'es equal, the particle ordering
    must match, up to swapping of both ingoing and outgoing particles.
    So these two are equal

        AB->CD  BA->DC

    but these two aren't

        AB->CD  AB->DC
    .

    In addition, ScatterProcess instance may carry a phase, which
    means, should an amplitude of this process be evaluated,
    additional factor of this phase should be inserted. This phase
    factor is set on crossing (see .T and .U property methods on
    this class).
    '''

    def __init__(self, ab, cd, phase=sp.S.One):
        self.ab, self.cd = ab, cd
        self.phase = phase

    def __eq__(self, other):
        return isinstance(other, type(self)) and \
            self.phase == other.phase and \
            ((self.ab == other.ab and self.cd == other.cd) \
            or (self.ab[::-1] == other.ab and self.cd[::-1] == other.cd))

    @property
    def wo_phase(self):
        '''Returns the same process but drops any phase it may carry.'''
        return type(self)(self.ab, self.cd)

    @property
    def rev(self):
        '''Returns a reversed process, i.e. CD->AB from AB->CD.'''
        cd, _ = _part_conj(self.cd)
        ab, _ = _part_conj(self.ab)
        return type(self)(cd, ab, self.phase)

    @property
    def tu(self):
        '''Returns process obtained by swapping the outgoing particles.'''
        return type(self)(self.ab, self.cd[::-1], self.phase)

    @property
    def a(self):
        '''Returns the A symbol from process AB->CD.'''
        return self.ab[0]

    @property
    def b(self):
        '''Returns the B symbol from process AB->CD.'''
        return self.ab[1]

    @property
    def c(self):
        '''Returns the C symbol from process AB->CD.'''
        return self.cd[0]

    @property
    def d(self):
        '''Returns the D symbol from process AB->CD.'''
        return self.cd[1]

    @property
    def parts(self):
        '''Returns string with particle symbols of the process.'''
        return self.ab + self.cd

    @property
    def abcd(self):
        '''Returns string with particle symbols of the process.'''
        return self.ab + self.cd

    @property
    def U(self):
        '''
        Returns a U-crossed process.

        Suppose we have a process AB -> CD, we then do
        crossing into A'B'->C'D' by

           A -> A'      s' = (p_A' + p_B')^2 = (p_A + p_D)^2 = u
           B -> C'      t' = (p_A' + p_C')^2 = (p_A + p_B)^2 = s 
           C -> D'      u' = (p_A' + p_D')^2 = (p_A + p_C)^2 = t
           D -> B'

        .

        Phase is set accordingly.
        '''
        d_conj, dph = _part_conj(self.d)
        b_conj, bph = _part_conj(self.b)
        ab = self.a + d_conj
        cd = b_conj + self.c
        return type(self)(ab, cd, phase=self.phase*dph*bph)

    @property
    def T(self):
        '''
        Returns a T-crossed process

        This is inverse of .U, phase is set accordingly.
        '''
        c_conj, cph = _part_conj(self.c)
        b_conj, bph = _part_conj(self.b)
        ab = self.a + c_conj
        cd = self.d + b_conj
        return type(self)(ab, cd, phase=self.phase*cph*bph)

    ARROW = '→'

    def __str__(self):
        return f"{self.ab}{self.ARROW}{self.cd}" + \
                (f"(phase: {self.phase})" if self.phase != sp.S.One else "")

    def __repr__(self):
        return f"{self.ab}{self.ARROW}{self.cd}" + \
                (f"(phase: {self.phase})" if self.phase != sp.S.One else "")

    @property
    def tex_label(self):
        transl = {
            "p": "\\pi^{+}",
            "m": "\\pi^{-}",
            "0": "\\pi^{0}"
        }
        return f"{transl.get(self.a, self.a)} {transl.get(self.b, self.b)}" \
               f" {self.ARROW} {transl.get(self.c, self.c)} {transl.get(self.d, self.d)}"

    def __hash__(self):
        return hash((frozenset(self.ab), frozenset(self.cd)))

def scatter(ab, cd):
    return ScatterProcess(ab, cd)

def _basic_lookup(label, table):
    if label in table:
        return table[label]

    if label.rev in table:
        return table[label.rev]

def lookup_process(label, table):
    subs_ident   = {}
    subs_tu      = {t: u, u: t}
    subs_T_cross = {s: u, t: s, u: t}
    subs_U_cross = {s: t, t: u, u: s}

    for l1, l1s in zip([label, label.tu],
                    [subs_ident, subs_tu]):
        for l2, l2s in zip([l1, l1.T, l1.U],
                    [subs_ident, subs_U_cross, subs_T_cross]):
            if (found := _basic_lookup(l2.wo_phase, table)) is not None:
                return found.subs(l2s, simultaneous=True) \
                        .subs(l1s, simultaneous=True) * l2.phase

    raise ValueError("process not found im table: ", label)

def m_for_part1(ident):
    if ident in "pm0":
        return M_pi
    elif ident == "E":
        return M_eta
    else:
        raise NotImplementedError(ident)

def m_for_part(*parts):
    if len(parts) == 1:
        return m_for_part1(parts[0])
    else:
        return [m_for_part1(part) for part in parts]

def extract_SP_wave(expr, plabel, in_ks=True):
    t_u = sp.Dummy()

    m_A, m_B, m_C, m_D \
        = [m_for_part(x) for x in plabel.parts]

    if len(set([m_A, m_B, m_C, m_D])) == 1:
        expr_ = sp.expand(expr).subs({
                t: sp.S.Half*(-sigma(s,m_A)**2*s + (t_u)),
                u: sp.S.Half*(-sigma(s,m_A)**2*s - (t_u))
        }, simultaneous=True)
    else:
        s_0_val = m_A**2 + m_B**2 + m_C**2 + m_D**2
        expr_ = sp.expand(expr).subs({
                t: sp.S.Half*(s_0_val-s + (t_u)),
                u: sp.S.Half*(s_0_val-s - (t_u))
        }, simultaneous=True)

    S, P = sp.S.Zero, sp.S.Zero

    for decomp in single_out_special_terms(expr_, t_u):
        poly, special = decomp[0], decomp[1:]
        if not special:
            cos_theta = sp.Dummy()

            if len(set([m_A, m_B, m_C, m_D])) == 1:
                poly_ = poly.subs({
                    t_u: (cos_theta*s*sigma(s, m_A)**2) 
                })
            else:
                poly_ = poly.subs({
                    t_u: (cos_theta*lambda_sqrt(s, m_A, m_B)*lambda_sqrt(s, m_C, m_D) \
                        - (m_A**2 - m_B**2)*(m_C**2 - m_D**2))/s
                })

            dp = legendre_decompose(sp.Poly(poly_, cos_theta)) + [0, 0]
            S += dp[0] / sp.sqrt(2)
            P += dp[1] / sp.sqrt(6)
        elif len(special) == 1 and isinstance(special[0], JBar) \
                and len(set([m_A, m_B, m_C, m_D])) == 1:
            assert len(set([m_A, m_B, m_C, m_D])) == 1
            jbar = special[0]

            # reparametrize the integrand in terms of the argument to JBar
            xi = sp.Dummy("xi")
            sol = sp.linsolve([jbar.args[0] - xi], (t_u,))
            if len(sol) != 1 or jbar.args[1] != m_A:
                raise NotImplementedError(jbar)
            tu_sub = list(sol)[0][0]

            ppoly = (poly.subs({t_u: tu_sub}) * sp.diff(tu_sub, xi)).as_poly(xi)
            for order, coeff in enumerate(reversed(ppoly.all_coeffs())):
                primi = JBar_PRIMIS[order]
                sigma2 = sigma(s,m_A)**2 # 1 - 4m**2/s
                a = jbar.args[0].subs({t_u: -s*sigma2})
                b = jbar.args[0].subs({t_u: s*sigma2})
                S += sp.S.Half / (s*sigma2) * coeff \
                        * (primi(b, m_A) - primi(a, m_A))/(16*pi**2)

            ppoly = (poly.subs({t_u: tu_sub}) * tu_sub * sp.diff(tu_sub, xi)).as_poly(xi)
            for order, coeff in enumerate(reversed(ppoly.all_coeffs())):
                primi = JBar_PRIMIS[order]
                sigma2 = sigma(s,m_A)**2
                a = jbar.args[0].subs({t_u: -s*sigma2})
                b = jbar.args[0].subs({t_u: s*sigma2})

                P += sp.S.Half / (s*sigma2) **2 * coeff \
                        * (primi(b, m_A) - primi(a, m_A))/(16*pi**2)
        elif len(special) == 1 and isinstance(special[0], JBar) \
                and len(set([m_A, m_B, m_C, m_D])) == 2:

            jbar = special[0]

            if jbar.args[1] not in set([m_A, m_B, m_C, m_D]):
                raise NotImplementedError(jbar)

            m_a = jbar.args[1]
            m_b = (set([m_A, m_B, m_C, m_D]) - set([m_a])).pop()

            xi = sp.Dummy("xi")
            sol = sp.linsolve([jbar.args[0] - xi], (t_u,))
            if len(sol) != 1:
                raise NotImplementedError(jbar)
            tu_sub = list(sol)[0][0]

            ppoly = (poly.subs({t_u: tu_sub}) * sp.diff(tu_sub, xi)).as_poly(xi)
            for order, coeff in enumerate(reversed(ppoly.all_coeffs())):
                ratio = lambda_sqrt(s, m_b, m_a)*lambda_sqrt(s, m_a, m_a)/s  
                if jbar.args[0] == t_u/2 + m_b**2/2 + 3*m_a**2/2 - s/2:
                    S += sp.S.Half / ratio \
                            * lookup_5_27_kwave_sum(order, m_a, m_b) * coeff / sp.S.Pi
                elif jbar.args[0] == -t_u/2 + m_b**2/2 + 3*m_a**2/2 - s/2:
                    S += -sp.S.Half / ratio \
                            * lookup_5_27_kwave_sum(order, m_a, m_b) * coeff / sp.S.Pi
                else:
                    raise NotImplementedError()

            ppoly = (poly.subs({t_u: tu_sub}) * tu_sub * sp.diff(tu_sub, xi)).as_poly(xi)
            for order, coeff in enumerate(reversed(ppoly.all_coeffs())):
                ratio = lambda_sqrt(s, m_b, m_a)*lambda_sqrt(s, m_a, m_a)/s
                if jbar.args[0] == t_u/2 + m_b**2/2 + 3*m_a**2/2 - s/2:
                    P += sp.S.Half / ratio**2 \
                            * lookup_5_27_kwave_sum(order, m_a, m_b) * coeff / sp.S.Pi
                elif jbar.args[0] == -t_u/2 + m_b**2/2 + 3*m_a**2/2 - s/2:
                    P += -sp.S.Half / ratio**2 \
                            * lookup_5_27_kwave_sum(order, m_a, m_b) * coeff / sp.S.Pi
                else:
                    raise NotImplementedError()
        else:
            raise NotImplementedError(special, poly)

    S = sp.cancel(S/(16*pi))
    P = sp.cancel(P/(16*pi))

    if in_ks:
        S, P = expand_in_ks(S), expand_in_ks(P)

    S = S.replace(lambda expr: type(expr) in [k_0, k_1, k_2, k_3, k_4, sigma],
                    lambda expr: type(expr)(*[sp.cancel(sp.expand_func(arg)) \
                                             for arg in expr.args]))
    P = P.replace(lambda expr: type(expr) in [k_0, k_1, k_2, k_3, k_4, sigma],
                    lambda expr: type(expr)(*[sp.cancel(sp.expand_func(arg)) \
                                             for arg in expr.args]))

    return S, P

def lookup_SP_wave(plabel, table, **kwargs):
    return extract_SP_wave(lookup_process(plabel, table), plabel, **kwargs)

def render_ptable(table, simplify=True):
    return sp.Array([
        [sp.Dummy(proc.tex_label), sp.simplify(ampl) if simplify else ampl]
        for proc, ampl in table.items()
    ])

class TestProcessLabels(unittest.TestCase):
    def test_scatter_crossings(self):
        p = ScatterProcess("AB", "CD")

        self.assertEqual(p.T.U, p)
        self.assertEqual(p.U.T, p)
        self.assertEqual(p.U.U.U, p)
        self.assertEqual(p.rev.rev, p)
        self.assertEqual(p.tu.tu, p)

        self.assertEqual(
            ScatterProcess("pm", "00").T,
            ScatterProcess("p0", "0p", phase=-1)
        )
        self.assertEqual(
            ScatterProcess("pm", "00").U,
            ScatterProcess("p0", "p0", phase=-1)
        )

        self.assertNotEqual(p.T, p)
        self.assertNotEqual(p.U, p)
        self.assertNotEqual(p.tu, p)
        self.assertNotEqual(p.rev, p)
