"""
The 2-sphere, S^2
"""
import numpy as np
from numpy.polynomial.legendre import leggauss


def change_coordinates(coords, p_from='C', p_to='S'):
    """
    Change Spherical to Cartesian coordinates and vice versa, for points x in S^2.

    :param coords:
    :param p_from:
    :param p_to:
    :return:
    """
    if p_from == p_to:
        return coords
    elif p_from == 'S' and p_to == 'C':

        theta = coords[..., 0]
        phi = coords[..., 1]
        r = 1.

        out = np.empty(theta.shape + (3,))

        ct = np.cos(theta)
        cp = np.cos(phi)
        st = np.sin(theta)
        sp = np.sin(phi)
        out[..., 0] = r * st * cp  # x
        out[..., 1] = r * st * sp  # y
        out[..., 2] = r * ct       # z
        return out

    elif p_from == 'C' and p_to == 'S':

        x = coords[..., 0]
        y = coords[..., 1]
        z = coords[..., 2]

        out = np.empty(x.shape + (2,))
        out[..., 0] = np.arccos(z)         # theta
        out[..., 1] = np.arctan2(y, x)     # phi
        return out

    else:
        raise ValueError('Unknown conversion:' + str(p_from) + ' to ' + str(p_to))


def meshgrid(b, convention='Driscoll-Healy'):
    """
    Create a coordinate grid for the 2-sphere.
    There are various ways to setup a grid on the sphere.

    if convention == 'Driscoll-Heally', we follow the convention from [4], which is also used in [5].
    The Driscoll-Healy convention is:
    theta_j = pi j / (2 b)
    phi_k = pi k / b
    where j, k in {0, ..., 2 b - 1}

    if convention == 'SOFT', we follow the convention from [1]
    The SOFT convention is:
    theta_j = pi (2 j + 1) / (4 b)
    phi_k = pi  k / b
    (in the SOFT documentation, substitute theta = beta, phi = alpha)

    if convention == 'Clenshaw-Curtis', we use the Clenshaw-Curtis grid, as defined in [2].

    if convention == 'Gauss-Legendre', we use the Gauss-Legendre grid, as defined in [2].

    if convention == 'HEALPix', we use the HEALPix grid.

    [1] SOFT: SO(3) Fourier Transforms
    Kostelec, Peter J & Rockmore, Daniel N.

    [2] Fast evaluation of quadrature formulae on the sphere
    Jens Keiner, Daniel Potts

    [3] A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature
    Zydrunas Gimbutas Shravan Veerapaneni

    [4] Computing Fourier transforms and convolutions on the 2-sphere
    Driscoll, JR & Healy, DM

    [5] Engineering Applications of Noncommutative Harmonic Analysis
    Chrikjian, G.S. & Kyatkin, A.B.

    :param b:
    :return:
    """
    return np.meshgrid(*linspace(b, convention))


def linspace(b, convention='Driscoll-Healy'):
    if convention == 'Driscoll-Healy':
        theta = np.arange(2 * b) * np.pi / (2. * b)
        phi = np.arange(2 * b) * np.pi / b
    elif convention == 'SOFT':
        theta = np.pi * (2 * np.arange(2 * b) + 1) / (4. * b)
        phi = np.arange(2 * b) * np.pi / b
    elif convention == 'Clenshaw-Curtis':
        # theta = np.arange(2 * b + 1) * np.pi / (2 * b)
        # =phi = np.arange(2 * b + 2) * np.pi / (b + 1)
        # Must use np.linspace to prevent numerical errors that cause theta > pi
        theta = np.linspace(0, np.pi, 2 * b + 1)
        phi = np.linspace(0, 2 * np.pi, 2 * b + 2, endpoint=False)
    elif convention == 'Gauss-Legendre':
        # For details, see:
        # "A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature"
        # Zydrunas Gimbutas, Shravan Veerapaneni
        x, _ = leggauss(b + 1)
        theta = np.arccos(x)
        phi = np.arange(2 * b + 2) * np.pi / (b + 1)
    elif convention == 'HEALPix':
        # TODO: implement this here so that we don't need the dependency on healpy / healpix_compat
        from healpix_compat import healpy_sphere_meshgrid
        return healpy_sphere_meshgrid(b)
    elif convention == 'equidistribution':
        raise NotImplementedError('Not implemented yet; see Fast evaluation of quadrature formulae on the sphere.')
    else:
        raise ValueError('Unknown convention:' + convention)
    return theta, phi


