import numpy as np
from math import sqrt
from scipy.special import eval_genlaguerre

def gaussian_laguerre(p, l, mode_field_diameter=1, grid=None):
    r'''Creates a Gaussian-Laguerre mode.

    This function evaluates a (p,l) order Gaussian-Laguerre mode on a grid.
    The definition of the modes are the following,

    .. math::
        \exp{\left(-\frac{r^2}{w_0^2}\right)} L_p^{|l|}\left(\frac{2r^2}{w_0^2} \right) \left(\sqrt{2}\frac{r}{w_0}\right)

    Here :math:`w_0` is the mode_field_radius, which is :math:`\mathrm{MFD}/2`.
    And :math:`L_p^{|l|}` are the generalized Laguerre Polynomials.
    All modes are numerical normalized to have a total power of 1.

    More details on generalized Laguerre Polynomials can be found on:
    http://mathworld.wolfram.com/AssociatedLaguerrePolynomial.html

    Parameters
    ----------
    p : int
        The radial order.
    l : int
        The azimuthal order.
    mode_field_diameter : scalar
        The mode field diameter of the mode.
    grid : Grid
        The grid on which to evaluate the mode.

    Returns
    -------
    Field
        The evaluated mode.
    '''
    from ..field import Field

    if grid.is_separated and grid.is_('polar'):
        R, Theta = grid.separated_coords
    else:
        R, Theta = grid.as_('polar').coords

    # Easy access
    r = 2 * R / mode_field_diameter
    r2 = r**2

    # Compute the mode
    lg = (r * sqrt(2))**abs(l) * np.exp(-r2) * np.exp(-1j * l * Theta) * eval_genlaguerre(p, abs(l), 2 * r2)

    # Numerically normalize the mode
    lg /= np.sqrt(np.sum(np.abs(lg)**2 * grid.weights))

    return Field(lg, grid)

# High level functions
def make_gaussian_laguerre_basis(grid, pmax, lmax, mode_field_diameter, pmin=0):
    '''Creates a Gaussian-Laguerre mode basis.

    This function evaluates Gaussian-Laguerre modes. For each radial order
    within [pmin, pmax] it will calculate the azimuthal order [-lmax, lmax] inclusive.
    This function returns a ModeBasis made out of these Gaussian-Laguerre modes.

    Parameters
    ----------
    grid : Grid
        The grid on which to evaluate the Gaussian-Laguerre mode.
    pmax : int
        The maximum radial order of the modes.
    lmax : int
        The maximum azimuthal order.
    mode_field_diameter : scalar
        The mode field diameter of the Gaussian-Laguerre mode.
    pmin : int
        The minimal radial order.

    Returns
    -------
    ModeBasis
        The Gaussian-Laguerre modes.
    '''
    from .mode_basis import ModeBasis

    modes = [gaussian_laguerre(pi, li, mode_field_diameter, grid) for li in range(-lmax, lmax + 1) for pi in range(pmin, pmax)]
    return ModeBasis(modes, grid)
