import numpy as np

from ..optics import Wavefront, AgnosticOpticalElement, make_agnostic_forward, make_agnostic_backward
from ..field import Field, evaluate_supersampled
from ..fourier import FastFourierTransform, make_fft_grid, FourierFilter

class AngularSpectrumPropagator(AgnosticOpticalElement):
	'''The monochromatic angular spectrum propagator for scalar fields.

	The scalar Angular Spectrum propagator is implemented as described by
	[McLeod2014]_. The propagation of an electric field can be described as a transfer
	function in frequency space. The transfer function is taken from
	equation 9 of [McLeod2014]_, and the related impulse response is taken from
	equation 6 of [McLeod2014]_.

	.. [McLeod2014] Robert R. McLeod and Kelvin H. Wagner 2014, "Vector Fourier optics of
		anisotropic materials," Adv. Opt. Photon. 6, 368-412 (2014)

	Parameters
	----------
	input_grid : anything
		This argument is ignored. The input grid is taken from the incoming wavefront.
	distance : scalar
		The distance to propagate
	num_oversampling : int
		The number of times the transfer function is oversampled. Default is 2.
	wavelength : scalar
		The wavelength of the wavefront.
	refractive_index : scalar
		The refractive index of the medium that the wavefront is propagating in.

	Raises
	------
	ValueError
		If the `input_grid` is not regular and Cartesian.
	'''
	def __init__(self, input_grid, distance, num_oversampling=2, refractive_index=1):
		self._distance = distance

		self._num_oversampling = num_oversampling
		self._refractive_index = refractive_index

		AgnosticOpticalElement.__init__(self, grid_dependent=True, wavelength_dependent=True)

	def make_instance(self, instance_data, input_grid, output_grid, wavelength):
		if not input_grid.is_regular or not input_grid.is_('cartesian'):
			raise ValueError('The input grid must be a regular, Cartesian grid.')

		k = 2 * np.pi / wavelength * self.evaluate_parameter(self.refractive_index, input_grid, output_grid, wavelength)
		L_max = np.max(input_grid.dims * input_grid.delta)

		if np.any(input_grid.delta < wavelength * abs(self.distance) / L_max):
			def transfer_function(fourier_grid):
				enlarged_grid = make_fft_grid(fourier_grid)
				fft_upscale = FastFourierTransform(enlarged_grid)

				def impulse_response(grid):
					r_squared = grid.x**2 + grid.y**2 + self.distance**2
					r = np.sqrt(r_squared)
					cos_theta = self.distance / r

					return Field(cos_theta / (2 * np.pi) * np.exp(1j * k * r) * (1 / r_squared - 1j * k / r), grid)

				impulse_response = evaluate_supersampled(impulse_response, enlarged_grid, self.num_oversampling)

				return fft_upscale.forward(impulse_response)
		else:
			def transfer_function_native(fourier_grid):
				k_squared = fourier_grid.as_('polar').r**2
				k_z = np.sqrt(k**2 - k_squared + 0j)

				return Field(np.exp(1j * k_z * self.distance), fourier_grid)

			def transfer_function(fourier_grid):
				return evaluate_supersampled(transfer_function_native, fourier_grid, self.num_oversampling)

		instance_data.fourier_filter = FourierFilter(input_grid, transfer_function, q=2)

	@property
	def distance(self):
		return self._distance

	@distance.setter
	def distance(self, distance):
		self._distance = distance

		self.clear_cache()

	@property
	def num_oversampling(self):
		return self._num_oversampling

	@num_oversampling.setter
	def num_oversampling(self, num_oversampling):
		self._num_oversampling = num_oversampling

		self.clear_cache()

	@property
	def refractive_index(self):
		return self._refractive_index

	@refractive_index.setter
	def refractive_index(self, refractive_index):
		self._refractive_index = refractive_index

		self.clear_cache()

	def get_input_grid(self, output_grid, wavelength):
		return output_grid

	def get_output_grid(self, input_grid, wavelength):
		return input_grid

	@make_agnostic_forward
	def forward(self, instance_data, wavefront):
		'''Propagate a wavefront forward by a certain distance.

		Parameters
		----------
		wavefront : Wavefront
			The incoming wavefront.

		Returns
		-------
		Wavefront
			The wavefront after the propagation.
		'''
		filtered = instance_data.fourier_filter.forward(wavefront.electric_field)

		return Wavefront(filtered, wavefront.wavelength, wavefront.input_stokes_vector)

	@make_agnostic_backward
	def backward(self, instance_data, wavefront):
		'''Propagate a wavefront backward by a certain distance.

		Parameters
		----------
		wavefront : Wavefront
			The incoming wavefront.

		Returns
		-------
		Wavefront
			The wavefront after the propagation.
		'''
		filtered = instance_data.fourier_filter.backward(wavefront.electric_field)

		return Wavefront(filtered, wavefront.wavelength, wavefront.input_stokes_vector)
