from __future__ import division

import numpy as np
from .coordinates import RegularCoords, SeparatedCoords, UnstructuredCoords
from .field import Field
from .cartesian_grid import CartesianGrid
from .grid import Grid

def make_uniform_grid(dims, extent, center=0, has_center=False):
    '''Create a uniformly-spaced :class:`Grid` of a certain shape and size.

    Parameters
    ----------
    dims : scalar or ndarray
        The number of points in each dimension. If this is a scalar, it will
        be multiplexed over all dimensions.
    extent : scalar or ndarray
        The total extent of the grid in each dimension.
    center : scalar or ndarray
        The center point. The grid will by symmetric around this point.
    has_center : boolean
        Does the grid has to have the center as one of its points. If this is
        False, this does not mean that the grid will not have the center.

    Returns
    -------
    Grid
        A :class:`Grid` with :class:`RegularCoords`.
    '''
    num_dims = max(np.array([dims]).shape[-1], np.array([extent]).shape[-1], np.array([center]).shape[-1])

    dims = (np.ones(num_dims) * dims).astype('int')
    extent = (np.ones(num_dims) * extent).astype('float')
    center = (np.ones(num_dims) * center).astype('float')

    delta = extent / dims
    zero = -extent / 2 + center + delta / 2

    if has_center:
        zero -= delta / 2 * (1 - np.mod(dims, 2))

    return CartesianGrid(RegularCoords(delta, dims, zero))

def make_pupil_grid(dims, diameter=1):
    '''Makes a new :class:`Grid`, meant for descretisation of a pupil-plane wavefront.

    This grid is symmetric around the origin, and therefore has no point exactly on
    the origin for an even number of pixels.

    Parameters
    ----------
    dims : ndarray or integer
        The number of pixels per dimension. If this is an integer, this number
        of pixels is used for all dimensions.
    diameter : ndarray or scalar
        The diameter of the grid in each dimension. If this is a scalar, this diameter
        is used for all dimensions.

    Returns
    -------
    Grid
        A :class:`CartesianGrid` with :class:`RegularCoords`.
    '''
    diameter = (np.ones(2) * diameter).astype('float')
    return make_uniform_grid(dims, diameter)

def make_focal_grid_from_pupil_grid(pupil_grid, q=1, num_airy=None, focal_length=1, wavelength=1):
    '''Make a grid for a focal plane from a pupil grid.

    Calculate the focal grid corresponding to the pupil grid, using an FFT grid as a guide. The
    resulting grid focal will always contain the origin (0, 0) point. The extent of the pupil
    grid will be used as the diameter of the pupil. If the pupil is undersized on the pupil
    grid, the resulting focal grid needs to be rescaled manually.

    .. note::
        In almost all cases, it is preferable to use :func:`make_focal_grid`. This
        function allows you to directly set the diameter, and doesn't require the user to pass
        the pupil grid as an argument. `make_focal_grid_from_pupil_grid()` retains old functionality
        and serves as a backwards compatibility function, due to its ubiquitous usage in HCIPy code.

    Parameters
    ----------
    pupil_grid : Grid
        The pupil grid for which the focal grid needs to be calculated. The extent of this Grid
        is used as the diameter of the pupil.
    q : scalar or array_like
        The number of pixels per resolution element (= lambda f / D).
    num_airy : scalar or array_like
        The spatial extent of the grid in radius in resolution elements (= lambda f / D).
    focal_length : scalar
        The focal length used for calculating the spatial resolution at the focal plane.
    wavelength : scalar
        The reference wavelength used for calculating the spatial resolution at the focal plane.

    Returns
    -------
    Grid
        A Grid describing the sampling for a focal plane.
    '''
    from ..fourier import make_fft_grid

    f_lambda = focal_length * wavelength
    if num_airy is None:
        fov = 1
    else:
        fov = (num_airy * np.ones(pupil_grid.ndim, dtype='float')) / (pupil_grid.shape / 2)

    if np.max(fov) > 1:
        import warnings
        warnings.warn('Focal grid is larger than the maximum allowed angle (fov=%.03f). You may see wrapping when doing propagations.' % np.max(fov), stacklevel=2)

    uv = make_fft_grid(pupil_grid, q, fov)
    focal_grid = uv.scaled(f_lambda / (2 * np.pi))

    return focal_grid

