<filename>pinchon/shared/gaunt.py
import numpy as np

import operator as op

from sympy import simplify
from sympy import pi
from sympy import factorial, sqrt

from pinchon.shared.spherical import intSpherik, associatedLegendre, YY, SS


# Gaunt coefficients

# triangle functions
def istriangle(l1, l2, l3) -> bool:
	return (l2 + l3 >= l1 >= abs(l2 - l3) and
	        l1 + l3 >= l2 >= abs(l1 - l3) and
	        l1 + l2 >= l3 >= abs(l1 - l2)
	        )


def Delta(j1, j2, j3):
	# TODO: fix the factorial
	return np.sqrt(
		factorial(j1 + j2 - j3) * factorial(j1 - j2 + j3) * factorial(-j1 + j2 + j3) / factorial(
			j1 + j2 + j3 + 1))


# completely explicit definition
def gaunt(l1, l2, l3, m1, m2, m3):
	if not ((m1 + m2 + m3) == 0):
		return 0

	if (l1 + l2 + l3) % 2 != 0:
		return 0

	L = op.floordiv(l1 + l2 + l3, 2)

	kmin = max(-l3 + l1 + m2, -l3 + l2 - m1, 0)
	kmax = min(l2 + m2, l1 - m1, l1 + l2 - l3)

	sss = 0
	for k in range(kmin, kmax):
		sss += (-1) ** k / (factorial(k) * factorial(k + l3 - l1 - m2) * factorial(k + l3 - l2 + m1)
		                    * factorial(k + l3 - l2 + m1) * factorial(l1 - m1 - k) * factorial(
					l2 + m2 - k))

	res = Delta(l1, l2, l3) ** 2 * (-1) ** (L + l3 + m1 - m2) * sqrt((2 * l1 + 1) * (2 * l2 + 1) * (2 * l3 + 1)) * \
	      factorial(L) / (factorial(L - l1) * factorial(L - l2) * factorial(L - l3)) * \
	      sqrt(factorial(l1 + m1) * factorial(l1 - m1) * factorial(l2 + m2) *
	                factorial(l2 - m2) * factorial(l3 + m3) * factorial(l3 - m3)) * sss
	# return simplify(res)
	return simplify(res)


# mathematical definition
def defGaunt(l1, l2, l3, m1, m2, m3):
	return simplify(2 * sqrt(pi) * intSpherik(Y(l1, m1) * Y(l2, m2) * Y(l3, m3)))


# definition with Wigner' 3j symbols
def wigner3j(j1, j2, j3, m1, m2, m3):
	if not (istriangle(j1, j2, j3)) or not (m1 + m2 + m3 == 0):
		return 0

	kmax = min(j1 + j2 - j3, j1 - m1, j2 + m2)
	kmin = max(0, j1 - j3 + m2, j2 - j3 - m1)
	sss = 0

	for k in range(kmin, kmax):
		den = factorial(k) * factorial(j1 + j2 - j3 - k) * factorial(j1 - m1 - k) * \
		      factorial(j2 + m2 - k) * factorial(j3 - j2 + m1 + k) * factorial(j3 - j1 - m2 + k)
		sss = sss + (-1) ** k / den

	res = (-1) ** (j1 - j2 - m3) * sqrt(triangle(j1, j2, j3)) * sqrt(factorial(j1 + m1)
																		* factorial(j1 - m1) * factorial(
		j2 + m2) * factorial(j2 - m2) * factorial(j3 + m3) * factorial(j3 - m3)) * sss
	# TODO define triangle
	return simplify(res)


def wigner3j0(j1, j2, j3):
	sj = j1 + j2 + j3
	if sj % 2 != 0:
		return 0
	g = op.floordiv(sj, 2)
	res = (-1) ** g * sqrt(
		factorial(2 * (g - j1)) * factorial(2 * (g - j2)) * factorial(2 * (g - j3)) / factorial(
			2 * g + 1)) * factorial(g) / (factorial(g - j1) * factorial(g - j2) * factorial(g - j3))
	return simplify(res)


def gauntY(l1, l2, l3, m1, m2, m3):
	return simplify(sqrt((2 * l1 + 1) * (2 * l2 + 1) * (2 * l3 + 1) / (4 * pi))
	                * wigner3j0(l1, l2, l3)
	                * wigner3j(l1, l2, l3, m1, m2, m3)
	                )


# real Gaunt coefficients
# mathematical definition
def realGaunt(l1, l2, l3, m1, m2, m3):
	return intSpherik(SS(l1, m1) * SS(l2, m2) * SS(l3, m3))
