import numpy as np
from ..field import Field

from ..optics import GeometricPhaseElement

def generate_app_keller(wavefront, propagator, contrast, num_iterations, beta=0):
	"""
	Accelerated Gerchberg-Saxton-like algorithm for APP design by
	Christoph Keller [Keller2016]_ and based on Douglas-Rachford operator splitting.
	The acceleration was inspired by the paper by Jim Fienup [Fienup1976]_. The
	acceleration can provide speed-ups of up to two orders of magnitude and
	produce better APPs.

	.. [Keller2016] Keller C.U., 2016, "Novel instrument concepts for
		characterizing directly imaged exoplanets", Proc. SPIE 9908,
		Ground-based and Airborne Instrumentation for Astronomy VI, 99089V
		doi: 10.1117/12.2232633; https://doi.org/10.1117/12.2232633

	.. [Fienup1976] J. R. Fienup, 1976, "Reconstruction of an object from the modulus
		of its Fourier transform," Opt. Lett. 3, 27-29

	Parameters
	----------
	wavefont : Wavefront
		The input aperture as a wavefront; a phase can be provided as a
		starting point for the APP generation.
	propagator : Propagator
		The propagator from the pupil grid to the focal plane grid.
	contrast : array_like
		The required contrast in the focal plane; this is a float mask
		that is 1.0 everywhere except for the dark zone where it is the
		contrast such as 1e-5.
	num_iterations : int
		The maximum number of iterations.
	beta : scalar
		The acceleration parameter. The default is 0 (no acceleration).
		Good values for beta are typically between 0.3 and 0.9. Values larger
		than 1.0 will not work.

	Returns
	-------
	Wavefront
		The APP as a wavefront.

	Raises
	------
	ValueError
		If beta is not between 0 and 1.
	"""
	if beta < 0 or beta > 1:
		raise ValueError('Beta should be between 0 and 1.')

	# initialize APP with wavefront
	app = wavefront.copy()

	# define pupil aperture where amplitude is larger than 1e-6
	aperture = wavefront.amplitude > 1e-6

	# define dark zone as location where contrast requirement is < 1e-1
	dark_zone = contrast < 0.1

	old_image = None
	for i in range(num_iterations):
		# calculate image plane electric field
		image = propagator.forward(app)

		# stop iteration if contrast requirement is met
		if not np.any(image.intensity / image.intensity.max() > contrast):
			break

		# modify focal plane electic field using acceleration
		new_image = image.copy()
		if beta != 0 and old_image is not None:
			new_image.electric_field[dark_zone] = old_image.electric_field[dark_zone] * beta - new_image.electric_field[dark_zone] * (1 + beta)
		else:
			new_image.electric_field[dark_zone] = 0
		old_image = new_image.copy()

		# determine corresponding aperture electric field
		app.electric_field = propagator.backward(new_image).electric_field

		# enforce aperture support
		app.electric_field[~aperture] = 0

		# enforce unity transmission within aperture support
		app.electric_field[aperture] *= wavefront.amplitude[aperture] / app.amplitude[aperture]

	return app

def generate_app_por(wavefront, propagator, propagator_max, contrast, num_iterations=1):
	'''Optimize a one-sided APP using a globally optimal algorithm. This algorithm does not
	apply any symmetries for two-sided dark zones or circularly-symmetric pupils. This
	function requires that you have installed the Gurobi optimizer.

	Parameters
	----------
	wavefont : Wavefront
		The input aperture as a wavefront.
	propagator : Propagator
		The propagator from `wavefront` to the focal plane grid.
	propagator_max : Propagator
		The propagator from `wavefront` to the point in the focal-plane that we want to maximize.
	contrast : array_like or scalar
		The required contrast in the focal plane.
	num_iterations : int
		The number of iterations for the Strehl ratio to converge.

	Returns
	-------
	Field
		The resultant electric field transmission function.
	'''
	import gurobipy as gp

	M = propagator.get_transformation_matrix_forward(wavefront.electric_field.grid, wavefront.wavelength)
	M_max = propagator_max.get_transformation_matrix_forward(wavefront.electric_field.grid, wavefront.wavelength)

	wf_0 = propagator.forward(wavefront).electric_field
	M /= wf_0[:,np.newaxis]

	M *= wavefront.electric_field
	M_max *= wavefront.electric_field

	M_max = np.sum(M_max, axis=0)

	M = np.array(np.bmat([[M.real, -M.imag], [M.imag, M.real]]))
	M_max = np.concatenate((M_max.real, -M_max.imag))

	m, n = M.shape
	contrast_requirement = np.ones(m) * np.sqrt(contrast)

	model = gp.Model('lp')
	model.Params.Threads = 4

	x = model.addVars(n, lb=-1, ub=1)

	obj = gp.quicksum((x[i] * M_max[i] for i in range(n)))
	model.setObjective(obj, gp.GRB.MAXIMIZE)

	for i in range(n//2):
		r2 = x[i]*x[i] + x[i+n//2]*x[i+n//2]
		model.addQConstr(r2 <= 1)

	for i, ee in enumerate(M):
		e = gp.quicksum((x[i] * ee[i] for i in range(n)))
		model.addConstr(e <= contrast_requirement[i])
		model.addConstr(e >= -contrast_requirement[i])

	model.optimize()

	solution = np.array([x[i].x for i in range(n)])
	solution = solution[:n//2] + 1j * solution[n//2:]

	return Field(solution, wavefront.electric_field.grid)

def generate_app_doelman(wavefront, propagator, propagator_max, contrast, num_iterations):
	raise NotImplementedError()

class VectorApodizingPhasePlate(GeometricPhaseElement):
	'''The vector-Apodizing Phase Plate (vAPP).

	Parameters
	----------
	phase_pattern : field or array_like
		The phase pattern generating the vAPP PSFs (radians).
	leakage : scalar
		The relative leakage strength (0 = no leakage, 1 = maximum leakage)
	retardance_offset : scalar
		The retardance offset from half wave in radians. This will result in leakage.
	'''
	def __init__(self, phase_pattern, leakage=None, retardance_offset=0):

		GeometricPhaseElement.__init__(self, phase_pattern, leakage, retardance_offset)
