import numpy as np

from .grid import Grid
from .coordinates import UnstructuredCoords

class PolarGrid(Grid):
	'''A grid representing a two-dimensional Polar coordinate system.
	'''
	_coordinate_system = 'polar'

	@property
	def r(self):
		'''The radial coordinate (dimension 0).
		'''
		return self.coords[0]

	@property
	def theta(self):
		''' The angular coordinate (dimension 1).
		'''
		return self.coords[1]

	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.
		'''
		self.coords *= np.array([scale, 1])
		self.weights *= np.abs(scale)**(self.ndim)
		return self

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

		.. caution::
		   All structure in the coordinates will be destroyed.

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

		Returns
		-------
		Grid
			Itself to allow for chaining these transformations.
		'''
		new_grid = PolarGrid(UnstructuredCoords(self.coords)).as_('cartesian')
		self.coords = new_grid.shifted(shift).as_('polar').coords

		return self

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

		.. caution::
		   The returned grid is a Cartesian grid.

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

		Returns
		-------
		Grid
			The scaled grid.
		'''
		grid = self.as_('cartesian')
		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 this (polar) grid, it is ignored.

		Returns
		-------
		Grid
			Itself to allow for chaining these transformations.
		'''
		self.theta += angle
		return self

	@staticmethod
	def _get_automatic_weights(coords):
		return None

def _cartesian_to_polar(self):
	from .coordinates import UnstructuredCoords

	x = self.x
	y = self.y
	r = np.hypot(x, y)
	theta = np.arctan2(y, x)
	return PolarGrid(UnstructuredCoords([r, theta]))

def _polar_to_cartesian(self):
	from .coordinates import UnstructuredCoords
	from .cartesian_grid import CartesianGrid

	r = self.r
	theta = self.theta
	x = r * np.cos(theta)
	y = r * np.sin(theta)
	return CartesianGrid(UnstructuredCoords([x, y]))

Grid._add_coordinate_system('polar', PolarGrid)

Grid._add_coordinate_system_transformation('cartesian', 'polar', _cartesian_to_polar)
Grid._add_coordinate_system_transformation('polar', 'cartesian', _polar_to_cartesian)
