import collections
import numbers
import numpy as np

from ._input_validation import _nonneg_int_or_fail

from ._special_ufuncs import (legendre_p, assoc_legendre_p,
                              sph_legendre_p, sph_harm_y)
from ._gufuncs import (legendre_p_all, assoc_legendre_p_all,
                       sph_legendre_p_all, sph_harm_y_all)

__all__ = [
    "assoc_legendre_p",
    "assoc_legendre_p_all",
    "legendre_p",
    "legendre_p_all",
    "sph_harm_y",
    "sph_harm_y_all",
    "sph_legendre_p",
    "sph_legendre_p_all",
]


class MultiUFunc:
    def __init__(self, ufunc_or_ufuncs, doc=None, *,
                 force_complex_output=False, **default_kwargs):
        if not isinstance(ufunc_or_ufuncs, np.ufunc):
            if isinstance(ufunc_or_ufuncs, collections.abc.Mapping):
                ufuncs_iter = ufunc_or_ufuncs.values()
            elif isinstance(ufunc_or_ufuncs, collections.abc.Iterable):
                ufuncs_iter = ufunc_or_ufuncs
            else:
                raise ValueError("ufunc_or_ufuncs should be a ufunc or a"
                                 " ufunc collection")

            # Perform input validation to ensure all ufuncs in ufuncs are
            # actually ufuncs and all take the same input types.
            seen_input_types = set()
            for ufunc in ufuncs_iter:
                if not isinstance(ufunc, np.ufunc):
                    raise ValueError("All ufuncs must have type `numpy.ufunc`."
                                     f" Received {ufunc_or_ufuncs}")
                seen_input_types.add(frozenset(x.split("->")[0] for x in ufunc.types))
            if len(seen_input_types) > 1:
                raise ValueError("All ufuncs must take the same input types.")

        self._ufunc_or_ufuncs = ufunc_or_ufuncs
        self.__doc = doc
        self.__force_complex_output = force_complex_output
        self._default_kwargs = default_kwargs
        self._resolve_out_shapes = None
        self._finalize_out = None
        self._key = None
        self._ufunc_default_args = lambda *args, **kwargs: ()
        self._ufunc_default_kwargs = lambda *args, **kwargs: {}

    @property
    def __doc__(self):
        return self.__doc

    def _override_key(self, func):
        """Set `key` method by decorating a function.
        """
        self._key = func

    def _override_ufunc_default_args(self, func):
        self._ufunc_default_args = func

    def _override_ufunc_default_kwargs(self, func):
        self._ufunc_default_kwargs = func

    def _override_resolve_out_shapes(self, func):
        """Set `resolve_out_shapes` method by decorating a function."""
        if func.__doc__ is None:
            func.__doc__ = \
                """Resolve to output shapes based on relevant inputs."""
        func.__name__ = "resolve_out_shapes"
        self._resolve_out_shapes = func

    def _override_finalize_out(self, func):
        self._finalize_out = func

    def _resolve_ufunc(self, **kwargs):
        """Resolve to a ufunc based on keyword arguments."""

        if isinstance(self._ufunc_or_ufuncs, np.ufunc):
            return self._ufunc_or_ufuncs

        ufunc_key = self._key(**kwargs)
        return self._ufunc_or_ufuncs[ufunc_key]

    def __call__(self, *args, **kwargs):
        kwargs = self._default_kwargs | kwargs

        args += self._ufunc_default_args(**kwargs)

        ufunc = self._resolve_ufunc(**kwargs)

        # array arguments to be passed to the ufunc
        ufunc_args = [np.asarray(arg) for arg in args[-ufunc.nin:]]

        ufunc_kwargs = self._ufunc_default_kwargs(**kwargs)

        if (self._resolve_out_shapes is not None):
            ufunc_arg_shapes = tuple(np.shape(ufunc_arg) for ufunc_arg in ufunc_args)
            ufunc_out_shapes = self._resolve_out_shapes(*args[:-ufunc.nin],
                                                        *ufunc_arg_shapes, ufunc.nout,
                                                        **kwargs)

            ufunc_arg_dtypes = tuple(ufunc_arg.dtype if hasattr(ufunc_arg, 'dtype')
                                     else np.dtype(type(ufunc_arg))
                                     for ufunc_arg in ufunc_args)

            if hasattr(ufunc, 'resolve_dtypes'):
                ufunc_dtypes = ufunc_arg_dtypes + ufunc.nout * (None,)
                ufunc_dtypes = ufunc.resolve_dtypes(ufunc_dtypes)
                ufunc_out_dtypes = ufunc_dtypes[-ufunc.nout:]
            else:
                ufunc_out_dtype = np.result_type(*ufunc_arg_dtypes)
                if (not np.issubdtype(ufunc_out_dtype, np.inexact)):
                    ufunc_out_dtype = np.float64

                ufunc_out_dtypes = ufunc.nout * (ufunc_out_dtype,)

            if self.__force_complex_output:
                ufunc_out_dtypes = tuple(np.result_type(1j, ufunc_out_dtype)
                                         for ufunc_out_dtype in ufunc_out_dtypes)

            out = tuple(np.empty(ufunc_out_shape, dtype=ufunc_out_dtype)
                        for ufunc_out_shape, ufunc_out_dtype
                        in zip(ufunc_out_shapes, ufunc_out_dtypes))

            ufunc_kwargs['out'] = out

        out = ufunc(*ufunc_args, **ufunc_kwargs)
        if (self._finalize_out is not None):
            out = self._finalize_out(out)

        return out