def make_focal_grid(q, num_airy, spatial_resolution=None, pupil_diameter=None, focal_length=None, f_number=None, reference_wavelength=None):
    r'''Make a grid for a focal plane.

    This grid will be a CartesianGrid with RegularCoords, and supports different resolutions, samplings,
    or extent in x and y. If `spatial_resolution` is 1, then the grid will be returned in normalized
    (ie. homogenized) coordinates. Otherwise, it will be in physical units. The spatial resolution is
    defined by:

    .. math:: \Delta x = \lambda f / D = \lambda F

    where :math:`\lambda` is the wavelength, :math:`f` is the effective focal length before the focal plane,
    :math:`D` is the diameter of the pupil, and :math:`F` is the F-number of the incoming light beam.
    You can supply either the spatial resolution or a set of focal length, reference wavelength and pupil diameter,
    or a set of F-number and reference wavelength. If none are supplied, a spatial resolution of 1 will be
    assumed, meaning normalized units.

    The grid will always contain the origin (0, 0) point.

    Parameters
    ----------
    q : scalar or array_like
        The number of pixels per resolution element (= lambda f / D).
    num_airy : scalar or array_like
        The spatial extent of the grid in radius in resolution elements (= lambda f / D).
    spatial_resolution : scalar    or array_like
        The physical size of a resolution element (= lambda f / D). It this is not given,
        the spatial resolution will be calculated from the given `focal_length`,
        `reference_wavelength` and `pupil_diameter`.
    pupil_diameter : scalar or array_like
        The diameter of the pupil. If it is an array, this indicates the diameter in x and y.
    focal_length : scalar
        The focal length used for calculating the spatial resolution at the focal plane.
    f_number : scalar or array_like
        The F number, also known as focal ratio, at the focal plane. If this is given, it overrides
        the given pupil diameter and focal length.
    reference_wavelength : scalar
        The reference wavelength used for calculating the spatial resolution at the focal plane.

    Returns
    -------
    Grid
        A Grid describing the sampling for a focal plane.

    Raises
    ------
    ValueError
        If both no spatial resolution and no complete set of (focal length, reference wavelength and pupil diameter) was supplied.
    '''
    if isinstance(q, Grid):
        raise ValueError('The function signature was changed as of HCIPy 0.3.0. Please use the new signature (prefered), or use make_focal_grid_from_pupil_grid() if you want to retain old behaviour.')

    if spatial_resolution is None:
        if f_number is None:
            if pupil_diameter is None or focal_length is None:
                if reference_wavelength is None:
                    spatial_resolution = 1
                else:
                    raise ValueError('You only supplied a reference wavelength and forgot to supply either an f_number or a (pupil_diameter, focal_length).')
            else:
                f_number = focal_length / pupil_diameter

        if spatial_resolution is None:
            if reference_wavelength is None:
                raise ValueError('You supplied an f_number or (pupil_diameter, focal_length), and forgot to supply a reference wavelength.')
            else:
                spatial_resolution = f_number * reference_wavelength

    delta = spatial_resolution / q * np.ones(2)
    dims = (2 * num_airy * q * np.ones(2)).astype('int')
    zero = delta * (-dims / 2 + np.mod(dims, 2) * 0.5)

    return CartesianGrid(RegularCoords(delta, dims, zero))

def make_hexagonal_grid(circum_diameter, n_rings, pointy_top=False, center=None):
    '''Make a regular hexagonal grid.

    Parameters
    ----------
    circum_diameter : scalar
        The circum diameter of the hexagons in the grid.
    n_rings : integer
        The number of rings in the grid.
    pointy_top : boolean
        If the hexagons contained in the grid.
    center : ndarray
        The center of the grid in cartesian coordinates.

    Returns
    -------
    Grid
        A :class:`CartesianGrid` with `UnstructuredCoords`, indicating the
        center of the hexagons.
    '''
    if center is None:
        center = np.zeros(2)

    apothem = circum_diameter * np.sqrt(3) / 4

    q = [0]
    r = [0]

    for n in range(1, n_rings + 1):
        # top
        q += list(range(n, 0, -1))
        r += list(range(0, n))
        # right top
        q += list(range(0, -n, -1))
        r += [n] * n
        # right bottom
        q += [-n] * n
        r += list(range(n, 0, -1))
        # bottom
        q += list(range(-n, 0))
        r += list(range(0, -n, -1))
        # left bottom
        q += list(range(0, n))
        r += [-n] * n
        # left top
        q += [n] * n
        r += list(range(-n, 0))

    x = (-np.array(q) + np.array(r)) * circum_diameter / 2 + center[0]
    y = (np.array(q) + np.array(r)) * apothem * 2 + center[1]

    weight = 2 * apothem**2 * np.sqrt(3)

    if pointy_top:
        return CartesianGrid(UnstructuredCoords((x, y)), weight)
    else:
        return CartesianGrid(UnstructuredCoords((y, x)), weight)

