import numpy as np
import os
import sys
import matplotlib.pyplot as plt
import matplotlib.cm as cm
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 centroid shift analysis.
					Script made for the SPIE manuscript.
'''


def plot_colourline(x,y,c, ax):
    '''
    	Source: https://stackoverflow.com/questions/36505587/color-line-by-third-variable-python
    	Function to colour a line according to a thrid variable.
    	In our case, this will be the Strehl ratio
    '''
    c = cm.plasma_r(c)
    for i in np.arange(len(x)-1):
        ax.plot([x[i],x[i+1]], [y[i],y[i+1]], c=c[i])
    return




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

	# Load in the results from centroid_shift_analysis.py
	rms_zern_nm, r_airy_zern_mas, r_err_zern_airy_mas, strehl_list_zern = load_variable('./data/centroid_shift_zernikes.pkl')
	rms_psd_nm, r_airy_psd_mas, r_err_psd_airy_mas, strehl_list_psd = load_variable('./data/centroid_shift_psd.pkl')

	# Convert the errors to logarithmic space, assuming small errors.
	log_err_zern_airy_mas = 0.434 * r_err_zern_airy_mas / r_airy_zern_mas
	log_err_psd_airy_mas = 0.434 * r_err_psd_airy_mas / r_airy_psd_mas


	# Plot the resulting WFE scaling relation
	fig = plt.figure(figsize=(9,7))
	ax1 = fig.add_subplot(111)
	ax1.errorbar((0.8138)*rms_zern_nm, r_airy_zern_mas, yerr=log_err_zern_airy_mas, color='k',
				linestyle='none', capsize=3, 
				markeredgewidth=1)
	ax1.errorbar((0.8138)*rms_psd_nm, r_airy_psd_mas, yerr=log_err_psd_airy_mas, color='k',
				linestyle='none', capsize=3, 
				markeredgewidth=1)
	ax1.scatter((0.8138)*rms_zern_nm, r_airy_zern_mas, c=strehl_list_zern, marker='v', label='Zernikes', cmap='plasma_r')
	
	im = ax1.scatter((0.8138)*rms_psd_nm, r_airy_psd_mas, c=strehl_list_psd, marker='o', label='PSD', cmap='plasma_r')

	# We colour the lines by their Strehl ratio
	plot_colourline((0.8138)*rms_zern_nm, r_airy_zern_mas, strehl_list_zern, ax=ax1)
	plot_colourline((0.8138)*rms_psd_nm, r_airy_psd_mas, strehl_list_psd, ax=ax1)
	ax1.set_xscale('log')
	ax1.set_yscale('log')
	ax1.set_xlabel('$\\sigma_{\\mathrm{WFE}}$ (nm)', fontsize=18)
	ax1.set_ylabel('Centroid shift (mas)', fontsize=18)
	ax1.set_xlim([np.min(rms_zern_nm)*(0.8138), np.max(rms_zern_nm)*(0.8138)])
	ax1.legend(fontsize=18)
	ax1.grid(True, alpha=0.3, which='major')
	ax1.grid(True, alpha=0.1, which='minor')
	ax1.tick_params(labelsize=18)
	cb = fig.colorbar(im, ax=ax1)
	cb.set_label(label='Strehl', size=18)
	cb.ax.tick_params(labelsize=18)
	plt.savefig('./figures/WFE_scaling_relation_obscured_airy.png', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/WFE_scaling_relation_obscured_airy.pdf', bbox_inches='tight', transparent=True)
	plt.tight_layout()
	plt.show()