"""Various utility functions.

.. include:: math-definitions.rst

"""

import collections
import numpy as np
from numpy.core.umath_tests import inner1d
from scipy.special import spherical_jn, spherical_yn
from . import default


def rotation_matrix(n1, n2):
    """Compute rotation matrix for rotation from *n1* to *n2*.

    Parameters
    ----------
    n1, n2 : (3,) array_like
        Two vectors.  They don't have to be normalized.

    Returns
    -------
    (3, 3) `numpy.ndarray`
        Rotation matrix.

    """
    n1 = normalize_vector(n1)
    n2 = normalize_vector(n2)
    I = np.identity(3)
    if np.all(n1 == n2):
        return I  # no rotation
    elif np.all(n1 == -n2):
        return -I  # flip
    # TODO: check for *very close to* parallel vectors

    # Algorithm from http://math.stackexchange.com/a/476311
    v = v0, v1, v2 = np.cross(n1, n2)
    s = np.linalg.norm(v)  # sine
    c = np.inner(n1, n2)  # cosine
    vx = [[0, -v2, v1],
          [v2, 0, -v0],
          [-v1, v0, 0]]  # skew-symmetric cross-product matrix
    return I + vx + np.dot(vx, vx) * (1 - c) / s**2


def wavenumber(omega, c=None):
    """Compute the wavenumber for a given radial frequency."""
    if c is None:
        c = default.c
    return omega / c


def direction_vector(alpha, beta=np.pi/2):
    """Compute normal vector from azimuth, colatitude."""
    return sph2cart(alpha, beta, 1)


def sph2cart(alpha, beta, r):
    r"""Spherical to cartesian coordinate transform.

    .. math::

        x = r \cos \alpha \sin \beta \\
        y = r \sin \alpha \sin \beta \\
        z = r \cos \beta

    with :math:`\alpha \in [0, 2\pi), \beta \in [0, \pi], r \geq 0`

    Parameters
    ----------
    alpha : float or array_like
            Azimuth angle in radiants
    beta : float or array_like
            Colatitude angle in radiants (with 0 denoting North pole)
    r : float or array_like
            Radius

    Returns
    -------
    x : float or `numpy.ndarray`
        x-component of Cartesian coordinates
    y : float or `numpy.ndarray`
        y-component of Cartesian coordinates
    z : float or `numpy.ndarray`
        z-component of Cartesian coordinates

    """
    x = r * np.cos(alpha) * np.sin(beta)
    y = r * np.sin(alpha) * np.sin(beta)
    z = r * np.cos(beta)
    return x, y, z


def cart2sph(x, y, z):
    r"""Cartesian to spherical coordinate transform.

    .. math::

        \alpha = \arctan \left( \frac{y}{x} \right) \\
        \beta = \arccos \left( \frac{z}{r} \right) \\
        r = \sqrt{x^2 + y^2 + z^2}

    with :math:`\alpha \in [-pi, pi], \beta \in [0, \pi], r \geq 0`

    Parameters
    ----------
    x : float or array_like
        x-component of Cartesian coordinates
    y : float or array_like
        y-component of Cartesian coordinates
    z : float or array_like
        z-component of Cartesian coordinates

    Returns
    -------
    alpha : float or `numpy.ndarray`
            Azimuth angle in radiants
    beta : float or `numpy.ndarray`
            Colatitude angle in radiants (with 0 denoting North pole)
    r : float or `numpy.ndarray`
            Radius

    """
    r = np.sqrt(x**2 + y**2 + z**2)
    alpha = np.arctan2(y, x)
    beta = np.arccos(z / r)
    return alpha, beta, r


def asarray_1d(a, **kwargs):
    """Squeeze the input and check if the result is one-dimensional.

    Returns *a* converted to a `numpy.ndarray` and stripped of
    all singleton dimensions.  Scalars are "upgraded" to 1D arrays.
    The result must have exactly one dimension.
    If not, an error is raised.

    """
    result = np.squeeze(np.asarray(a, **kwargs))
    if result.ndim == 0:
        result = result.reshape((1,))
    elif result.ndim > 1:
        raise ValueError("array must be one-dimensional")
    return result


def asarray_of_rows(a, **kwargs):
    """Convert to 2D array, turn column vector into row vector.

    Returns *a* converted to a `numpy.ndarray` and stripped of
    all singleton dimensions.  If the result has exactly one dimension,
    it is re-shaped into a 2D row vector.

    """
    result = np.squeeze(np.asarray(a, **kwargs))
    if result.ndim == 1:
        result = result.reshape(1, -1)
    return result


def as_xyz_components(components, **kwargs):
    r"""Convert *components* to `XyzComponents` of `numpy.ndarray`\s.

    The *components* are first converted to NumPy arrays (using
    :func:`numpy.asarray`) which are then assembled into an
    `XyzComponents` object.

    Parameters
    ----------
    components : triple or pair of array_like
        The values to be used as X, Y and Z arrays.  Z is optional.
    **kwargs
        All further arguments are forwarded to :func:`numpy.asarray`,
        which is applied to the elements of *components*.

    """
    return XyzComponents([np.asarray(c, **kwargs) for c in components])


