"""
Module implementing or wrapping special functions required for higher level
classes and functions of the spharpy package.
"""

from itertools import count
import numpy as np
import scipy.special as _spspecial
from scipy.optimize import brentq


def spherical_bessel(n, z, derivative=False):
    r"""
    Spherical Bessel function of the first kind and of order n evaluated at z.

    The spherical Bessel function of the first kind is defined as
    ([#]_, Eq. 2.29)

    .. math::

        j_n(z) = \sqrt{\frac{\pi}{2z}} J_{n+\frac{1}{2}} (z)

    with the Bessel function of the first kind :math:`J_\alpha`.

    Parameters
    ----------
    n : int, ndarray
        Order of the spherical Bessel function
    z : double, ndarray
        Argument of the spherical Bessel function. Has to be real valued.
    derivative : bool
        Return the derivative of the spherical Bessel function. Default is
        ``False``.


    Returns
    -------
    jn : double, ndarray
        Spherical Bessel function. Array with dimensions [N x Z], where N is
        the number of elements in n and Z is the number of elements in z.

    Note
    ----
    This is a wrapper around the Scipy implementation of the
    :py:func:`spherical Bessel function <scipy.special.spherical_jn>`.

    References
    ----------
    .. [#] Rafaely, B. (2019). Fundamentals of spherical array processing
           (2nd ed.). Springer. https://doi.org/10.1007/978-3-319-99561-8
    """

    ufunc = _spspecial.spherical_jn
    n = np.asarray(n, dtype=int)
    z = np.asarray(z, dtype=float)

    bessel = np.zeros((n.size, z.size), dtype=complex)

    if n.size > 1:
        for idx, order in zip(count(), n):
            bessel[idx, :] = ufunc(order, z, derivative=derivative)
    else:
        bessel = ufunc(n, z, derivative=derivative)

    if z.ndim <= 1 or n.ndim <= 1:
        bessel = np.squeeze(bessel)

    return bessel


def spherical_bessel_zeros(n_max, n_zeros):
    r"""Compute the zeros of the spherical Bessel function.

    This function will always start at order zero which is equal
    to :math:`\sin(z)/z` and iteratively compute the roots for higher orders.
    The roots are computed using Brents algorithm from the scipy package. See
    :py:func:`~spherical_bessel` for more information.

    Parameters
    ----------
    n_max : int
        The order of the spherical bessel function
    n_zeros : int
        The number of roots to be computed

    Returns
    -------
    roots : ndarray, double
        The roots of the spherical bessel function

    """
    def func(x, n):
        return _spspecial.spherical_jn(n, x)

    zerosj = np.zeros((n_max+1, n_zeros), dtype=float)
    zerosj[0] = np.arange(1, n_zeros+1)*np.pi
    points = np.arange(1, n_zeros+n_max+1)*np.pi

    roots = np.zeros(n_zeros+n_max, dtype=float)
    for i in range(1, n_max+1):
        for j in range(n_zeros+n_max-i):
            roots[j] = brentq(func, points[j], points[j+1], (i,), maxiter=5000)
        points = roots
        zerosj[i, :n_zeros] = roots[:n_zeros]

    return zerosj


def spherical_hankel(n, z, kind=2, derivative=False):
    r"""
    Spherical Hankel function of the first kind and order n evaluated at z.

    The spherical Hankel function of the first kind is defined as
    ([#]_, Eq. 2.30)

    .. math::

        h_n(z) = \sqrt{\frac{\pi}{2z}} H_{n+\frac{1}{2}} (z)

    with the Hankel function of the first kind :math:`H_\alpha`.

    Parameters
    ----------
    n : int, ndarray
        Order of the spherical bessel function
    z : double, ndarray
        Argument of the spherical bessel function. Has to be real valued.
    kind : int, optional
        The kind of the spherical Hankel function. Must be ``1`` or ``2``. The
        default is ``2``.
    derivative : bool, optional
        Flag to indicate if the derivative os the spherical Hankel function
        should be returned. The default is ``False``.

    Returns
    -------
    hn : double, ndarray
        Spherical bessel function. Array with dimensions [N x Z], where N is
        the number of elements in n and Z is the number of elements in z.

    References
    ----------
    .. [#] Rafaely, B. (2019). Fundamentals of spherical array processing
           (2nd ed.). Springer. https://doi.org/10.1007/978-3-319-99561-8

    Note
    ----
    This is based on the Hankel functions implemented in the scipy package.
    """

    if kind not in (1, 2):
        raise ValueError("The spherical hankel function can \
            only be of first or second kind.")

    n = np.asarray(n, dtype=int)
    z = np.asarray(z, dtype=float)

    ufunc = _spherical_hankel_derivative if derivative else _spherical_hankel

    if n.size > 1:
        hankel = np.zeros((n.size, z.size), dtype=complex)
        for idx, order in zip(count(), n):
            hankel[idx, :] = ufunc(order, z, kind)
    else:
        hankel = ufunc(n, z, kind)

    if z.ndim <= 1 or n.ndim <= 1:
        hankel = np.squeeze(hankel)

    return hankel


