import numpy as np
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
from fourierOptics import *
from pdb import set_trace

def field_angle(hx, hy):
	fx = hx * 0.0075
	fy = hy * 0.0075
	return fx, fy


if __name__ == '__main__':
	# Initializing MICADO paraxial model
	mcd = ParaxialMICADO()
	wfe = WavefrontError(mcd.n, mcd.m)

	# Uncomment to load in WFE (coma)
	# Problem here is that the field of view is too small to see 
	# a radial effect...
	# xin, yin = load_variable('../adc_data_coordinates.pkl')
	# adcObj = ADC_surfaces(xin,yin)
	# zern = wfe.zernike_cartesian([0,0,0,0,0,0,0,100e-5], adcObj.xz, adcObj.yz, rescale=False)
	# mcd.ELT_ADC.p1.add_wavefront_error(zern)
	
	# Setup
	l = 1.49e-3
	Elist = []

	# Set field points
	hxList, hyList = np.mgrid[-0.03:0.03:3j,-0.03:0.03:3j]
	hxList = hxList.flatten()
	hyList = hyList.flatten()
	
	# Iterating over the field positions.
	for i in range(len(hxList)):
		hx = hxList[i]
		hy = hyList[i]

		# Calculate the relative wavefront tilt for the field position		
		fieldx, fieldy = field_angle(hx,hy)
		field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=fieldx, ty=fieldy)

		# Apply the field tilt
		Ein = Profiles().ELT_pupil(mcd.n,mcd.m,mcd.dpix).astype('complex128')\
				* np.exp(-1j * 2*np.pi/l * field_tilt)
		mcd.apply_mask(Ein, replace=True)

		# Propagate through the MICADO system
		E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
								print_progress=False)
		Elist.append(E)

	# Average over all the wavelengths to get the final field.
	Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) 	# Amplitude
	Eci = Eca**2 												# Intensity
	Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase

	# Convert the scale to sky-angles using the plate scale
	x *= mcd.platescale*1000 			# x in milli arcseconds
	y *= mcd.platescale*1000 			# y in milli arcseconds

	# Plotting the results
	fig = plt.figure(figsize=(8,3))
	ax = fig.add_subplot(121)
	im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
	ax.set_xlim([x.min(),x.max()])
	ax.set_ylim([y.min(),y.max()])
	ax.set_xlabel('x [mas]')
	ax.set_ylabel('y [mas]')
	fig.colorbar(im, ax=ax, label='Amplitude')
	ax2 = fig.add_subplot(122)
	im2 = ax2.imshow(Ecp, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()])
	ax2.set_xlim([x.min(),x.max()])
	ax2.set_ylim([y.min(),y.max()])
	ax2.set_xlabel('x [mas]')
	ax2.set_ylabel('y [mas]')
	fig.colorbar(im2, ax=ax2, label='Phase')
	plt.tight_layout()
	plt.show()