def make_chebyshev_grid(dims, minimum=None, maximum=None):
    if minimum is None:
        minimum = -1

    if maximum is None:
        maximum = 1

    dims = np.array(dims)
    minimum = np.ones(len(dims)) * minimum
    maximum = np.ones(len(dims)) * maximum

    middles = (minimum + maximum) / 2
    intervals = (maximum - minimum) / 2

    sep_coords = []
    for dim, middle, interval in zip(dims, middles, intervals):
        c = np.cos(np.pi * (2 * np.arange(dim) + 1) / (2.0 * dim))
        c = middle + interval * c
        sep_coords.append(c)

    return CartesianGrid(SeparatedCoords(sep_coords))

def make_supersampled_grid(grid, oversampling):
    '''Make a new grid that oversamples by a factor `oversampling`.

    .. note ::
        The Grid `grid` must be a grid with separable coordinates.

    Parameters
    ----------
    grid : Grid
        The grid that we want to oversample.
    oversampling : integer or scalar or ndarray
        The factor by which to oversample. If this is a scalar, it will be rounded to
        the nearest integer. If this is an array, a different oversampling factor will
        be used for each dimension.

    Returns
    -------
    Grid
        The oversampled grid.
    '''
    oversampling = (np.round(oversampling)).astype('int')

    if grid.is_regular:
        delta_new = grid.delta / oversampling
        zero_new = grid.zero - grid.delta / 2 + delta_new / 2
        dims_new = grid.dims * oversampling

        return grid.__class__(RegularCoords(delta_new, dims_new, zero_new))
    elif grid.is_separated:
        raise NotImplementedError()

    raise ValueError('Cannot create a supersampled grid from a non-separated grid.')

def make_subsampled_grid(grid, undersampling):
    '''Make a new grid that undersamples by a factor `undersampling`.

    .. note ::
        The dimensions of the `grid` must be divisible by `undersampling`.

    Parameters
    ----------
    grid : Grid
        The grid that we want to oversample.
    undersampling : integer or scalar or ndarray
        The factor by which to undersample. If this is a scalar, it will be rounded to
        the nearest integer. If this is an array, a different undersampling factor will
        be used for each dimension.

    Returns
    -------
    Grid
        The undersampled grid.
    '''
    undersampling = (np.round(undersampling)).astype('int')

    if grid.is_regular:
        delta_new = grid.delta * undersampling
        zero_new = grid.zero - grid.delta / 2 + delta_new / 2
        dims_new = grid.dims // undersampling

        return grid.__class__(RegularCoords(delta_new, dims_new, zero_new))
    elif grid.is_separated:
        raise NotImplementedError()

    raise ValueError("Cannot create a subsampled grid from a non-separated grid.")

def subsample_field(field, subsampling, new_grid=None, statistic='mean'):
    '''Average the field over subsampling pixels in each dimension.

    .. note ::
        The dimensions of the grid of `field` must be divisible by `subsampling`.

    Parameters
    ----------
    field : Field
        The field to subsample. The grid of this field must have the right
        dimensions to be able to be subsampled.
    subsampling : integer or scalar or ndarray
        The subsampling factor. If this is a scalar, it will be rounded to the
        nearest integer. If this is an array, the subsampling factor will be
        different for each dimension.
    new_grid : Grid
        If this grid is given, no new grid will be calculated and this grid will
        be used instead. This saves on calculation time if your new grid is already
        known beforehand.
    statistic : string
        The statistic to compute (default is 'mean').
        The following statistics are available:
        * 'mean': compute the mean of values for points within each superpixel.
        * 'sum': compute the sum of values for points within each superpixel. This is identical to a weighted histogram.
        * 'min': compute the minimum of values for points within each superpixel.
        * 'max': compute the maximum of values for points within each superpixel.
        * 'median': compute the median of values for points within each superpixel.
        * 'nanmedian': compute the median of values for points within each superpixel while ignoring NaN values.

    Returns
    -------
    Field
        The subsampled field.
    '''
    subsampling = (np.round(subsampling)).astype('int')

    if new_grid is None:
        new_grid = make_subsampled_grid(field.grid, subsampling)

    reshape = []
    axes = []
    for i, s in enumerate(new_grid.shape):
        reshape.extend([s, subsampling])
        axes.append(2 * i + 1)

    if field.tensor_order > 0:
        reshape = list(field.tensor_shape) + reshape
        axes = np.array(axes) + field.tensor_order
        new_shape = list(field.tensor_shape) + [-1]
    else:
        new_shape = [-1]

    available_statistics = {
        'mean': np.mean,
        'max': np.max,
        'min': np.min,
        'sum': np.sum,
        'median': np.median,
        'nanmedian': np.nanmedian
    }

    if statistic not in available_statistics:
        raise ValueError('This statistic is not recognized.')

    if field.grid.is_regular:
        # All weights will be the same, so the array can be combined without taking the weights into account.
        return Field(available_statistics[statistic](field.reshape(tuple(reshape)), axis=tuple(axes)).reshape(tuple(new_shape)), new_grid)
    else:
        # Some weights will be different so calculate weighted mean instead.
        if statistic in ['min', 'max', 'sum', 'median', 'nanmedian']:
            f = available_statistics[statistic](field.reshape(tuple(reshape)), axis=tuple(axes))
            return Field(f.reshape(tuple(new_shape)), new_grid)
        else:
            # Statistic is mean
            weights = field.grid.weights
            w = weights.reshape(tuple(reshape)).sum(axis=tuple(axes))
            f = np.sum((field * weights).reshape(tuple(reshape)), axis=tuple(axes))
            return Field((f / w).reshape(tuple(new_shape)), new_grid)