def _spherical_hankel(n, z, kind):
    if kind == 1:
        hankel = _spspecial.hankel1(n+0.5, z)
    elif kind == 2:
        hankel = _spspecial.hankel2(n+0.5, z)
    hankel = np.sqrt(np.pi/2/z) * hankel

    return hankel


def _spherical_hankel_derivative(n, z, kind):
    return _spherical_hankel(n - 1, z, kind) - (n + 1) / z * _spherical_hankel(
        n, z, kind)


def spherical_harmonic(n, m, theta, phi):
    """The spherical harmonics of order n and degree m.
    The spherical harmonic functions are fully normalized (N3D) and
    include the Condon-Shortley phase according to [#]_.

    n : unsigned int
        The spherical harmonic order
    m : int
        The spherical harmonic degree
    theta : ndarray, double
        The colatitude angle
    phi : ndarray, double
        The azimuth angle

    Returns
    -------
    spherical_harmonic : ndarray, double
        The complex valued spherical harmonic of order n and degree m

    Note
    ----
    This function wraps the spherical harmonic implementation from scipy.

    References
    ----------
    .. [#] B. Rafaely, "Fundamentals of Spherical Array Processing", (2015),
           Springer-Verlag
    """

    return _spspecial.sph_harm_y(n, m, theta, phi)


def spherical_harmonic_real(n, m, theta, phi):
    r"""Real valued spherical harmonic function of order n and degree m
    evaluated at the angles theta and phi.
    The spherical harmonic functions are fully normalized (N3D) and follow
    the AmbiX phase convention, which does not include the Condon-Shortley
    phase [#]_.

    .. math::

        Y_n^m(\theta, \phi) = \sqrt{\frac{2n+1}{4\pi}
        \frac{(n-|m|)!}{(n+|m|)!}} P_n^{|m|}(\cos \theta)
        \begin{cases}
            \displaystyle \cos(|m|\phi),  & \text{if $m \ge 0$} \newline
            \displaystyle \sin(|m|\phi) ,  & \text{if $m < 0$}
        \end{cases}

    References
    ----------
    .. [#]  C. Nachbar, F. Zotter, E. Deleflie, and A. Sontacchi, “Ambix - A
            Suggested Ambisonics Format (revised by F. Zotter),” International
            Symposium on Ambisonics and Spherical Acoustics,
            vol. 3, pp. 1–11, 2011.

    Parameters
    ----------
    n : unsigned int
        The spherical harmonic order
    m : int
        The spherical harmonic degree
    theta : ndarray, double
        The colatitude angle
    phi : ndarray, double
        The azimuth angle

    Returns
    -------
    spherical_harmonic : ndarray, double
        The real valued spherial harmonic of order n and degree m
    """
    # careful here, scipy uses phi as the elevation angle and
    # theta as the azimuth angle
    Y_nm_cplx = spherical_harmonic(n, m, theta, phi)

    if m == 0:
        Y_nm = np.real(Y_nm_cplx)
    elif m > 0:
        Y_nm = np.real(Y_nm_cplx) * np.sqrt(2)
    elif m < 0:
        Y_nm = np.imag(Y_nm_cplx) * np.sqrt(2) * float(-1)**(m+1)

    Y_nm *= float(-1)**(m)

    return Y_nm


def spherical_harmonic_derivative_phi(n, m, theta, phi):
    """Calculate the derivative of the spherical harmonics with respect to
    the azimuth angle phi.

    Parameters
    ----------
    n : int
        Spherical harmonic order
    m : int
        Spherical harmonic degree
    theta : double
        Elevation angle 0 < theta < pi
    phi : double
        Azimuth angle 0 < phi < 2*pi

    Returns
    -------
    sh_diff : complex double
        Spherical harmonic derivative
    """
    if m == 0 or n == 0:
        res = np.zeros(phi.shape, dtype=complex)
    else:
        res = spherical_harmonic(n, m, theta, phi) * 1j * m

    return res


