import numpy as np
import copy
import warnings
import xxhash

class Grid(object):
    '''A set of points on some coordinate system.

    Parameters
    ----------
    coords : Coords
        The actual definition of the coordinate values.
    weights : array_like or None
        The interval size, area, volume or hypervolume of each point, depending on the number of dimensions.
        If this is None (default), the weights will be attempted to be calculated on the fly when needed.

    Attributes
    ----------
    coords
        The coordinate values for each dimension.
    '''

    _coordinate_system = 'none'
    _coordinate_system_transformations = {}
    _coordinate_systems = {}

    def __init__(self, coords, weights=None):
        self.coords = coords
        self.weights = weights

    def copy(self):
        '''Create a copy.
        '''
        return copy.deepcopy(self)

    @classmethod
    def from_dict(cls, tree):
        '''Make a Grid from a dictionary, previously created by `to_dict()`.

        Parameters
        ----------
        tree : dictionary
            The dictionary from which to make a new Grid object.

        Returns
        -------
        Grid
            The created object.

        Raises
        ------
        ValueError
            If the dictionary is not formatted correctly.
        '''
        from .coordinates import Coords

        coords = Coords.from_dict(tree['coords'])
        grid_class = Grid._coordinate_systems[tree['coordinate_system']]

        if 'weights' in tree:
            weights = copy.deepcopy(tree['weights'])
        else:
            weights = None

        return grid_class(coords, weights)

    def to_dict(self):
        '''Convert the object to a dictionary for serialization.

        Returns
        -------
        dictionary
            The created dictionary.
        '''
        tree = {
            'coordinate_system': self._coordinate_system,
            'coords': self.coords.to_dict(),
            'weights': self._weights
        }

        return tree

    def subset(self, criterium):
        '''Construct a subset of the current sampling, based on `criterium`.

        Parameters
        ----------
        criterium : function or array_like
            The criterium used to select points. A function will be evaluated for every point.
            Otherwise, this must be a boolean array of integer array, used for slicing the points.

        Returns
        -------
        Grid
            A new grid with UnstructuredCoords that includes only the points for which the criterium
            was true.
        '''
        from .coordinates import UnstructuredCoords

        if hasattr(criterium, '__call__'):
            indices = criterium(self) != 0
        else:
            indices = criterium

        new_coords = [c[indices] for c in self.coords]

        if np.isscalar(self.weights):
            new_weights = self.weights
        else:
            new_weights = self.weights[indices]

        return self.__class__(UnstructuredCoords(new_coords), new_weights)

    @property
    def ndim(self):
        '''The number of dimensions.
        '''
        return len(self.coords)

    @property
    def size(self):
        '''The number of points in this grid.
        '''
        return self.coords.size

    def __len__(self):
        '''The number of points in this grid.
        '''
        return self.size

    @property
    def dims(self):
        '''The number of elements in each dimension for a separated grid.

        Raises
        ------
        ValueError
            If the grid is not separated.
        '''
        if not self.is_separated:
            raise ValueError('A non-separated grid does not have dims.')
        return self.coords.dims

    @property
    def shape(self):
        '''The shape of a reshaped ``numpy.ndarray`` using this grid.

        Raises
        ------
        ValueError
            If the grid is not separated.
        '''
        if not self.is_separated:
            raise ValueError('A non-separated grid does not have a shape.')
        return self.coords.shape

    @property
    def delta(self):
        '''The spacing between points in regularly-spaced grid.

        Raises
        ------
        ValueError
            If the grid is not regular.
        '''
        if not self.is_regular:
            raise ValueError('A non-regular grid does not have a delta.')
        return self.coords.delta

    @property
    def zero(self):
        '''The zero point of a regularly-spaced grid.

        Raises
        ------
        ValueError
            If the grid is not regular.
        '''
        if not self.is_regular:
            raise ValueError('A non-regular grid does not have a zero.')
        return self.coords.zero

    @property
    def separated_coords(self):
        '''A list of coordinates for each dimension in a separated grid.

        Raises
        ------
        ValueError
            If the grid is not separated.
        '''
        if not self.is_separated:
            raise ValueError('A non-separated grid does not have separated coordinates.')
        return self.coords.separated_coords

    @property
    def regular_coords(self):
        '''The tuple (delta, dims, zero) for a regularly-spaced grid.

        Raises
        ------
        ValueError
            If the grid is not regular.
        '''
        if not self.is_regular:
            raise ValueError('A non-regular grid does not have regular coordinates.')
        return self.coords.regular_coords

    @property
    def weights(self):
        '''The interval size, area, volume or hypervolume of each point, depending on the number of dimensions.

        The weights are attempted to be calculated on the fly if not set. If this fails, a warning is emitted and all points will be given an equal weight of one.
        '''
        if self._weights is None:
            self._weights = self.__class__._get_automatic_weights(self.coords)

            if self._weights is None:
                self._weights = 1
                warnings.warn('No automatic weights could be calculated for this grid.', stacklevel=2)

        return self._weights

    @weights.setter
    def weights(self, weights):
        self._weights = weights

    @property
    def points(self):
        '''A list of points of this grid.

        This can be used for easier iteration::

            for p in grid.points:
                print(p)
        '''
        return np.array(self.coords).T

    @property
    def is_separated(self):
        '''True if the grid is separated, False otherwise.
        '''
        return self.coords.is_separated

    @property
    def is_regular(self):
        '''True if the grid is regularly-spaced, False otherwise.
        '''
        return self.coords.is_regular

    @property
    def is_unstructured(self):
        '''True if the grid is unstructured, False otherwise.
        '''
        return self.coords.is_unstructured

    def is_(self, system):
        '''Check if the coordinate system is `system`.

        Parameters
        ----------
        system : str
            The name of the coordinate system to check for.

        Returns
        -------
        bool
            If the coordinate system of the grid is equal to `system`.
        '''
        return system == self._coordinate_system

    def as_(self, system):
        '''Convert the grid to the new coordinate system `system`.

        If the grid is already in the right coordinate system, this function doesn't do anything.

        Parameters
        ----------
        system : str
            The name of the coordinate system to check for.

        Returns
        -------
        Grid
            A new :class:`Grid` in the required coordinate system.

        Raises
        ------
        ValueError
            If the conversion to the coordinate system `system` isn't known.
        '''
        if self.is_(system):
            return self
        else:
            return Grid._coordinate_system_transformations[self._coordinate_system][system](self)

    @staticmethod
    def _add_coordinate_system_transformation(source, dest, func):
        if source in Grid._coordinate_system_transformations:
            Grid._coordinate_system_transformations[source][dest] = func
        else:
            Grid._coordinate_system_transformations[source] = {dest: func}

    @staticmethod
    def _add_coordinate_system(name, grid_class):
        Grid._coordinate_systems[name] = grid_class

    def __getitem__(self, i):
        '''The `i`-th point in this grid.
        '''
        return self.points[i]

    def scale(self, scale):
        '''Scale the grid in-place.

        Parameters
        ----------
        scale : array_like
            The factor with which to scale the grid.

        Returns
        -------
        Grid
            Itself to allow for chaining these transformations.
        '''
        raise NotImplementedError()

    def scaled(self, scale):
        '''A scaled copy of this grid.

        Parameters
        ----------
        scale : array_like
            The factor with which to scale the grid.

        Returns
        -------
        Grid
            The scaled grid.
        '''
        grid = self.copy()
        grid.scale(scale)
        return grid

    def shift(self, shift):
        '''Shift the grid in-place.

        Parameters
        ----------
        shift : array_like
            The amount with which to shift the grid.

        Returns
        -------
        Grid
            Itself to allow for chaining these transformations.
        '''
        raise NotImplementedError()

    def shifted(self, shift):
        '''A shifted copy of this grid.

        Parameters
        ----------
        shift : array_like
            The amount with which to shift the grid.

        Returns
        -------
        Grid
            The scaled grid.
        '''
        grid = self.copy()
        grid.shift(shift)
        return grid

    def rotate(self, angle, axis=None):
        '''Rotate the grid in-place.

        Parameters
        ----------
        angle : scalar
            The angle in radians.
        axis : ndarray or None
            The axis of rotation. For two-dimensional grids, it is ignored. For
            three-dimensional grids it is required.

        Returns
        -------
        Grid
            Itself to allow for chaining these transformations.
        '''
        raise NotImplementedError()

    def rotated(self, angle, axis=None):
        '''A rotated copy of this grid.

        Parameters
        ----------
        angle : scalar
            The angle in radians.
        axis : ndarray or None
            The axis of rotation. For two-dimensional grids, it is ignored. For
            three-dimensional grids it is required.

        Returns
        -------
        Grid
            The rotated grid.
        '''
        grid = self.copy()
        grid.rotate(angle, axis)
        return grid

    def reverse(self):
        '''Reverse the order of the points in-place.

        Returns
        -------
        Grid
            Itself to allow for chaining these transformations.
        '''
        self.coords.reverse()
        return self

    def reversed(self):
        '''Make a copy of the grid with the order of the points reversed.

        Returns
        -------
        Grid
            The reversed grid.
        '''
        grid = self.copy()
        grid.reverse()
        return grid

    @staticmethod
    def _get_automatic_weights(coords):
        raise NotImplementedError()

    def __str__(self):
        return str(self.__class__.__name__) + '(' + str(self.coords.__class__.__name__) + ')'

    def __hash__(self):
        h = xxhash.xxh64()
        h.update(self._coordinate_system)

        if self.is_regular:
            h.update(self.delta)
            h.update(self.dims)
            h.update(self.zero)
        elif self.is_separated:
            for s in self.separated_coords:
                h.update(s)
        else:
            for s in self.coords:
                h.update(s)

        return h.intdigest()

    def __eq__(self, other):
        '''Check equality of two grids.

        .. note::
            This will return False if the two grids have different coordinate systems or coordinate types,
            even if the underlying points are identical

        Parameters
        ----------
        other : object
            The object to which to compare to.

        Returns
        -------
        boolean
            Whether the two objects are identical.
        '''
        if not isinstance(other, Grid):
            return False

        if self._coordinate_system != other._coordinate_system:
            return False

        return self.coords == other.coords

    def closest_to(self, p):
        '''Get the index of the point closest to point `p`.

        Point `p` is assumed to have the same coordinate system as the grid itself.

        Parameters
        ----------
        p : array_like
            The point at which to search for.

        Returns
        -------
        int
            The index of the closest point.
        '''
        rel_points = self.points - np.array(p) * np.ones(self.ndim)
        return np.argmin(np.sum(rel_points**2, axis=-1))

    def zeros(self, tensor_shape=None, dtype=None):
        '''Create a field of zeros from this `Grid`.

        Parameters
        ----------
        tensor_shape : array_like or None
            The shape of the tensors in the to be created field. If this is None,
            a scalar field will be created.
        dtype : data-type
            The numpy data-type with which to create the field.

        Returns
        -------
        Field
            A zeros field.
        '''
        from .field import Field

        if tensor_shape is None:
            shape = [self.size]
        else:
            shape = np.concatenate((tensor_shape, [self.size]))

        return Field(np.zeros(shape, dtype), self)

    def ones(self, tensor_shape=None, dtype=None):
        '''Create a field of ones from this `Grid`.

        Parameters
        ----------
        tensor_shape : array_like or None
            The shape of the tensors in the to be created field. If this is None,
            a scalar field will be created.
        dtype : data-type
            The numpy data-type with which to create the field.

        Returns
        -------
        Field
            A ones field.
        '''
        from .field import Field

        if tensor_shape is None:
            shape = [self.size]
        else:
            shape = np.concatenate((tensor_shape, [self.size]))

        return Field(np.ones(shape, dtype=dtype), self)

    def empty(self, tensor_shape=None, dtype=None):
        '''Create an empty Field from this `Grid`.

        Parameters
        ----------
        tensor_shape : array_like or None
            The shape of the tensors in the to be created field. If this is None,
            a scalar field will be created.
        dtype : data-type
            The numpy data-type with which to create the field.

        Returns
        -------
        Field
            A empty field.
        '''
        from .field import Field

        if tensor_shape is None:
            shape = [self.size]
        else:
            shape = np.concatenate((tensor_shape, [self.size]))

        return Field(np.empty(shape, dtype=dtype), self)
