import numpy as np
import os
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 generate the plots of the relation between the
					Strehl ratio and the RMS of the Wavefront Error. A comparison
					with the Marechal approximation is made.
	 				Script made for the SPIE manuscript.
'''

def expected_strehl2(rms, l):
	return np.exp( -(2*np.pi*rms/l)**2 )


def append_to_existing_data(A, filename):
	if os.path.exists(filename):
		A_o = load_variable(filename)
		A 	= np.append(A, A_o)
	return A


def plot_some_stuff():
	fig = plt.figure(figsize=(9,7))
	ax = fig.add_subplot(111)
	ax2 = ax.twiny()

	rmsValues = rmsList[:25]
	for i in range(len(rms_zern_list_complete)):
		ax.plot((0.8138)*rms_zern_list_complete[i]*1e6, strehl_zern_complete[i], 'y-', 
			alpha=0.1, linewidth=.5)
		ax.plot((0.8138)*rms_psd_list_complete[i]*1e6, strehl_psd_complete[i], 'b-', 
			alpha=0.1, linewidth=.5)

	ax.plot(np.linspace(1e-6,1e-3,100)*1e6, expected_strehl2(np.linspace(1e-6,1e-3,100),l),'k--', 
				label='$S=\\exp[-(2\\pi\\sigma/\\lambda)^2]$')
	ax.errorbar((0.8138)*rmsValues*1e6, mean_strehl_zern, yerr=err_strehl_zern, color='y',
				linestyle='-', label='Zernikes', capsize=3, markeredgewidth=1)#,
				# linewidth=1)
	ax.errorbar((0.8138)*rmsValues*1e6, mean_strehl_psd, yerr=err_strehl_psd, color='b',
				linestyle='-', label='PSD', capsize=3, markeredgewidth=1)#,
				# linewidth=1)
	ax.set_xlabel('$\\sigma_{\\mathrm{WFE}}$ (nm)', fontsize=18)
	ax.set_ylabel('Strehl', fontsize=18)
	ax.set_ylim([0,1])
	ax.set_xlim([np.min(rmsValues)*1e6*(0.8138), np.max(rmsValues)*1e6*(0.8138)])
	ax2.set_xlim([np.min(rmsValues*(0.8138)/l), np.max(rmsValues*(0.8138)/l)])
	ax2.set_xlabel('$\\sigma/\\lambda$', fontsize=18)
	ax.legend(fontsize=18)
	ax.grid(True, alpha=0.3)
	ax.tick_params(labelsize=18)
	ax2.tick_params(labelsize=18)
	plt.tight_layout()
	plt.savefig('./figures/SR_relation_WFE_on_ADC_p1_obscured.png', bbox_inches='tight', transparent=True)
	plt.savefig('./figures/SR_relation_WFE_on_ADC_p1_obscured.pdf', bbox_inches='tight', transparent=True)
	plt.show()

if __name__ == '__main__':
	plt.rc('font', family='serif')
	plt.rc('text', usetex=True)
	plt.rc('axes', labelsize='large')
	
	l = 1.49e-3
	rmsList = load_variable('./data/SR_RMS_values_obscured.pkl')
	strehl_zern_complete = load_variable('./data/SR_strehl_zern_obscured.pkl')
	strehl_psd_complete  = load_variable('./data/SR_strehl_psd_obscured.pkl') 
	rms_zern_list_complete = load_variable('./data/SR_rms_zern_complete_obscured.pkl')
	rms_psd_list_complete = load_variable('./data/SR_rms_psd_complete_obscured.pkl')
	
	mean_strehl_zern 	= np.mean(strehl_zern_complete, axis=0)
	err_strehl_zern 	= np.std(strehl_zern_complete, axis=0)
	mean_strehl_psd 	= np.mean(strehl_psd_complete, axis=0)
	err_strehl_psd 		= np.std(strehl_psd_complete, axis=0)

	plot_some_stuff()