import numpy as np
import copy
import warnings

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

	Parameters
	----------
	coords : CoordsBase
		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 = {}
	
	def __init__(self, coords, weights=None):
		self.coords = coords
		self.weights = weights
	
	def copy(self):
		'''Create a copy.
		'''
		return copy.deepcopy(self)
	
	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.
		'''
		from .coordinates import UnstructuredCoords
		
		if hasattr(criterium, '__call__'):
			indices = criterium(self) != 0
		else:
			indices = criterium
		new_coords = [c[indices] for c in self.coords]
		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
	
	@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)
		
		if np.isscalar(self._weights):
			return np.ones(self.size) * self._weights
		else:
			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
	
	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}
	
	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 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__) + '(' + str(self.coords.__class__) + ')'

	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))