sph_legendre_p = MultiUFunc(
    sph_legendre_p,
    r"""sph_legendre_p(n, m, theta, *, diff_n=0)

    Spherical Legendre polynomial of the first kind.

    Parameters
    ----------
    n : ArrayLike[int]
        Degree of the spherical Legendre polynomial. Must have ``n >= 0``.
    m : ArrayLike[int]
        Order of the spherical Legendre polynomial.
    theta : ArrayLike[float]
        Input value.
    diff_n : Optional[int]
        A non-negative integer. Compute and return all derivatives up
        to order ``diff_n``. Default is 0.

    Returns
    -------
    p : ndarray or tuple[ndarray]
        Spherical Legendre polynomial with ``diff_n`` derivatives.

    Notes
    -----
    The spherical counterpart of an (unnormalized) associated Legendre polynomial has
    the additional factor

    .. math::

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

    It is the same as the spherical harmonic :math:`Y_{n}^{m}(\theta, \phi)`
    with :math:`\phi = 0`.
    """, diff_n=0
)


@sph_legendre_p._override_key
def _(diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return diff_n


@sph_legendre_p._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


sph_legendre_p_all = MultiUFunc(
    sph_legendre_p_all,
    """sph_legendre_p_all(n, m, theta, *, diff_n=0)

    All spherical Legendre polynomials of the first kind up to the
    specified degree ``n`` and order ``m``.

    Output shape is ``(n + 1, 2 * m + 1, ...)``. The entry at ``(j, i)``
    corresponds to degree ``j`` and order ``i`` for all  ``0 <= j <= n``
    and ``-m <= i <= m``.

    See Also
    --------
    sph_legendre_p
    """, diff_n=0
)


@sph_legendre_p_all._override_key
def _(diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return diff_n


@sph_legendre_p_all._override_ufunc_default_kwargs
def _(diff_n):
    return {'axes': [()] + [(0, 1, -1)]}


@sph_legendre_p_all._override_resolve_out_shapes
def _(n, m, theta_shape, nout, diff_n):
    if not isinstance(n, numbers.Integral) or (n < 0):
        raise ValueError("n must be a non-negative integer.")

    return ((n + 1, 2 * abs(m) + 1) + theta_shape + (diff_n + 1,),)


@sph_legendre_p_all._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


assoc_legendre_p = MultiUFunc(
    assoc_legendre_p,
    r"""assoc_legendre_p(n, m, z, *, branch_cut=2, norm=False, diff_n=0)

    Associated Legendre polynomial of the first kind.

    Parameters
    ----------
    n : ArrayLike[int]
        Degree of the associated Legendre polynomial. Must have ``n >= 0``.
    m : ArrayLike[int]
        order of the associated Legendre polynomial.
    z : ArrayLike[float | complex]
        Input value.
    branch_cut : Optional[ArrayLike[int]]
        Selects branch cut. Must be 2 (default) or 3.
        2: cut on the real axis ``|z| > 1``
        3: cut on the real axis ``-1 < z < 1``
    norm : Optional[bool]
        If ``True``, compute the normalized associated Legendre polynomial.
        Default is ``False``.
    diff_n : Optional[int]
        A non-negative integer. Compute and return all derivatives up
        to order ``diff_n``. Default is 0.

    Returns
    -------
    p : ndarray or tuple[ndarray]
        Associated Legendre polynomial with ``diff_n`` derivatives.

    Notes
    -----
    The normalized counterpart of an (unnormalized) associated Legendre
    polynomial has the additional factor

    .. math::

        \sqrt{\frac{(2 n + 1) (n - m)!}{2 (n + m)!}}
    """, branch_cut=2, norm=False, diff_n=0
)


@assoc_legendre_p._override_key
def _(branch_cut, norm, diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return norm, diff_n


@assoc_legendre_p._override_ufunc_default_args
def _(branch_cut, norm, diff_n):
    return branch_cut,


@assoc_legendre_p._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


assoc_legendre_p_all = MultiUFunc(
    assoc_legendre_p_all,
    """assoc_legendre_p_all(n, m, z, *, branch_cut=2, norm=False, diff_n=0)

    All associated Legendre polynomials of the first kind up to the
    specified degree ``n`` and order ``m``.

    Output shape is ``(n + 1, 2 * m + 1, ...)``. The entry at ``(j, i)``
    corresponds to degree ``j`` and order ``i`` for all  ``0 <= j <= n``
    and ``-m <= i <= m``.

    See Also
    --------
    assoc_legendre_p
    """, branch_cut=2, norm=False, diff_n=0
)


@assoc_legendre_p_all._override_key
def _(branch_cut, norm, diff_n):
    if not ((isinstance(diff_n, numbers.Integral))
            and diff_n >= 0):
        raise ValueError(
            f"diff_n must be a non-negative integer, received: {diff_n}."
        )
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return norm, diff_n


@assoc_legendre_p_all._override_ufunc_default_args
def _(branch_cut, norm, diff_n):
    return branch_cut,


@assoc_legendre_p_all._override_ufunc_default_kwargs
def _(branch_cut, norm, diff_n):
    return {'axes': [(), ()] + [(0, 1, -1)]}


@assoc_legendre_p_all._override_resolve_out_shapes
def _(n, m, z_shape, branch_cut_shape, nout, **kwargs):
    diff_n = kwargs['diff_n']

    if not isinstance(n, numbers.Integral) or (n < 0):
        raise ValueError("n must be a non-negative integer.")
    if not isinstance(m, numbers.Integral) or (m < 0):
        raise ValueError("m must be a non-negative integer.")

    return ((n + 1, 2 * abs(m) + 1) +
        np.broadcast_shapes(z_shape, branch_cut_shape) + (diff_n + 1,),)


@assoc_legendre_p_all._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


legendre_p = MultiUFunc(
    legendre_p,
    """legendre_p(n, z, *, diff_n=0)

    Legendre polynomial of the first kind.

    Parameters
    ----------
    n : ArrayLike[int]
        Degree of the Legendre polynomial. Must have ``n >= 0``.
    z : ArrayLike[float]
        Input value.
    diff_n : Optional[int]
        A non-negative integer. Compute and return all derivatives up
        to order ``diff_n``. Default is 0.

    Returns
    -------
    p : ndarray or tuple[ndarray]
        Legendre polynomial with ``diff_n`` derivatives.

    See Also
    --------
    legendre

    References
    ----------
    .. [1] Zhang, Shanjie and Jin, Jianming. "Computation of Special
           Functions", John Wiley and Sons, 1996.
           https://people.sc.fsu.edu/~jburkardt/f77_src/special_functions/special_functions.html
    """, diff_n=0
)


@legendre_p._override_key
def _(diff_n):
    if (not isinstance(diff_n, numbers.Integral)) or (diff_n < 0):
        raise ValueError(
            f"diff_n must be a non-negative integer, received: {diff_n}."
        )
    if not 0 <= diff_n <= 2:
        raise NotImplementedError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return diff_n


@legendre_p._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


legendre_p_all = MultiUFunc(
    legendre_p_all,
    """legendre_p_all(n, z, *, diff_n=0)

    All Legendre polynomials of the first kind up to the
    specified degree ``n``.

    Output shape is ``(n + 1, ...)``. The entry at ``j``
    corresponds to degree ``j`` for all  ``0 <= j <= n``.

    See Also
    --------
    legendre_p
    """, diff_n=0
)


@legendre_p_all._override_key
def _(diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return diff_n


@legendre_p_all._override_ufunc_default_kwargs
def _(diff_n):
    return {'axes': [(), (0, -1)]}


@legendre_p_all._override_resolve_out_shapes
def _(n, z_shape, nout, diff_n):
    n = _nonneg_int_or_fail(n, 'n', strict=False)

    return nout * ((n + 1,) + z_shape + (diff_n + 1,),)


@legendre_p_all._override_finalize_out
def _(out):
    return np.moveaxis(out, -1, 0)


sph_harm_y = MultiUFunc(
    sph_harm_y,
    r"""sph_harm_y(n, m, theta, phi, *, diff_n=0)

    Spherical harmonics. They are defined as

    .. math::

        Y_n^m(\theta,\phi) = \sqrt{\frac{2 n + 1}{4 \pi} \frac{(n - m)!}{(n + m)!}}
            P_n^m(\cos(\theta)) e^{i m \phi}

    where :math:`P_n^m` are the (unnormalized) associated Legendre polynomials.

    Parameters
    ----------
    n : ArrayLike[int]
        Degree of the harmonic. Must have ``n >= 0``. This is
        often denoted by ``l`` (lower case L) in descriptions of
        spherical harmonics.
    m : ArrayLike[int]
        Order of the harmonic.
    theta : ArrayLike[float]
        Polar (colatitudinal) coordinate; must be in ``[0, pi]``.
    phi : ArrayLike[float]
        Azimuthal (longitudinal) coordinate; must be in ``[0, 2*pi]``.
    diff_n : Optional[int]
        A non-negative integer. Compute and return all derivatives up
        to order ``diff_n``. Default is 0.

    Returns
    -------
    y : ndarray[complex] or tuple[ndarray[complex]]
       Spherical harmonics with ``diff_n`` derivatives.

    Notes
    -----
    There are different conventions for the meanings of the input
    arguments ``theta`` and ``phi``. In SciPy ``theta`` is the
    polar angle and ``phi`` is the azimuthal angle. It is common to
    see the opposite convention, that is, ``theta`` as the azimuthal angle
    and ``phi`` as the polar angle.

    Note that SciPy's spherical harmonics include the Condon-Shortley
    phase [2]_ because it is part of `sph_legendre_p`.

    With SciPy's conventions, the first several spherical harmonics
    are

    .. math::

        Y_0^0(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{1}{\pi}} \\
        Y_1^{-1}(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{3}{2\pi}}
                                    e^{-i\phi} \sin(\theta) \\
        Y_1^0(\theta, \phi) &= \frac{1}{2} \sqrt{\frac{3}{\pi}}
                                 \cos(\theta) \\
        Y_1^1(\theta, \phi) &= -\frac{1}{2} \sqrt{\frac{3}{2\pi}}
                                 e^{i\phi} \sin(\theta).

    References
    ----------
    .. [1] Digital Library of Mathematical Functions, 14.30.
           https://dlmf.nist.gov/14.30
    .. [2] https://en.wikipedia.org/wiki/Spherical_harmonics#Condon.E2.80.93Shortley_phase
    """, force_complex_output=True, diff_n=0
)


@sph_harm_y._override_key
def _(diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 0, 1, and 2,"
            f" received: {diff_n}."
        )
    return diff_n


@sph_harm_y._override_finalize_out
def _(out):
    if (out.shape[-1] == 1):
        return out[..., 0, 0]

    if (out.shape[-1] == 2):
        return out[..., 0, 0], out[..., [1, 0], [0, 1]]

    if (out.shape[-1] == 3):
        return (out[..., 0, 0], out[..., [1, 0], [0, 1]],
            out[..., [[2, 1], [1, 0]], [[0, 1], [1, 2]]])


sph_harm_y_all = MultiUFunc(
    sph_harm_y_all,
    """sph_harm_y_all(n, m, theta, phi, *, diff_n=0)

    All spherical harmonics up to the specified degree ``n`` and order ``m``.

    Output shape is ``(n + 1, 2 * m + 1, ...)``. The entry at ``(j, i)``
    corresponds to degree ``j`` and order ``i`` for all  ``0 <= j <= n``
    and ``-m <= i <= m``.

    See Also
    --------
    sph_harm_y
    """, force_complex_output=True, diff_n=0
)


@sph_harm_y_all._override_key
def _(diff_n):
    diff_n = _nonneg_int_or_fail(diff_n, "diff_n", strict=False)
    if not 0 <= diff_n <= 2:
        raise ValueError(
            "diff_n is currently only implemented for orders 2,"
            f" received: {diff_n}."
        )
    return diff_n


@sph_harm_y_all._override_ufunc_default_kwargs
def _(diff_n):
    return {'axes': [(), ()] + [(0, 1, -2, -1)]}


@sph_harm_y_all._override_resolve_out_shapes
def _(n, m, theta_shape, phi_shape, nout, **kwargs):
    diff_n = kwargs['diff_n']

    if not isinstance(n, numbers.Integral) or (n < 0):
        raise ValueError("n must be a non-negative integer.")

    return ((n + 1, 2 * abs(m) + 1) + np.broadcast_shapes(theta_shape, phi_shape) +
        (diff_n + 1, diff_n + 1),)


@sph_harm_y_all._override_finalize_out
def _(out):
    if (out.shape[-1] == 1):
        return out[..., 0, 0]

    if (out.shape[-1] == 2):
        return out[..., 0, 0], out[..., [1, 0], [0, 1]]

    if (out.shape[-1] == 3):
        return (out[..., 0, 0], out[..., [1, 0], [0, 1]],
            out[..., [[2, 1], [1, 0]], [[0, 1], [1, 2]]])