import unittest
import sympy as sp
from .sym import *

def manifest_tu_antisym(poly, s0_val=None):
    if type(poly) is dict:
        # We were given a process table,
        # apply to amplitudes individually
        return {p: manifest_tu_antisym(a, s0_val=s0_val)
                for p, a in poly.items()}

    ret = sp.simplify(poly.subs({
        t: sp.S.Half*(3*s_0-s + (t-u)),
        u: sp.S.Half*(3*s_0-s - (t-u))
    }, simultaneous=True))

    if s0_val is not None:
        ret = ret.subs({s_0: s0_val})
    return ret

def legendre_compose(coord, gen):
    return sum([
        sp.legendre_poly(deg, gen, polys=True) \
            * sp.sqrt((sp.Integer(deg)*2+1)/2) * c
        for deg, c in enumerate(coord)
    ])

def legendre_decompose(poly):
    poly_ = sp.Poly(poly)
    coord = [sp.S.Zero] * (poly.degree() + 1)
    while (deg := poly_.degree()) >= 0:
        legnorm = sp.legendre_poly(deg, poly.gen, polys=True) \
                    * sp.sqrt((sp.Integer(deg)*2+1)/2)
        lambd = poly_.LC() / legnorm.LC()
        poly_ -= legnorm*lambd
        assert poly_.degree() < deg
        coord[deg] = lambd
    return coord

class TestLegendre(unittest.TestCase):
    def test_decompose_compose(self):
        x = sp.Dummy()

        for p in [1, x, x**2+x**3, x**4, x**7]:
            poly = sp.Poly(p, x)
            coords = legendre_decompose(poly)
            poly_redone = legendre_compose(coords, x)
            self.assertEqual(poly.coeffs(), poly_redone.coeffs())
