
from sympy import S, sqrt, Sum, symbols
from sympy.physics.quantum.cg import Wigner3j, Wigner6j, Wigner9j, CG, cg_simp
from sympy.functions.special.tensor_functions import KroneckerDelta


def test_cg_simp_add():
    j, m1, m1p, m2, m2p = symbols('j m1 m1p m2 m2p')
    # Test Varshalovich 8.7.1 Eq 1
    a = CG(S(1)/2,S(1)/2,0,0,S(1)/2,S(1)/2)
    b = CG(S(1)/2,-S(1)/2,0,0,S(1)/2,-S(1)/2)
    c = CG(1,1,0,0,1,1)
    d = CG(1,0,0,0,1,0)
    e = CG(1,-1,0,0,1,-1)
    assert cg_simp(a+b) == 2
    assert cg_simp(c+d+e) == 3
    assert cg_simp(a+b+c+d+e) == 5
    assert cg_simp(a+b+c) == 2+c
    assert cg_simp(2*a+b) == 2+a
    assert cg_simp(2*c+d+e) == 3+c
    assert cg_simp(5*a+5*b) == 10
    assert cg_simp(5*c+5*d+5*e) == 15
    assert cg_simp(-a-b) == -2
    assert cg_simp(-c-d-e) == -3
    assert cg_simp(-6*a-6*b) == -12
    assert cg_simp(-4*c-4*d-4*e) == -12
    a = CG(S(1)/2,S(1)/2,j,0,S(1)/2,S(1)/2)
    b = CG(S(1)/2,-S(1)/2,j,0,S(1)/2,-S(1)/2)
    c = CG(1,1,j,0,1,1)
    d = CG(1,0,j,0,1,0)
    e = CG(1,-1,j,0,1,-1)
    assert cg_simp(a+b) == 2*KroneckerDelta(j,0)
    assert cg_simp(c+d+e) == 3*KroneckerDelta(j,0)
    assert cg_simp(a+b+c+d+e) == 5*KroneckerDelta(j,0)
    assert cg_simp(a+b+c) == 2*KroneckerDelta(j,0)+c
    assert cg_simp(2*a+b) == 2*KroneckerDelta(j,0)+a
    assert cg_simp(2*c+d+e) == 3*KroneckerDelta(j,0)+c
    assert cg_simp(5*a+5*b) == 10*KroneckerDelta(j,0)
    assert cg_simp(5*c+5*d+5*e) == 15*KroneckerDelta(j,0)
    assert cg_simp(-a-b) == -2*KroneckerDelta(j,0)
    assert cg_simp(-c-d-e) == -3*KroneckerDelta(j,0)
    assert cg_simp(-6*a-6*b) == -12*KroneckerDelta(j,0)
    assert cg_simp(-4*c-4*d-4*e) == -12*KroneckerDelta(j,0)
    # Test Varshalovich 8.7.1 Eq 2
    a = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,0,0)
    b = CG(S(1)/2,-S(1)/2,S(1)/2,S(1)/2,0,0)
    c = CG(1,1,1,-1,0,0)
    d = CG(1,0,1,0,0,0)
    e = CG(1,-1,1,1,0,0)
    assert cg_simp(a-b) == sqrt(2)
    assert cg_simp(c-d+e) == sqrt(3)
    assert cg_simp(a-b+c-d+e) == sqrt(2)+sqrt(3)
    assert cg_simp(a-b+c) == sqrt(2)+c
    assert cg_simp(2*a-b) == sqrt(2)+a
    assert cg_simp(2*c-d+e) == sqrt(3)+c
    assert cg_simp(5*a-5*b) == 5*sqrt(2)
    assert cg_simp(5*c-5*d+5*e) == 5*sqrt(3)
    assert cg_simp(-a+b) == -sqrt(2)
    assert cg_simp(-c+d-e) == -sqrt(3)
    assert cg_simp(-6*a+6*b) == -6*sqrt(2)
    assert cg_simp(-4*c+4*d-4*e) == -4*sqrt(3)
    a = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,j,0)
    b = CG(S(1)/2,-S(1)/2,S(1)/2,S(1)/2,j,0)
    c = CG(1,1,1,-1,j,0)
    d = CG(1,0,1,0,j,0)
    e = CG(1,-1,1,1,j,0)
    assert cg_simp(a-b) == sqrt(2)*KroneckerDelta(j,0)
    assert cg_simp(c-d+e) == sqrt(3)*KroneckerDelta(j,0)
    assert cg_simp(a-b+c-d+e) == sqrt(2)*KroneckerDelta(j,0)+sqrt(3)*KroneckerDelta(j,0)
    assert cg_simp(a-b+c) == sqrt(2)*KroneckerDelta(j,0)+c
    assert cg_simp(2*a-b) == sqrt(2)*KroneckerDelta(j,0)+a
    assert cg_simp(2*c-d+e) == sqrt(3)*KroneckerDelta(j,0)+c
    assert cg_simp(5*a-5*b) == 5*sqrt(2)*KroneckerDelta(j,0)
    assert cg_simp(5*c-5*d+5*e) == 5*sqrt(3)*KroneckerDelta(j,0)
    assert cg_simp(-a+b) == -sqrt(2)*KroneckerDelta(j,0)
    assert cg_simp(-c+d-e) == -sqrt(3)*KroneckerDelta(j,0)
    assert cg_simp(-6*a+6*b) == -6*sqrt(2)*KroneckerDelta(j,0)
    assert cg_simp(-4*c+4*d-4*e) == -4*sqrt(3)*KroneckerDelta(j,0)
    # Test Varshalovich 8.7.2 Eq 9
    # alpha=alphap,beta=betap case
    # numerical
    a = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,1,0)**2
    b = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,0,0)**2
    c = CG(1,0,1,1,1,1)**2
    d = CG(1,0,1,1,2,1)**2
    assert cg_simp(a+b) == 1
    assert cg_simp(c+d) == 1
    assert cg_simp(a+b+c+d) == 2
    assert cg_simp(4*a+4*b) == 4
    assert cg_simp(4*c+4*d) == 4
    assert cg_simp(5*a+3*b) == 3+2*a
    assert cg_simp(5*c+3*d) == 3+2*c
    assert cg_simp(-a-b) == -1
    assert cg_simp(-c-d) == -1
    # symbolic
    a = CG(S(1)/2,m1,S(1)/2,m2,1,1)**2
    b = CG(S(1)/2,m1,S(1)/2,m2,1,0)**2
    c = CG(S(1)/2,m1,S(1)/2,m2,1,-1)**2
    d = CG(S(1)/2,m1,S(1)/2,m2,0,0)**2
    assert cg_simp(a+b+c+d) == 1
    assert cg_simp(4*a+4*b+4*c+4*d) == 4
    assert cg_simp(3*a+5*b+3*c+4*d) == 3+2*b+d
    assert cg_simp(-a-b-c-d) == -1
    a = CG(1,m1,1,m2,2,2)**2
    b = CG(1,m1,1,m2,2,1)**2
    c = CG(1,m1,1,m2,2,0)**2
    d = CG(1,m1,1,m2,2,-1)**2
    e = CG(1,m1,1,m2,2,-2)**2
    f = CG(1,m1,1,m2,1,1)**2
    g = CG(1,m1,1,m2,1,0)**2
    h = CG(1,m1,1,m2,1,-1)**2
    i = CG(1,m1,1,m2,0,0)**2
    assert cg_simp(a+b+c+d+e+f+g+h+i) == 1
    assert cg_simp(4*(a+b+c+d+e+f+g+h+i)) == 4
    assert cg_simp(a+b+2*c+d+4*e+f+g+h+i) == 1+c+3*e
    assert cg_simp(-a-b-c-d-e-f-g-h-i) == -1
    # alpha!=alphap or beta!=betap case
    # numerical
    a = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,1,0)*CG(S(1)/2,-S(1)/2,S(1)/2,S(1)/2,1,0)
    b = CG(S(1)/2,S(1)/2,S(1)/2,-S(1)/2,0,0)*CG(S(1)/2,-S(1)/2,S(1)/2,S(1)/2,0,0)
    c = CG(1,1,1,0,2,1)*CG(1,0,1,1,2,1)
    d = CG(1,1,1,0,1,1)*CG(1,0,1,1,1,1)
    assert cg_simp(a+b) == 0
    assert cg_simp(c+d) == 0
    # symbolic
    a = CG(S(1)/2,m1,S(1)/2,m2,1,1)*CG(S(1)/2,m1p,S(1)/2,m2p,1,1)
    b = CG(S(1)/2,m1,S(1)/2,m2,1,0)*CG(S(1)/2,m1p,S(1)/2,m2p,1,0)
    c = CG(S(1)/2,m1,S(1)/2,m2,1,-1)*CG(S(1)/2,m1p,S(1)/2,m2p,1,-1)
    d = CG(S(1)/2,m1,S(1)/2,m2,0,0)*CG(S(1)/2,m1p,S(1)/2,m2p,0,0)
    assert cg_simp(a+b+c+d) == KroneckerDelta(m1,m1p)*KroneckerDelta(m2,m2p)
    a = CG(1,m1,1,m2,2,2)*CG(1,m1p,1,m2p,2,2)
    b = CG(1,m1,1,m2,2,1)*CG(1,m1p,1,m2p,2,1)
    c = CG(1,m1,1,m2,2,0)*CG(1,m1p,1,m2p,2,0)
    d = CG(1,m1,1,m2,2,-1)*CG(1,m1p,1,m2p,2,-1)
    e = CG(1,m1,1,m2,2,-2)*CG(1,m1p,1,m2p,2,-2)
    f = CG(1,m1,1,m2,1,1)*CG(1,m1p,1,m2p,1,1)
    g = CG(1,m1,1,m2,1,0)*CG(1,m1p,1,m2p,1,0)
    h = CG(1,m1,1,m2,1,-1)*CG(1,m1p,1,m2p,1,-1)
    i = CG(1,m1,1,m2,0,0)*CG(1,m1p,1,m2p,0,0)
    assert cg_simp(a+b+c+d+e+f+g+h+i) == KroneckerDelta(m1,m1p)*KroneckerDelta(m2,m2p)

