import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.ticker as plticker
from matplotlib.gridspec import GridSpec
import pandas as pd
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *
from pdb import set_trace

'''
	Date: 			2022-06-20
	Author: 		Joost van den Born
	Description: 	Plotting the geometric distortions, that were calculated
					in an earlier simulation.
					Here we also include the detector layout in the plot.
'''

def add_detector_layout(figureObj, lw=0.8, alpha=0.5, show_label=True):
	'''
	Adds the detector layout to an existing figure+axes object.
	The detector outlines are based on the MICADO Hawaii 4RG mosaic
	and the dimensions are in mm.
	
	Parameters
	----------
	figureObj : tuple
	    Tuple containing the figure and axes objects (fig,ax)
	lw : float, optional
	    Linewidth of the drawn detectors. Default: 0.8
	alpha : float, optional
	    Alpha value to draw the rectangles at
	show_label : bool, optional
	    If True, shows the detector ID number in the rectangle
	
	Returns
	-------
	TYPE
	    Updated fig, axes objects
	'''
	fig, ax = figureObj

	# Detector layout
	det_hdr = ['id','x_cen', 'y_cen', 'x_size', 'y_size', 'pixel_size', 'angle', 'gain']
	det_props = np.array([[1, -63.94,  63.94, 4096, 4096, 0.015,   0.0, 1.0],
						  [2,   0.00,  63.94, 4096, 4096, 0.015, 270.0, 1.0],
						  [3,  63.94,  63.94, 4096, 4096, 0.015, 180.0, 1.0],
						  [6, -63.94,   0.00, 4096, 4096, 0.015,   0.0, 1.0],
						  [5,   0.00,   0.00, 4096, 4096, 0.015, 180.0, 1.0],
						  [4,  67.94,   0.00, 4096, 4096, 0.015, 180.0, 1.0],
						  [7, -63.94, -63.94, 4096, 4096, 0.015,   0.0, 1.0],
						  [8,   0.00, -63.94, 4096, 4096, 0.015,  90.0, 1.0],
						  [9,  63.94, -63.94, 4096, 4096, 0.015, 180.0, 1.0]])
	detectors = pd.DataFrame(det_props, columns=det_hdr)
	detectors["id"] 	= detectors["id"].astype(int)
	detectors["x_size"] = detectors["x_size"].astype(int)
	detectors["y_size"] = detectors["y_size"].astype(int)

	for ind, det in detectors.iterrows():
		x_mm = det['x_cen'] - det['x_size']/2*det['pixel_size']
		y_mm = det['y_cen'] - det['y_size']/2*det['pixel_size']
		width = det['x_size']*det['pixel_size']
		height = det['y_size']*det['pixel_size']
		x_cen, y_cen = det['x_cen'], det['y_cen']
		
		rect = patches.Rectangle((x_mm, y_mm), width, height, 
				linewidth=lw, edgecolor='k', facecolor='none', alpha=alpha, zorder=1)
		ax.add_patch(rect)
		if show_label:
			ax.text(x_cen, y_cen, det['id'])
	return fig, ax

def convert_detector_size(x, y, old=True):
	old_ps 		= 100.13151872258219
	new_ps 		= ParaxialMICADO(n=3001, m=3001, dpix=50, l=1.49e-3).platescale * 1e3
	conv_fact 	= old_ps / new_ps
	
	if old:
		x *= conv_fact
		y *= conv_fact
	elif not old:
		x /= conv_fact
		y /= conv_fact
	
	return x, y
	

