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 plot the results of the geometric distortion analysis.
					Script made for the SPIE manuscript.
'''

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

	# Full field
	xe = load_variable('../SPIE/data/GD full field (Obscured Airy zern psd)/GD_obscured_PSF_xcentroid_expected_zern_psd.pkl')
	ye = load_variable('../SPIE/data/GD full field (Obscured Airy zern psd)/GD_obscured_PSF_ycentroid_expected_zern_psd.pkl')
	xp = load_variable('../SPIE/data/GD full field (Obscured Airy zern psd)/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd.pkl')
	yp = load_variable('../SPIE/data/GD full field (Obscured Airy zern psd)/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd.pkl')

	# Shape of the arrays
	n = xe.shape[0]

	# Reshaping of the arrays so that they can be used for the quiver plot.
	X = xe.flatten()
	Y = ye.flatten()
	dX = xp-xe
	dY = yp-ye

	dx = dX.flatten()
	dy = dY.flatten()

	# Two values are comparatively (extremely large), they are set to 0
	# to be able to see the other values.
	R = (np.sqrt(dx**2 + dy**2) * 100.13151872258219 * 1000) # Magnitude of the shift in micro arcseconds
	print('RMS of the PSF shifts: {:.4f} uas.'.format(np.sqrt(np.mean((R)**2))))


	# Dummy figure
	im = plt.imshow(np.reshape(R, (n,n)))
	plt.close()

	# Actual plotting of the figure
	fig = plt.figure(figsize=(9,7))
	ax = fig.add_subplot()
	for i in range(len(xe)):
	    ax.plot(xe[:,i], ye[:,i], 'k-', alpha=0.2, linewidth=.5)
	    ax.plot(xe[i,:], ye[i,:], 'k-', alpha=0.2, linewidth=.5)
	ax.quiver(X,Y,dx,dy, angles='xy', edgecolor='k', facecolor='None', linewidth=1)
	ax.quiver(X,Y,dx,dy,R, alpha=1)
	ax.set_aspect('equal', adjustable='box')
	ax.set_xlabel('x (mm)', fontsize=18)
	ax.set_ylabel('y (mm)', fontsize=18)
	ax.set_xlim([-340,340])
	ax.set_ylim([-340,340])
	ax.tick_params(labelsize=18)
	cb = fig.colorbar(im, ax=ax)
	cb.set_label(label='PSF centroid shift ($\\mu$as)', size=18)
	cb.ax.tick_params(labelsize=18)
	plt.savefig('./figures/GD_full_field_obscured_Airy_psd_zern.pdf', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/GD_full_field_obscured_Airy_psd_zern.png', bbox_inches='tight', transparent=True)
	plt.tight_layout()
	plt.show()