def test_cg_simp_sum():
    x, a, b, c, cp, alpha, beta, gamma, gammap = symbols('x a b c cp alpha beta gamma gammap')
    # Varshalovich 8.7.1 Eq 1
    assert cg_simp(x * Sum(CG(a,alpha,b,0,a,alpha), (alpha,-a,a))) == x*(2*a+1)*KroneckerDelta(b,0)
    assert cg_simp(x * Sum(CG(a,alpha,b,0,a,alpha), (alpha,-a,a))+CG(1,0,1,0,1,0)) == x*(2*a+1)*KroneckerDelta(b,0)+CG(1,0,1,0,1,0)
    assert cg_simp(2 * Sum(CG(1,alpha,0,0,1,alpha), (alpha,-1,1))) == 6
    # Varshalovich 8.7.1 Eq 2
    assert cg_simp(x*Sum((-1)**(a-alpha) * CG(a,alpha,a,-alpha,c,0), (alpha,-a,a))) == x*sqrt(2*a+1)*KroneckerDelta(c,0)
    assert cg_simp(3*Sum((-1)**(2-alpha) * CG(2,alpha,2,-alpha,0,0), (alpha,-2,2))) == 3*sqrt(5)
    # Varshalovich 8.7.2 Eq 4
    assert cg_simp(Sum(CG(a,alpha,b,beta,c,gamma)*CG(a,alpha,b,beta,cp,gammap),(alpha,-a,a),(beta,-b,b))) == KroneckerDelta(c,cp)*KroneckerDelta(gamma,gammap)
    assert cg_simp(Sum(CG(a,alpha,b,beta,c,gamma)*CG(a,alpha,b,beta,c,gammap),(alpha,-a,a),(beta,-b,b))) == KroneckerDelta(gamma,gammap)
    assert cg_simp(Sum(CG(a,alpha,b,beta,c,gamma)*CG(a,alpha,b,beta,cp,gamma),(alpha,-a,a),(beta,-b,b))) == KroneckerDelta(c,cp)
    assert cg_simp(Sum(CG(a,alpha,b,beta,c,gamma)**2,(alpha,-a,a),(beta,-b,b))) == 1
    assert cg_simp(Sum(CG(2,alpha,1,beta,2,gamma)*CG(2,alpha,1,beta,2,gammap), (alpha,-2,2), (beta,-1,1))) == KroneckerDelta(gamma,gammap)

def test_doit():
    assert Wigner3j(1/2,-1/2,1/2,1/2,0,0).doit() == -sqrt(2)/2
    assert Wigner6j(1,2,3,2,1,2).doit() == sqrt(21)/105
    assert Wigner9j(2,1,1,S(3)/2,S(1)/2,1,S(1)/2,S(1)/2,0).doit() == sqrt(2)/12
    assert CG(1/2,1/2,1/2,-1/2,1,0).doit() == sqrt(2)/2
