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
from matplotlib.colors import Normalize, LogNorm

'''
	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 plot_distortions2(xe, ye, xp, yp, lims=False, add_detectors=True, 
						save=False, figtitle='geom_distortions',
						show=True, cmap='magma'):
	vmin, vmax = 0.1, 100
	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)))
	im = plt.imshow(np.reshape(R[:121], (11,11)))
	im = plt.scatter(X, Y, s=R, c=np.log10(R), cmap=cmap, vmin=np.log10(vmin), vmax=np.log10(vmax))
	# im = plt.scatter(X, Y, s=R, c=R, cmap='jet')
	plt.close()

	# Initiate the figure
	# fig = plt.figure(figsize=(9.6,4.8))
	# ax = fig.add_subplot(121, aspect='equal')
	# ax2 = fig.add_subplot(122, aspect='equal')
	fig, axes = plt.subplots(ncols=2, 
						   nrows=1, 
						   gridspec_kw={'width_ratios': [1,1]},
						   figsize=(9.6,4.8))
	ax, ax2 = axes
	ax.set_aspect('equal')
	ax2.set_aspect('equal')
	# Add the detector layout if so desired.
	if add_detectors:
		add_detector_layout((fig,ax), lw=.8, alpha=1, show_label=False)
	
	# Plot the vector plot
	xe = np.asarray(xe[:121]).reshape((11,11))
	ye = np.asarray(ye[:121]).reshape((11,11))
	# 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,R, alpha=1, units='dots', zorder=300)
	
	for k in range(1):
		print()
		ind_end = (k+1) * 121 
		ind_start = k * 121
		dpix = np.diff(Y)[ind_start:ind_end][1]
		print(k, ind_start, ind_end, dpix)

		xplot = np.asarray(X[ind_start:ind_end]).reshape((11,11)) 
		yplot = np.asarray(Y[ind_start:ind_end]).reshape((11,11))
		yplot, xplot = np.mgrid[yplot.min()+dpix/2:yplot.max()-dpix/2:12j,
								xplot.min()+dpix/2:xplot.max()-dpix/2:12j]


		Rtemp  = np.asarray(R[ind_start:ind_end]).reshape((11,11))
		ax.pcolor(-1*yplot.T, -1*xplot.T, Rtemp.T, cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
		# ax2.pcolor(-1*xplot.T, -1*yplot.T, Rtemp.T, cmap=cmap, norm=LogNorm(vmin=1, vmax=100))
		shiftx = X[ind_start:ind_end][61]#np.abs(dpix)/2
		shifty = Y[ind_start:ind_end][60]
		print(Y[ind_start:ind_end])
		print('shiftx:', shiftx)
		print('shifty:', shifty)

		# ax2.pcolor(k*shifty-1*yplot.T, k*shiftx-1*xplot.T, Rtemp.T*10, cmap=cmap, norm=LogNorm(vmin=1, vmax=100))
		if k == 1: ax2.pcolor(shifty-1*yplot.T, shiftx-1*xplot.T, Rtemp.T, cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
		elif k == 2: ax2.pcolor(2*shifty-1*yplot.T, 2*shiftx-1*xplot.T, Rtemp.T, cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
		elif k == 3: ax2.pcolor(2*shifty-1*yplot.T, 2*shiftx-1*xplot.T, Rtemp.T, cmap=cmap, norm=LogNorm(vmin=vmin, vmax=vmax))
		if k == 0:
			ax.quiver(X[ind_start:ind_end],Y[ind_start:ind_end],
					 dx[ind_start:ind_end],dy[ind_start:ind_end], 
					 angles='xy', edgecolor='k', facecolor='k', 
					 linewidth=.8, zorder=300, alpha=1, scale=1.4e-5, 
					 units='dots')
		elif k in [2,3]:
			ax2.quiver(X[ind_start:ind_end],Y[ind_start:ind_end],
					 dx[ind_start:ind_end],dy[ind_start:ind_end], 
					 angles='xy', edgecolor='k', facecolor='k', 
					 linewidth=.8, zorder=300, alpha=1, scale=0.2e-5, 
					 units='dots')
		# ax.pcolor(xplot * (1 - dpix/105), yplot * (1 - dpix/105), Rtemp, cmap='jet', norm=LogNorm(vmin=1, vmax=100))
				   
	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)
	ax2.set_xlabel('x (mm)', fontsize=16)
	ax2.set_ylabel('y (mm)', fontsize=16)
	ax2.tick_params(axis='both', direction='in', labelsize=14, zorder=-1,
										top=True, right=True)
	
	
	cb = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.04)
	cb.set_label(label='PSF centroid shift (μas)', size=14)
	cb.ax.tick_params(labelsize=14)
	
	lbls = 10**np.asarray(cb.ax.get_yticks())
	print(lbls)
	newlbls = ['{:.0f}'.format(x) for x in 10**np.asarray(cb.ax.get_yticks())]
	newlbls[1] = '{:.1f}'.format(lbls[1])
	newlbls[0] = '{:.1f}'.format(lbls[0])
	cb.ax.set_yticklabels(newlbls)
	print(cb.ax.get_yticks())
	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])
	ax2.set_xlim([-2.01,2.75])
	ax2.set_ylim([-2.01,2.75])
	ax.indicate_inset_zoom(ax2, edgecolor="black") 
	fig.subplots_adjust(top=1.0,
						bottom=0.0,
						left=0.1,
						right=0.8,
						hspace=0.135,
						wspace=0.3)
	# 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


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
	data_dir = '../data/GD Focal Plane WFEs/'
	xe_og = load_variable(data_dir + 'GD_IFP_obscured_PSF_xcentroid_expected_zern.pkl')
	ye_og = load_variable(data_dir + 'GD_IFP_obscured_PSF_ycentroid_expected_zern.pkl')
	xp_og = load_variable(data_dir + 'GD_IFP_obscured_PSF_xcentroid_found_with_wfe_airy_zern.pkl')
	yp_og = load_variable(data_dir + 'GD_IFP_obscured_PSF_ycentroid_found_with_wfe_airy_zern.pkl')

	# # Zoom 1: 
	# xe_og1 = load_variable(data_dir + 'GD_IFP_zoom_obscured_PSF_xcentroid_expected_zern.pkl')
	# ye_og1 = load_variable(data_dir + 'GD_IFP_zoom_obscured_PSF_ycentroid_expected_zern.pkl')
	# xp_og1 = load_variable(data_dir + 'GD_IFP_zoom_obscured_PSF_xcentroid_found_with_wfe_airy_zern.pkl')
	# yp_og1 = load_variable(data_dir + 'GD_IFP_zoom_obscured_PSF_ycentroid_found_with_wfe_airy_zern.pkl')

	xe = xe_og#np.append(xe_og, xe_og1)
	ye = ye_og#np.append(ye_og, ye_og1)
	xp = xp_og#np.append(xp_og, xp_og1)
	yp = yp_og#np.append(yp_og, yp_og1)

	# xe = np.append(xe_og2, xe_og3)
	# ye = np.append(ye_og2, ye_og3)
	# xp = np.append(xp_og2, xp_og3)
	# yp = np.append(yp_og2, yp_og3)


	figtitle = '../figures adapted for thesis/Focal Plane/GD_IFP_zern'
	plot_distortions(xe, ye, xp, yp, show=True, save=False, figtitle=figtitle,
					lims=[[-120,120],[-120,120]])