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

'''
	Author: 		Joost van den Born, PhD student at University of Groningen
	Contact: 		born@astro.rug.nl
	Description: 	Script to generate two plots of the footprint on MCD-CAM lens, 
					showing the impact of the WFEs on the amplitude and phase
					near the exit pupil.
	 				Script made for the SPIE manuscript.
'''

if __name__ == '__main__':
	plt.rc('text', usetex=True)
	plt.rc('font', family='serif')

	# Initializing MICADO paraxial model
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
	wfe = WavefrontError(mcd.n, mcd.m)
	
	# Setup
	l = 1.49e-3

	# Set field points
	hxList = [-1,-1,0, 1,1]
	hyList = [-1, 1,0,-1,1]
	
	# Propagate through the field positions without any WFEs	
	for i in range(len(hxList)):
		hx = hxList[i]
		hy = hyList[i]
		print('hx: {}	- 	hy: {}'.format(hx,hy) )

		# Propagate through the MICADO system
		prop_surf = 4
		E, x, y = mcd.propagate(l=l,surface=prop_surf, zoom=False, show=False, 
								print_progress=False, hx=hx, hy=hy)
		# Save the resulting 'perfect' PSF for later
		if prop_surf == 4:
			# Don't show anything outside the clear aperture.
			mask = np.sqrt(x**2 + y**2) > mcd.ELT_ADC.p1.CA/2

			# Plot P1
			footprint_p1 = E.copy()
			footprint_p1[mask]=np.nan
			# figp1, axp1, axp2 = plot_results(footprint_p1[850:2150,850:2150], x[850:2150,850:2150], y[850:2150,850:2150],
			# 							fig_size=(12,5), show=False)
			figp1, axp1 = plot_real_valued_results(np.abs(footprint_p1)[850:2150,850:2150], x[850:2150,850:2150], y[850:2150,850:2150],
										labeltext='Amplitude', show=False, fig_size=None)
			plt.savefig('./figures/ELT-MCD-ADC_footprint_no_wfe_{}_{}.pdf'.format(hx,hy), bbox_inches='tight', transparent=True)
			plt.savefig('./figures/ELT-MCD-ADC_footprint_no_wfe_{}_{}.png'.format(hx,hy), bbox_inches='tight', transparent=True)	
			plt.show()



	# Load in the WFEs and set them to the corresponding ADC surfaces
	wfe_p1 = load_variable('./data/ADC_tWFE_p1_psd_25nm.pkl')
	wfe_p2 = load_variable('./data/ADC_tWFE_p2_psd_25nm.pkl')
	wfe_p3 = load_variable('./data/ADC_tWFE_p3_psd_25nm.pkl')
	wfe_p4 = load_variable('./data/ADC_tWFE_p4_psd_25nm.pkl')

	mcd.ELT_ADC.p1.add_wavefront_error(wfe_p1)
	mcd.ELT_ADC.p2.add_wavefront_error(wfe_p2)
	mcd.ELT_ADC.p3.add_wavefront_error(wfe_p3)
	mcd.ELT_ADC.p4.add_wavefront_error(wfe_p4)



	# Iterating over the field positions.
	print('\nPropagation with Wavefront Errors applied')
	prop_surf = 4
	for i in range(len(hxList)):
		hx = hxList[i]
		hy = hyList[i]
		print('hx: {}	- 	hy: {}'.format(hx,hy) )

		# Propagate through the MICADO system
		Ep, xp, yp = mcd.propagate(l=l,surface=prop_surf+4, zoom=False, show=False, 
								print_progress=False, hx=hx, hy=hy)

		if prop_surf == 4:
			# Don't show anything outside the clear aperture.
			mask = np.sqrt(x**2 + y**2) > mcd.ELT_ADC.p1.CA/2

			# Plot P1
			footprint_p1 = Ep.copy()
			figp1, axp1, axp2 = plot_results(footprint_p1[1050:1350,1150:1450], x[1050:1350,1150:1450], y[1050:1350,1150:1450],
			 						   fig_size=(12,5), show=False)
			plt.savefig('./figures/ELT-MCD-COL_footprint_with_wfe_{}_{}_zoom.pdf'.format(hx,hy), bbox_inches='tight', transparent=True)
			plt.savefig('./figures/ELT-MCD-COL_footprint_with_wfe_{}_{}_zoom.png'.format(hx,hy), bbox_inches='tight', transparent=True)	
			plt.show()