import numpy as np

from sympy import symbols, factorial, simplify, integrate, diff, exp
from sympy import cos, sin, sqrt, pi, I
from sympy import legendre
from sympy.abc import theta, phi, x


# with(orthopoly):


def intSpherik(f):
	"""
	Compute the integral of:
	.. math :: \int_0^{\pi}\int_0^{2\pi} f \mathrm{d}\phi \mathrm{d}\theta

	Parameters
	----------
	f

	Returns
	-------


	Notes
	-----
	theta, phi are sympy symbols
	"""
	# g = int(f, phi=0..2*pi);
	g = integrate(f, (phi, 0, 2*pi))
	# simplify(int_(g * sin(theta), theta=0..Pi))
	return simplify(integrate(g * sin(theta), (theta, 0, pi)))


def associatedLegendre(l: int, m: int):
	"""
	Compute the associated legendre as (with replacing x by cos(theta) in the derivative:

	.. math :: (-1)^m  \frac{1}{(2^l\cdot l!)} (\sin(\theta))^m
				\frac{\mathrm{d}^m P_n(x)}{\mathrm{d} x^m}\\
				x = \cos(\theta)

	Parameters
	----------
	l : int
	m : int

	Returns
	-------
	associated_legendre :

	Notes
	-----
	Result is simplify compared to the maple script

	x, theta are sympy symbols

	Raises
	------
	ArithemticError if m < 0

	"""
	# TODO: look into sympy.assoc_legendre it seems behave the same, and it is way faster
	if l == 0:
		return 1
	if m < 0:
		raise ArithmeticError(f"m should be non negative, m is {m}")

	associated_legendre = (-1) ** m * 1 / (2 ** l * factorial(l)) * (sin(theta)) ** m * \
	diff(((x**2 - 1)**l), x, l+m) .subs(x, cos(theta))
	# subs(x=cos(theta), diff((x ^ 2 - 1) ^ l, seq(x, i=1..(l + m))))
	# f.subs(x, cos(theta))

	return associated_legendre


def YY(l: int, m: int):
	"""

	Parameters
	----------
	l
	m

	Returns
	-------

	Notes
	-----
	x, theta are sympy symbols
	"""
	# TODO: check
	if l == 0 and m == 0:
		return 1 / (4 * np.pi)

	mm = abs(m)

	res = I ** (m + mm) * \
		simplify(sqrt((2*l + 1) * factorial(l - mm) / factorial(l + mm) / 4 / pi) * 1/2 ** l / factorial(l)) * \
		sin(theta) ** mm * exp(I * m * phi) * \
		diff(((x**2 - 1) ** l), x, l + mm).subs(x, cos(theta))
	# subs(x=cos(theta), diff( (x * 2 - 1) ** l, x$(l + mm)))
	# f.subs(x, cos(theta))

	return res


def SS(l: int, m: int):
	"""

	Parameters
	----------
	l
	m

	Returns
	-------

	Notes
	-----
	phi is a sympy symbol
	"""

	if m == 0:
		return sqrt((2 * l + 1) / (4 * pi)) * legendre(l, cos(theta))

	cte = sqrt((2 * + l + 1) * factorial(l - abs(m)) / factorial(l + abs(m)) / 2 / pi)

	if m > 0:
		# note : it was associatedLegendre(l, m, theta) in the maple code
		return cte * associatedLegendre(l, m) * cos(m * phi)
	return cte * associatedLegendre(l, -m) * sin(-m * phi)