def spherical_harmonic_gradient_phi(n, m, theta, phi):
    """Calculate the derivative of the spherical harmonics with respect to
    the azimuth angle phi divided by sin(theta).

    Parameters
    ----------
    n : int
        Spherical harmonic order
    m : int
        Spherical harmonic degree
    theta : double
        Elevation angle 0 < theta < pi
    phi : double
        Azimuth angle 0 < phi < 2*pi

    Returns
    -------
    sh_diff : complex double
        Spherical harmonic derivative
    """
    if m == 0:
        return np.zeros(theta.shape, dtype=complex)

    factor = np.sqrt((2*n+1)/(2*n-1))/2
    exp_phi = np.exp(1j*phi)
    first = np.sqrt((n+m)*(n+m-1)) * exp_phi * spherical_harmonic(
        n-1, m-1, theta, phi)
    second = np.sqrt((n-m) * (n-m-1)) / exp_phi * spherical_harmonic(
        n-1, m+1, theta, phi)
    Ynm_sin_theta = (-1) * factor * (first + second)

    return Ynm_sin_theta * 1j


def spherical_harmonic_derivative_theta(n, m, theta, phi):
    """Calculate the derivative of the spherical harmonics with respect to
    the elevation angle theta.

    Parameters
    ----------
    n : int
        Spherical harmonic order
    m : int
        Spherical harmonic degree
    theta : double
        Elevation angle 0 < theta < pi
    phi : double
        Azimuth angle 0 < phi < 2*pi

    Returns
    -------
    sh_diff : complex double
        Spherical harmonic derivative
    """
    if n == 0:
        return np.zeros(theta.shape, dtype=complex)

    exp_phi = np.exp(1j*phi)
    first = np.sqrt((n-m+1) * (n+m)) * exp_phi * spherical_harmonic(
        n, m-1, theta, phi)
    second = np.sqrt((n-m) * (n+m+1)) / exp_phi * spherical_harmonic(
        n, m+1, theta, phi)
    return (first-second)/2 * (-1)


def legendre_function(n, m, z, cs_phase=True):
    r"""
    Legendre function of order n and degree m with argument z.

    The Legendre function is defined as ([#]_ Eq. 1.30, 1.33, and 1.34)

    .. math::

        P_n^m(z) = (-1)^m(1-z^2)^{m/2}\frac{d^m}{dz^m}P_n(z),

    where the Condon-Shortley phase term :math:`(-1)^m` is dropped when
    ``cs_phase = False`` and the Legendre Polynomial :math:`P_n(z)`.

    Parameters
    ----------
    n : int
        The order
    m : int
        The degree
    z : ndarray, double
        The argument as an array
    cs_phase : bool, optional
        Whether to include the Condon-Shortley phase term :math:`(-1)^m`
        or not. Default is ``True``.

    Returns
    -------
    legendre : ndarray, double
        The Legendre function. This will return zeros if :math:`|m| > n`.

    References
    ----------
    .. [#] Rafaely, B. (2019). Fundamentals of spherical array processing
           (2nd ed.). Springer. https://doi.org/10.1007/978-3-319-99561-8

    Note
    ----
    This is a wrapper for the Legendre function implementation from scipy. The
    scipy implementation uses the Condon-Shortley phase. Therefore, the sign
    needs to be flipped here for uneven degrees when dropping the
    Condon-Shortley phase.
    """
    z = np.atleast_1d(z)

    # squeeze required because the legendre function introduced in scipy 1.15
    # returns a 2D array, whereas the previous function `lpmn` returned a 1D
    # array
    legendre = np.squeeze(_spspecial.assoc_legendre_p(n, m, z))

    # remove Condon-Shortley phase
    if not cs_phase and m % 2:
        legendre *= -1

    return legendre


def spherical_harmonic_normalization(n, m, norm='full'):
    r"""The normalization factor for real valued spherical harmonics.

    .. math::

        N_n^m = \sqrt{\frac{2n+1}{4\pi}\frac{(n-m)!}{(n+m)!}}

    Parameters
    ----------
    n : int
        The spherical harmonic order.
    m : int
        The spherical harmonic degree.
    norm : 'full', 'semi', optional
        Normalization to use. Can be either fully normalized on the sphere or
        semi-normalized. Default is ``'full'``.

    Returns
    -------
    norm : double
        The normalization factor.
    """
    if np.abs(m) > n:
        factor = 0.0
    elif norm == 'full':
        z = n+m+1
        factor = _spspecial.poch(z, -2*m)
        factor *= (2*n+1)/(4*np.pi)
        if int(m) != 0:
            factor *= 2
        factor = np.sqrt(factor)
    elif norm == 'semi':
        z = n+m+1
        factor = _spspecial.poch(z, -2*m)
        if int(m) != 0:
            factor *= 2
        factor = np.sqrt(factor)
    else:
        raise ValueError("Unknown normalization.")

    return factor