def plot_distortions(xe, ye, xp, yp, lims=False, add_detectors=True, 
						save=False, figtitle='geom_distortions',
						show=True):
	
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50, l=1.49e-3)
	n = xe.shape[0]
	X = xe.flatten()
	Y = ye.flatten()
	dX = xp - xe
	dY = yp - ye
	dx = dX.flatten()
	dy = dY.flatten()

	# Magnitude of the distortions
	R = np.sqrt(dx**2 + dy**2) * mcd.platescale * 1e6
	print('RMS of the PSF shifts: {:.4f} μas.'.format(np.sqrt(np.mean((R)**2))))

	# Dummy figure, used for the colorbar colors in the quiver plot
	im = plt.imshow(np.reshape(R, (n,n)))
	plt.close()

	# Initiate the figure
	fig = plt.figure(figsize=(6.4,4.8))
	ax = fig.add_subplot(aspect='equal')
	
	# Add the detector layout if so desired.
	if add_detectors:
		add_detector_layout((fig,ax), lw=.8, alpha=0.5, show_label=False)
	
	# Plot the vector plot
	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=.8, units='dots', zorder=300)
	ax.quiver(X,Y,dx,dy,R, alpha=1, units='dots', zorder=300)
	ax.set_xlabel('x (mm)', fontsize=16)
	ax.set_ylabel('y (mm)', fontsize=16)
	ax.tick_params(axis='both', direction='in', labelsize=14, zorder=-1,
										top=True, right=True)
	cb = fig.colorbar(im, ax=ax)
	cb.set_label(label='PSF centroid shift (μas)', size=14)
	cb.ax.tick_params(labelsize=14)
	if not lims:
		ax.set_xlim([-125, 115])
		ax.set_ylim([-115, 115]) 
	elif lims:
		ax.set_xlim(lims[0])
		ax.set_ylim(lims[1]) 
	plt.tight_layout()
	if save:
		plt.savefig('{}.pdf'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
		plt.savefig('{}.png'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
	if show:
		plt.show()
	return fig, ax, R

def get_differential_distortion_vectors(xe, ye, xp, yp, xe_ref, ye_ref, xp_ref, yp_ref):
	# Reference distortions:
	# Shape of the arrays
	n = xe_ref.shape[0]

	# Reshaping of the arrays so that they can be used for the quiver plot.
	X_ref  = xe_ref.flatten()
	Y_ref  = ye_ref.flatten()
	dX_ref = xp_ref - xe_ref
	dY_ref = yp_ref - ye_ref

	dx_ref = dX_ref.flatten()
	dy_ref = dY_ref.flatten()

	# Distortions to compare
	# 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()

	ddx = dX - dX_ref
	ddy = dY - dY_ref

	return ddx, ddy


def plot_differential_distortions(xe, ye, xp, yp, xe_ref, ye_ref, xp_ref, yp_ref,
				lims=False, add_detectors=True,	save=False, 
				figtitle='diff_geom_distortions', show=True):
	# Reference distortions:
	# Shape of the arrays
	n = xe_ref.shape[0]

	# Reshaping of the arrays so that they can be used for the quiver plot.
	X_ref  = xe_ref.flatten()
	Y_ref  = ye_ref.flatten()
	dX_ref = xp_ref - xe_ref
	dY_ref = yp_ref - ye_ref

	dx_ref = dX_ref.flatten()
	dy_ref = dY_ref.flatten()

	# Distortions to compare
	# 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()

	ddx = dX - dX_ref
	ddy = dY - dY_ref

	# Two values are comparatively (extremely large), they are set to 0
	# to be able to see the other values.
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50, l=1.49e-3)
	R = (np.sqrt(ddx**2 + ddy**2) * mcd.platescale * 1e6) # Magnitude of the shift in micro arcseconds
	print('RMS of the PSF shifts compared to the reference position: {:.4f} μas.'.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=(6.4, 4.8))
	ax = fig.add_subplot()

	# Add the detector layout if so desired.
	if add_detectors:
		add_detector_layout((fig,ax), lw=.8, alpha=0.5, show_label=False)
	
	# Plot the vector plot
	for i in range(len(xe_ref)):
	    ax.plot(xe_ref[:,i], ye_ref[:,i], 'k-', alpha=0.2, linewidth=.5)
	    ax.plot(xe_ref[i,:], ye_ref[i,:], 'k-', alpha=0.2, linewidth=.5)
	ax.quiver(X,Y,ddx,ddy, angles='xy', edgecolor='k', 
						facecolor='None', linewidth=.8, units='dots', zorder=300)
	ax.quiver(X,Y,ddx,ddy,R, alpha=1, units='dots', zorder=300)
	ax.set_xlabel('x (mm)', fontsize=16)
	ax.set_ylabel('y (mm)', fontsize=16)
	ax.tick_params(axis='both', direction='in', labelsize=14, zorder=-1,
										top=True, right=True)
	cb = fig.colorbar(im, ax=ax)
	cb.set_label(label='PSF centroid shift (μas)', size=14)
	cb.ax.tick_params(labelsize=14)
	if not lims:
		ax.set_xlim([-125, 115])
		ax.set_ylim([-115, 115]) 
	elif lims:
		ax.set_xlim(lims[0])
		ax.set_ylim(lims[1])  
	plt.tight_layout()
	if save:
		plt.savefig('{}.pdf'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
		plt.savefig('{}.png'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
	if show:
		plt.show()
	return fig, ax, R

def get_data():
	# Coordinates without WFEs
	xe = load_variable('../data/GD_obscured_PSF_xcentroid_expected_zern_psd.pkl')
	ye = load_variable('../data/GD_obscured_PSF_ycentroid_expected_zern_psd.pkl')

	# Coordinates with WFEs
	xf = load_variable('../data/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd.pkl')
	yf = load_variable('../data/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd.pkl')
	return xe, ye, xf, yf


def normalize_relative_angle(theta, degrees=True):
	if type(theta) in [list, np.ndarray]:
		values = np.array([normalize_relative_angle(i) for i in theta])
		return values
	circle = 360 if degrees else 2*np.pi
	theta_mod = theta % circle
	# theta_mod_half_circle = theta % (circle/2)
	if theta_mod > (circle / 2):
		rel_angle = theta_mod - (circle)
	else:
		rel_angle = theta_mod
	return rel_angle

if __name__ == '__main__':
	ps = ParaxialMICADO(n=3001, m=3001, dpix=50, l=1.49e-3).platescale
	print(ps)

	plt.rc('font', family='serif')
	plt.rc('axes', labelsize='large')
	
	# convert_detector_size()
	# Plotting the original distortions
	xe_og = load_variable('../data/GD full field (Airy small offset)/GD_PSF_xcentroid_expected.pkl')
	ye_og = load_variable('../data/GD full field (Airy small offset)/GD_PSF_ycentroid_expected.pkl')
	xp_og = load_variable('../data/GD full field (Airy small offset)/GD_PSF_xcentroid_found_with_wfe_airy.pkl')
	yp_og = load_variable('../data/GD full field (Airy small offset)/GD_PSF_ycentroid_found_with_wfe_airy.pkl')
	xe_og, ye_og = convert_detector_size(xe_og, ye_og)
	xp_og, yp_og = convert_detector_size(xp_og, yp_og)
	figtitle = '../figures adapted for thesis/GD_full_field_psd'
	plot_distortions(xe_og, ye_og, xp_og, yp_og, show=False, save=False, figtitle=figtitle,
					lims=[[-120,130],[-120,120]])

	# Plotting the original distortions
	xe_og = load_variable('../data/GD_obscured_PSF_xcentroid_expected_zern_psd.pkl')
	ye_og = load_variable('../data/GD_obscured_PSF_ycentroid_expected_zern_psd.pkl')
	xp_og = load_variable('../data/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd.pkl')
	yp_og = load_variable('../data/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd.pkl')
	figtitle = '../figures adapted for thesis/GD_full_field_zern_psd'
	plot_distortions(xe_og, ye_og, xp_og, yp_og, show=False, save=False, figtitle=figtitle)


	# Plotting the distortions when the ADC can rotate.
	# First for a very small rotation, this will serve as the reference distortions
	# that are calibrated for
	xe = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_expected_zern_psd_rot0.001.pkl')
	ye = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_expected_zern_psd_rot0.001.pkl')
	xp = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd_rot0.001.pkl')
	yp = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd_rot0.001.pkl')
	figtitle = '../figures adapted for thesis/ADC rotation/GD_full_field_zern_psd_0'
	plot_distortions(xe, ye, xp, yp, show=False, save=False, figtitle=figtitle, lims=None)
	
	# And comparing it to the original (i.e. no rotation)
	figtitle = '../figures adapted for thesis/ADC rotation/DGD_full_field_zern_psd_0'
	plot_differential_distortions(xe, ye, xp, yp, xe_og, ye_og, xp_og, yp_og, 
				show=False, save=False, figtitle=figtitle, lims=None)

	# Next we plot all the simulated rotation values
	# rotangles = [0.001, 0.125, 0.25, 0.5, 0.75, 1, 2, 5, 10, 30, 45, 90, 180]
	rotangles = [0.001, 0.125, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 3, 3.5, 4, 4.5, 5, 6, 7, 8, 9, 10]
	rms_of_distortions, max_of_distortions = [], []
	vector_magnitude = np.zeros((121,len(rotangles)))
	for i,rx in enumerate(rotangles):
		# Full field
		xe2 = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_expected_zern_psd_rot{}.pkl'.format(rx))
		ye2 = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_expected_zern_psd_rot{}.pkl'.format(rx))
		xp2 = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(rx))
		yp2 = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(rx))
		
		figtitle = '../figures adapted for thesis/ADC rotation/DGD_full_field_zern_psd_{}'.format(rx)	
		_,_, R = plot_differential_distortions(xe2, ye2, xp2, yp2, xe, ye, xp, yp, show=False, save=False,
			figtitle=figtitle, lims=[[-120,120],[-120,120]])

		rms_of_distortions.append(np.sqrt(np.mean((R)**2)))
		max_of_distortions.append(np.max(R))
		Rflat = R.flatten()
		for j in range(len(Rflat)):
			vector_magnitude[j,i] = Rflat[j]


	gs = GridSpec(1, 2, width_ratios=[3,2])
	fig = plt.figure(figsize=(6.4*2,4.8))
	ax = fig.add_subplot(gs[0,0])
	ax2 = fig.add_subplot(gs[0,1], projection='polar')
	
	relangle = normalize_relative_angle(2*np.asarray(rotangles))
	for i in range(len(Rflat)):
		ax.plot(relangle, vector_magnitude[i], 'b-', linewidth=0.5, alpha=0.1)
	ax.plot(relangle, rms_of_distortions, 'b-', label='RMS', zorder=30)
	ax.plot(relangle, rms_of_distortions, 'bx', zorder=30)
	ax.plot(relangle, max_of_distortions, 'r-', label='Maximum', zorder=30)
	ax.plot(relangle, max_of_distortions, 'rx', zorder=30)

	ax.axhline(50, color='k', linestyle='-', linewidth=1, label='50 μas', zorder=4)
	ax.axhline(10, color='k', linestyle='-.', linewidth=1, label='10 μas', zorder=4)
	ax.grid(True, alpha=0.3)
	ax.tick_params(axis='both', direction='in', labelsize=14, zorder=-1,
										top=True, right=True)
	ax.set_xlabel('Relative prism rotation (deg)', fontsize=16)
	ax.set_ylabel('Differential distortion (μas)', fontsize=16)
	ax.set_xlim([0,20])
	ax.set_ylim([-3,100])
	# plt.tight_layout()
	# save = False
	# if save:
	# 	figtitle = '../figures adapted for thesis/rotation_relation_relangles'
	# 	plt.savefig('{}.pdf'.format(figtitle), 
	# 		bbox_inches='tight', 
	# 		transparent=False, 
	# 		facecolor=(1, 1, 0, 0.0))
	# 	plt.savefig('{}.png'.format(figtitle), 
	# 		bbox_inches='tight', 
	# 		transparent=False, 
	# 		facecolor=(1, 1, 0, 0.0))

	# Trying to plot the change in the distortion vector as a function of ADC rotation
	
	ax2.grid(True, alpha=0.3)
	vxs, vys = [], []
	for i,rx in enumerate(rotangles):
		xe2 = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_expected_zern_psd_rot{}.pkl'.format(rx))
		ye2 = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_expected_zern_psd_rot{}.pkl'.format(rx))
		xp2 = load_variable('../data/ADC rotation/GD_obscured_PSF_xcentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(rx))
		yp2 = load_variable('../data/ADC rotation/GD_obscured_PSF_ycentroid_found_with_wfe_airy_zern_psd_rot{}.pkl'.format(rx))

		# Find the distortion vector
		ddx, ddy = get_differential_distortion_vectors(xe2, ye2, xp2, yp2, xe, ye, xp, yp)
		cx = xe2.flatten()
		cy = ye2.flatten()
		# print(cx[ci], cy[ci])
		vx = ddx.flatten() * ps * 1e6
		vy = ddy.flatten() * ps * 1e6
		
		vxs.append([vx])
		vys.append([vy])

	phi, mag = [],[]
	for i in range(121):
		# set_trace()
		# Array slicing: rotation, 0, grid index
		vx_evo = np.asarray(vxs)[:,0,i]
		vy_evo = np.asarray(vys)[:,0,i]

		phi_evo = np.arctan2(vy_evo, vx_evo)
		mag_evo = np.sqrt(vx_evo**2 + vy_evo**2)

		phi.append(phi_evo[-1])
		mag.append(mag_evo[-1])

		# Plot the vectors on the central detectors separately
		if i in [48, 49, 50, 59, 60, 61, 70, 71, 72]:
			if i == 48:
				ax.plot(relangle, vector_magnitude[i], 'g-', linewidth=1, alpha=1, zorder=20, label='Central detector')
				ax2.plot(phi_evo, mag_evo, 'g-', linewidth=2, zorder=300)
			else:
				ax.plot(relangle, vector_magnitude[i], 'g-', linewidth=1, alpha=1, zorder=20)
				ax2.plot(phi_evo, mag_evo, 'g-', linewidth=2, zorder=300)
		else:
			ax2.plot(phi_evo, mag_evo, 'b-', linewidth=1, zorder=1)
	
	# Making a legend on the first sub figure in the right order
	handles, labels = ax.get_legend_handles_labels()
	order = [0,1,4,2,3]

	ax.legend([handles[idx] for idx in order],[labels[idx] for idx in order], fontsize=14)

	# Slicing the polar plot to the relevant part.
	phi = np.asarray(phi)
	ax2.set_thetamin(np.degrees(np.min(phi%(2*np.pi)))-10)
	ax2.set_thetamax(np.degrees(np.max(phi%(2*np.pi)))+10)

	label_position=ax2.get_rlabel_position()
	ax2.text(np.radians(label_position-10),ax2.get_rmax()/2.,'Differential\ndistortion (μas)',
	        rotation=np.degrees(np.min(phi%(2*np.pi)))-10,ha='center',va='center', fontsize=14)

	loc = plticker.MultipleLocator(base=np.radians(10)) # this locator puts ticks at regular intervals
	ax2.xaxis.set_major_locator(loc)
	ticks = [40, '', 60,'', 80, '', 100, '', 120, '', 140, '', 160, '', 180, '', 200, '', 220]
	ax2.xaxis.set_ticklabels(ticks)
	ax2.tick_params(axis='both', direction='in', labelsize=12, zorder=-1,
										top=True, right=False, left=True)
	plt.tight_layout()
	save = True
	if save:
		figtitle = '../figures adapted for thesis/rotation_relation_with_polar_plot'
		plt.savefig('{}.pdf'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
		plt.savefig('{}.png'.format(figtitle), 
			bbox_inches='tight', 
			transparent=False, 
			facecolor=(1, 1, 0, 0.0))
	plt.show()