import numpy as np
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *


'''
	Description:
	This script generates a point spread function of the MICADO optical model.
	Here we apply a Bahtinov mask to the pupil and simulate a broadband PSF,
	done through the summation of multiple (here: 15) wavelength slices. 
	When there is atmospheric dispersion, the diffraction speckles created by the 
	Bahtinov mask will start to point to a location outside the zeroth order 
	PSF core.
	Here, atmospheric dispersion is modelled as a wavelength dependent phase
	tilt.
'''

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

	# Create the Bahtinov mask and set it on the entrance pupil (ELT-M1)
	bahtinov = Profiles().bahtinov_mask(40, 10, mcd.n, mcd.m, dpix=0.33, N=30, oversample=1)
	mcd.ELT_CAM.add_mask(bahtinov)

	# Set up
	Elist = []
	lList = np.linspace(1.49e-3, 1.78e-3, 15).tolist()

	# Dispersion of 40 mas over the band
	dispMax = 40e-3 / 3600 		# maximum dispersion in degrees
	dispersion = np.linspace(0, dispMax, len(lList)).tolist()

	# Iterate over the wavelengths
	for i in range(len(lList)):
		l = lList[i]

		# Generating the tilt to the input field.
		d = dispersion[i]
		field_tilt = wfe.generate_tilt(mcd.x,mcd.y,tx=0, ty=d)
		
		# Adding the tilt to the input field
		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 the input field throug the paraxial model
		E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
								image_output=[-4,4,-4,4], print_progress=True)
		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()