import numpy as np
import sys
# sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *
import matplotlib.pyplot as plt
from pdb import set_trace

if __name__ == '__main__':
	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')

	# Initiate MICADO paraxial model
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)

	# Generate the PSF and plot it
	# E, x, y = mcd.propagate(l=1.35e-3,surface=-1, zoom=True, show=True, save='all')

	confList = [ ( 1, [-0.4, 0.4, -0.4, 0.4], 	'02_fourier_ELT-IFP.pdf'),
				 ( 2, [-50, 50, -50, 50], 		'03_fourier_ELT-COL.pdf'),
				 ( 4, [-50, 50, -50, 50], 		'04_fourier_ELT-MCD-ADC-P2.pdf'),
				 # (-1, [-1.2, 1.2, -1.2, 1.2], 	'05_fourier_ELT-MCD-DET.pdf')]
				 (-1, [-1.2/2.669, 1.2/2.669, -1.2/2.669, 1.2/2.669], 	'05_fourier_ELT-MCD-DET.pdf')]


	for i in range(len(confList)):
		s = confList[i][0]
		coordinates = confList[i][1]
		savename = confList[i][2]
		E, x, y = mcd.propagate(l=1.35e-3, surface=s, image_output=coordinates,
								 zoom=False, show=True, save=False)

	Eout, x, y = mcd.propagate(l=1.35e-3, surface=0, zoom=False, show=False)
	x /= 1e3 # convert to m
	y /= 1e3 # convert to m
	_fig = plt.figure(figsize=(8,3))
	_ax = _fig.add_subplot(121)
	_im = _ax.imshow(np.abs(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
	_ax.set_xlim([-30, 30])
	_ax.set_ylim([-30, 30])
	_ax.set_xlabel('x [m]')
	_ax.set_ylabel('y [m]')
	_fig.colorbar(_im, ax=_ax, label='Amplitude')
	_ax2 = _fig.add_subplot(122)
	_im2 = _ax2.imshow(np.angle(Eout), cmap='hot', extent=[x.min(), x.max(), y.min(), y.max()])
	_ax2.set_xlim([-30, 30])
	_ax2.set_ylim([-30, 30])
	_ax2.set_xlabel('x [m]')
	_ax2.set_ylabel('y [m]')
	_fig.colorbar(_im2, ax=_ax2, label='Phase')
	plt.tight_layout()
	# plt.savefig('01_fourier_ELT-M1.pdf', bbox_inches='tight', transparent=True)
	plt.close()#show()