def quadrature_weights(b, convention='Gauss-Legendre'):
    """
    Compute quadrature weights for a given grid-type.
    The function S2.meshgrid generates the points that correspond to the weights generated by this function.

    if convention == 'Gauss-Legendre':
    The quadrature formula is exact for polynomials up to degree M less than or equal to 2b + 1,
    so that we can compute exact Fourier coefficients for f a polynomial of degree at most b.

    if convention == 'Clenshaw-Curtis':
    The quadrature formula is exact for polynomials up to degree M less than or equal to 2b,
    so that we can compute exact Fourier coefficients for f a polynomial of degree at most b.

    :param b: the grid resolution. See S2.meshgrid
    :param convention: 'Gauss-Legendre' or 'Clenshaw-Curtis'
    :return:
    """
    if convention == 'Clenshaw-Curtis':
        # Use the fast fft based method to compute these weights
        # see "Fast evaluation of quadrature formulae on the sphere"
        w = _clenshaw_curtis_weights(n=2 * b)
        W = np.empty((2 * b + 2, 2 * b + 1))
        W[:] = w[None, :]
    elif convention == 'Gauss-Legendre':
        # We found this formula in:
        # "A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature"
        # eq. 10
        _, w = leggauss(b + 1)
        W = w[None, :] * (2 * np.pi / (2 * b + 2) * np.ones(2 * b + 2)[:, None])
    else:
        raise ValueError('Unknown convention:' + str(convention))

    return W


def integrate(f, normalize=True):
    """
    Integrate a function f : S^2 -> R over the sphere S^2, using the invariant integration measure
    mu((theta, phi)) = sin(theta) dtheta dphi
    i.e. this returns
    int_S^2 f(x) dmu(x) = int_0^2pi int_0^pi f(theta, phi) sin(theta) dtheta dphi

    :param f: a function of two scalar variables theta in [0, pi] and phi in [0, 2pi], returning a scalar.
    :param normalize: if True, use the normalized Haar measure mu so that int_S^1 dmu(x) = 1.
     If False, use the Haar measure for the unit sphere, so that int_S^2 dmu(x) = 4 pi (the surface area of S^2).
    :return: integral of f over S^2
    """
    from scipy.integrate import quad

    f2 = lambda phi: quad(lambda theta: f(theta, phi) * np.sin(theta),
                          a=0,
                          b=np.pi)[0]
    integral = quad(f2, 0, 2 * np.pi)[0]

    if normalize:
        return integral / (4 * np.pi)
    else:
        return integral


def plot_sphere_func(f, grid='Clenshaw-Curtis', theta=None, phi=None, colormap='jet', fignum=0):

    # Note: all grids except Clenshaw-Curtis have holes at the poles

    import matplotlib
    matplotlib.use('WxAgg')
    matplotlib.interactive(True)
    from mayavi import mlab

    if grid == 'Driscoll-Healy':
        b = f.shape[0] / 2
    elif grid == 'Clenshaw-Curtis':
        b = (f.shape[0] - 2) / 2
    elif grid == 'SOFT':
        b = f.shape[0] / 2
    elif grid == 'Gauss-Legendre':
        b = (f.shape[0] - 2) / 2

    if theta is None or phi is None:
        theta, phi = meshgrid(b=b, convention=grid)

    phi = np.r_[phi, phi[0, :][None, :]]
    theta = np.r_[theta, theta[0, :][None, :]]
    f = np.r_[f, f[0, :][None, :]]

    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)

    mlab.figure(fignum, bgcolor=(1, 1, 1), fgcolor=(0, 0, 0), size=(600, 400))
    mlab.clf()
    mlab.mesh(x, y, z, scalars=f, colormap=colormap)

    # mlab.view(90, 70, 6.2, (-1.3, -2.9, 0.25))
    mlab.show()


def _clenshaw_curtis_weights(n):
    """
    Computes the Clenshaw-Curtis quadrature using a fast FFT method.

    For details, see
    Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules
    Jorg Waldvogel, 2005

    :param n:
    :return:
    """
    from scipy.fftpack import ifft

    N = np.arange(start=1, stop=n, step=2)[:, None]
    l = N.size
    m = n - l

    v0 = np.vstack([2. / N / (N-2), 1. / N[-1]] + [0] * m)
    v2 = -v0[:-1] - v0[:0:-1]

    g0 = -np.ones((n, 1))
    g0[l] += n
    g0[m] += n
    g = g0 / (n ** 2 - 1 + n % 2)

    wcc = ifft((v2 + g).flatten()).real
    wcc = np.hstack([wcc, wcc[0]])

    return wcc * np.pi / (n / 2 + 1)