def spherical_harmonic_derivative_theta_real(n, m, theta, phi):
    r"""The derivative of the real valued spherical harmonics with respect
    to the elevation angle :math:`\theta`.

    Parameters
    ----------
    n : int
        The spherical harmonic order.
    m : int
        The spherical harmonic degree.
    theta : ndarray, double
        The elevation angle
    phi : ndarray, double
        The azimuth angle


    Returns
    -------
    derivative : ndarray, double
        The derivative

    Note
    ----
    This implementation does not include the Condon-Shortley phase term.
    """

    m_abs = np.abs(m)

    if n == 0:
        return np.zeros(theta.shape, dtype=float)

    first = (n+m_abs)*(n-m_abs+1) * \
        legendre_function(
            n,
            m_abs-1,
            np.cos(theta),
            cs_phase=False)
    second = legendre_function(
        n,
        m_abs+1,
        np.cos(theta),
        cs_phase=False)
    legendre_diff = 0.5*(first - second)

    N_nm = spherical_harmonic_normalization(n, m_abs)

    phi_term = np.sin(m_abs*phi) if m < 0 else np.cos(m_abs*phi)
    return N_nm * legendre_diff * phi_term


def spherical_harmonic_derivative_phi_real(n, m, theta, phi):
    r"""The derivative of the real valued spherical harmonics with respect
    to the azimuth angle :math:`\phi`.

    Parameters
    ----------
    n : int
        The spherical harmonic order.
    m : int
        The spherical harmonic degree.
    theta : ndarray, double
        The elevation angle
    phi : ndarray, double
        The azimuth angle


    Returns
    -------
    derivative : ndarray, double
        The derivative

    Note
    ----
    This implementation does not include the Condon-Shortley phase term.
    """
    m_abs = np.abs(m)
    if m == 0:
        return np.zeros(theta.shape, dtype=float)

    legendre = legendre_function(n, m_abs, np.cos(theta), cs_phase=False)
    N_nm = spherical_harmonic_normalization(n, m_abs)

    if m < 0:
        phi_term = np.cos(m_abs*phi) * m_abs
    else:
        phi_term = -np.sin(m_abs*phi) * m_abs

    return N_nm * legendre * phi_term


def spherical_harmonic_gradient_phi_real(n, m, theta, phi):
    r"""The gradient of the real valued spherical harmonics with respect
    to the azimuth angle :math:`\phi`.
    The singularities at the poles are avoided using the formulation proposed
    in [#]_.

    Parameters
    ----------
    n : int
        The spherical harmonic order.
    m : int
        The spherical harmonic degree.
    theta : ndarray, double
        The elevation angle
    phi : ndarray, double
        The azimuth angle


    Returns
    -------
    derivative : ndarray, double
        The derivative

    Note
    ----
    This implementation does not include the Condon-Shortley phase term.

    References
    ----------
    .. [#]  J. Du, C. Chen, V. Lesur, and L. Wang, “Non-singular spherical
            harmonic expressions of geomagnetic vector and gradient tensor
            fields in the local north-oriented reference frame,” Geoscientific
            Model Development, vol. 8, no. 7, pp. 1979–1990, Jul. 2015.
    """
    m_abs = np.abs(m)

    if m == 0:
        return np.zeros(theta.shape, dtype=float)

    first = (n+m_abs)*(n+m_abs-1) * \
        legendre_function(
            n-1,
            m_abs-1,
            np.cos(theta),
            cs_phase=False)
    second = legendre_function(
        n-1,
        m_abs+1,
        np.cos(theta),
        cs_phase=False)
    legendre_diff = 0.5*(first + second)
    N_nm = spherical_harmonic_normalization(n, m_abs)

    phi_term = np.cos(m_abs*phi) if m < 0 else -np.sin(m_abs*phi)
    return N_nm * legendre_diff * phi_term


def legendre_coefficients(order):
    """Calculate the coefficients of a Legendre polynomial of the
    specified order.

    Parameters
    ----------
    order : int
        The order of the polynomial

    Returns
    -------
    coefficients : ndarray, double
        The coefficients of the polynomial

    Note
    ----
    This uses :py:mod:`numpy.polynomial.legendre` to compute the coefficients.
    """
    leg = np.polynomial.legendre.Legendre.basis(order)
    return np.polynomial.legendre.leg2poly(leg.coef)


def chebyshev_coefficients(order):
    """Calculate the coefficients of a Chebyshev polynomial of the
    specified order.

    Parameters
    ----------
    order : int
        The order of the polynomial

    Returns
    -------
    coefficients : ndarray, double
        The coefficients of the polynomial

    Note
    ----
    This uses :py:mod:`numpy.polynomial.chebyshev` to compute the coefficients.
    """
    cheb = np.polynomial.chebyshev.Chebyshev.basis(order)
    return np.polynomial.chebyshev.cheb2poly(cheb.coef)