def evaluate_supersampled(field_generator, grid, oversampling, statistic='mean', make_sparse=True):
    '''Evaluate a Field generator on `grid`, with an oversampling.

    Parameters
    ----------
    field_generator : Field generator or list of Field generators
        The field generator to evaluate. If this is a list of Field generators,
        each Field generator will be evaluated and stored in a ModeBasis.
    grid : Grid
        The grid on which to evaluate `field_generator`.
    oversampling : integer or scalar or ndarray
        The factor by which to oversample. If this is a scalar, it will be rounded to
        the nearest integer. If this is an array, a different oversampling factor will
        be used for each dimension.
    statistic : string
        The statistic to compute (default is 'mean').
        The following statistics are available:
        * 'mean': compute the mean of values for points within each superpixel.
        * 'sum': compute the sum of values for points within each superpixel. This is identical to a weighted histogram.
        * 'min': compute the minimum of values for points within each superpixel.
        * 'max': compute the maximum of values for point within each superpixel.
    make_sparse : boolean
        If the resulting ModeBasis needs to be sparsified. This is ignored if
        only a single Field generator is provided.

    Returns
    -------
    Field or ModeBasis
        The evaluated field or mode basis.
    '''
    import scipy.sparse
    from ..mode_basis import ModeBasis

    if isinstance(field_generator, (list, tuple)):
        modes = []

        for fg in field_generator:
            field = evaluate_supersampled(fg, grid, oversampling, statistic)

            if make_sparse:
                field = scipy.sparse.csr_matrix(field)
                field.eliminate_zeros()

            modes.append(field)

        return ModeBasis(modes, grid)

    oversampling = (np.round(oversampling) * np.ones(grid.ndim)).astype('int')

    if grid.is_separated:
        # Use dithered sub grids to get the final supersampled field.
        deltas = []
        for i in range(grid.ndim):
            x = grid.coords.separated_coords[i]
            d = (x[2:] - x[:-2]) / 2.
            d = np.concatenate(([x[1] - x[0]], d, [x[-1] - x[-2]]))
            deltas.append(d)

        dithers = make_uniform_grid(oversampling, 1)

        separated_coords = grid.separated_coords

        if statistic in ['mean', 'sum']:
            field = 0
        else:
            field = None

        for dither in dithers.points:
            dithered_separated_coords = [c + d * delta for c, d, delta in zip(separated_coords, dither, deltas)]
            dithered_grid = grid.__class__(SeparatedCoords(dithered_separated_coords))

            if statistic in ['mean', 'sum']:
                field += field_generator(dithered_grid)
            else:
                if field is None:
                    field = field_generator(dithered_grid)
                elif statistic == 'min':
                    field = np.minimum(field, field_generator(dithered_grid))
                else:
                    field = np.maximum(field, field_generator(dithered_grid))

        if statistic == 'mean':
            field /= len(dithers)

        field.grid = grid
        return field
    else:
        # Cannot use sub grids, so fall back to evaluation of generator on the full
        # supersampled grid.
        supersampled_grid = make_supersampled_grid(grid, oversampling)
        field = field_generator(supersampled_grid)
        return subsample_field(field, oversampling, grid, statistic)

def make_uniform_vector_field(field, jones_vector):
    '''Make an uniform vector field from a scalar field and a jones vector.

    Parameters
    ----------
    field : Field
        An input scalar field that is expanded to a vector field
    jones_vector : array_like
        The output vector at every pixel

    Returns
    -------
    Field
        The expanded vector field
    '''
    if field.is_scalar_field():
        return Field([ei * field for ei in jones_vector], field.grid)

def make_uniform_vector_field_generator(field_generator, jones_vector):
    '''Make an uniform vector field generator from a scalar field generator and a jones vector.

    Parameters
    ----------
    field_generator : Field generator
        The field generator to evaluate.
    jones_vector : array_like
        The output vector at every grid coordinate

    Returns
    -------
    Field generator
        This function can be evaluated on a grid to get a Field.
    '''

    def func(grid):
        scalar_field = field_generator(grid)
        return Field([ei * scalar_field for ei in jones_vector], grid)

    return func
