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


'''
	Author: 		Joost van den Born, PhD student at University of Groningen
	Contact: 		born@astro.rug.nl
	Description: 	Script to generate the wavefront errors that will be used for
					the ADC prisms during all subsequent analyses. Script made for 
					the SPIE manuscript.
'''

def main():
	plt.rc('text', usetex=True)
	plt.rc('font', family='serif')
	# Initializing MICADO paraxial model

	l = 1.49e-3
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
	wfe = WavefrontError(mcd.n, mcd.m)

	# Get x3 and y3 (the coordinates to apply the ADC WFE to):
	_, x3, y3 = mcd.propagate(l=l,surface=4, zoom=False, show=False, 
								print_progress=False)

	# Apply a set of wavefront errors to each ADC surface.
	print('\nSetting the wavefront errors...')
	wfe_p1 = ADC_surfaces(x3,y3)
	wfe_p2 = ADC_surfaces(x3,y3)
	wfe_p3 = ADC_surfaces(x3,y3)
	wfe_p4 = ADC_surfaces(x3,y3)

	wfe_p1.generate_ADC_wfe(zernRMS=20e-6, psdRMS=15e-6)
	wfe_p2.generate_ADC_wfe(zernRMS=20e-6, psdRMS=15e-6)
	wfe_p3.generate_ADC_wfe(zernRMS=20e-6, psdRMS=15e-6)
	wfe_p4.generate_ADC_wfe(zernRMS=20e-6, psdRMS=15e-6)
	
	mask = np.sqrt(wfe_p1.xz**2 + wfe_p1.yz**2) > 1

	# Plot P1
	
	surf_p1 = wfe_p1.latest_surface.copy()
	surf_p1[mask]=np.nan
	figp1, axp1 = plot_real_valued_results(surf_p1[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
	 						 				colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
	plt.savefig('./figures/ADC_tWFE_p1_zern_20nm_psd_15nm.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/ADC_tWFE_p1_zern_20nm_psd_15nm.png', bbox_inches='tight', transparent=True)	
	plt.close()

	# Plot P2
	surf_p2 = wfe_p2.latest_surface.copy()
	surf_p2[mask]=np.nan
	figp2, axp2 = plot_real_valued_results(surf_p2[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
	 						 				colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
	plt.savefig('./figures/ADC_tWFE_p2_zern_20nm_psd_15nm.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/ADC_tWFE_p2_zern_20nm_psd_15nm.png', bbox_inches='tight', transparent=True)	
	plt.close()

	# Plot P3
	surf_p3 = wfe_p3.latest_surface.copy()
	surf_p3[mask]=np.nan
	figp3, axp3 = plot_real_valued_results(surf_p3[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
	 						 				colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
	plt.savefig('./figures/ADC_tWFE_p3_zern_20nm_psd_15nm.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/ADC_tWFE_p3_zern_20nm_psd_15nm.png', bbox_inches='tight', transparent=True)	
	plt.close()

	# Plot P4
	surf_p4 = wfe_p4.latest_surface.copy()
	surf_p4[mask]=np.nan
	figp4, axp4 = plot_real_valued_results(surf_p4[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
	 						 				colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
	plt.savefig('./figures/ADC_tWFE_p4_zern_20nm_psd_15nm.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/ADC_tWFE_p4_zern_20nm_psd_15nm.png', bbox_inches='tight', transparent=True)	
	plt.close()

	# Get the wavefront errors and save them for later use.
	save_variable('./data/ADC_tWFE_p1_zern_20nm_psd_15nm.pkl', wfe_p1.latest_surface)
	save_variable('./data/ADC_tWFE_p2_zern_20nm_psd_15nm.pkl', wfe_p2.latest_surface)
	save_variable('./data/ADC_tWFE_p3_zern_20nm_psd_15nm.pkl', wfe_p3.latest_surface)
	save_variable('./data/ADC_tWFE_p4_zern_20nm_psd_15nm.pkl', wfe_p4.latest_surface)


if __name__ == '__main__':
	main()