def as_delayed_signal(arg, **kwargs):
    """Make sure that the given argument can be used as a signal.

    Parameters
    ----------
    arg : sequence of 1 array_like followed by 1 or 2 scalars
        The first element is converted to a NumPy array, the second
        element is used as the sampling rate (in Hertz) and the optional
        third element is used as the starting time of the signal (in
        seconds).  Default starting time is 0.
    **kwargs
        All keyword arguments are forwarded to :func:`numpy.asarray`.

    Returns
    -------
    `DelayedSignal`
        A named tuple consisting of a `numpy.ndarray` containing the
        audio data, followed by the sampling rate (in Hertz) and the
        starting time (in seconds) of the signal.

    Examples
    --------
    Typically, this is used together with tuple unpacking to assign the
    audio data, the sampling rate and the starting time to separate
    variables:

    >>> import sfs
    >>> sig = [1], 44100
    >>> data, fs, signal_offset = sfs.util.as_delayed_signal(sig)
    >>> data
    array([1])
    >>> fs
    44100
    >>> signal_offset
    0

    """
    try:
        data, samplerate, *time = arg
        time, = time or [0]
    except (IndexError, TypeError, ValueError):
        pass
    else:
        valid_arguments = (not np.isscalar(data) and
                           np.isscalar(samplerate) and
                           np.isscalar(time))
        if valid_arguments:
            data = np.asarray(data, **kwargs)
            return DelayedSignal(data, samplerate, time)
    raise TypeError('expected audio data, samplerate, optional start time')


def strict_arange(start, stop, step=1, *, endpoint=False, dtype=None,
                  **kwargs):
    """Like :func:`numpy.arange`, but compensating numeric errors.

    Unlike :func:`numpy.arange`, but similar to :func:`numpy.linspace`,
    providing ``endpoint=True`` includes both endpoints.

    Parameters
    ----------
    start, stop, step, dtype
        See :func:`numpy.arange`.
    endpoint
        See :func:`numpy.linspace`.

        .. note:: With ``endpoint=True``, the difference between *start*
           and *end* value must be an integer multiple of the
           corresponding *spacing* value!
    **kwargs
        All further arguments are forwarded to :func:`numpy.isclose`.

    Returns
    -------
    `numpy.ndarray`
        Array of evenly spaced values.  See :func:`numpy.arange`.

    """
    remainder = (stop - start) % step
    if np.any(np.isclose(remainder, (0.0, step), **kwargs)):
        if endpoint:
            stop += step * 0.5
        else:
            stop -= step * 0.5
    elif endpoint:
        raise ValueError("Invalid stop value for endpoint=True")
    return np.arange(start, stop, step, dtype)


def xyz_grid(x, y, z, *, spacing, endpoint=True, **kwargs):
    """Create a grid with given range and spacing.

    Parameters
    ----------
    x, y, z : float or pair of float
        Inclusive range of the respective coordinate or a single value
        if only a slice along this dimension is needed.
    spacing : float or triple of float
        Grid spacing.  If a single value is specified, it is used for
        all dimensions, if multiple values are given, one value is used
        per dimension.  If a dimension (*x*, *y* or *z*) has only a
        single value, the corresponding spacing is ignored.
    endpoint : bool, optional
        If ``True`` (the default), the endpoint of each range is
        included in the grid.  Use ``False`` to get a result similar to
        :func:`numpy.arange`.  See `strict_arange()`.
    **kwargs
        All further arguments are forwarded to `strict_arange()`.

    Returns
    -------
    `XyzComponents`
        A grid that can be used for sound field calculations.

    See Also
    --------
    strict_arange, numpy.meshgrid

    """
    if np.isscalar(spacing):
        spacing = [spacing] * 3
    ranges = []
    scalars = []
    for i, coord in enumerate([x, y, z]):
        if np.isscalar(coord):
            scalars.append((i, coord))
        else:
            start, stop = coord
            ranges.append(strict_arange(start, stop, spacing[i],
                                        endpoint=endpoint, **kwargs))
    grid = np.meshgrid(*ranges, sparse=True, copy=False)
    for i, s in scalars:
        grid.insert(i, s)
    return XyzComponents(grid)


def normalize(p, grid, xnorm):
    """Normalize sound field wrt position *xnorm*."""
    return p / np.abs(probe(p, grid, xnorm))


def probe(p, grid, x):
    """Determine the value at position *x* in the sound field *p*."""
    grid = as_xyz_components(grid)
    x = asarray_1d(x)
    r = np.linalg.norm(grid - x)
    idx = np.unravel_index(r.argmin(), r.shape)
    return p[idx]


def broadcast_zip(*args):
    """Broadcast arguments to the same shape and then use :func:`zip`."""
    return zip(*np.broadcast_arrays(*args))


def normalize_vector(x):
    """Normalize a 1D vector."""
    x = asarray_1d(x)
    return x / np.linalg.norm(x)


def db(x, *, power=False):