import numpy as np
import os
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
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 analysis where the shift
					of the PSF centroid as a function of the RMS value of the 
					WFE was investigated.
					Script made for the SPIE manuscript.
'''


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

	# Optical system constant
	platescale = 100.13151872258219

	###############################
	### Lowpass filter analysis ###
	###############################

	# Load in the data
	r_airy, r_err_airy = load_variable('./data/WFE frequency analysis/lowpass_frequency_filter_airy_distance.pkl')
	x0_airy, y0_airy, x0_airy_err, y0_airy_err = load_variable('./data/WFE frequency analysis/lowpass_frequency_filter_airy_locations.pkl')
	central_frequencies = np.asarray([0.5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 25,
			 30, 35, 40, 45, 50, 60, 70, 80, 90, 100, 120, 140, 160, 180, 200,
			 220, 250, 280, 310, 340, 370, 400, 450, 500])

	r_airy_mas = r_airy * platescale
	r_err_airy_mas = r_err_airy * platescale

	# Plotting the centroid shift for all bandpasses
	fig = plt.figure(figsize=(6.4,4.8))
	ax1 = fig.add_subplot(111)
	ax1.scatter(central_frequencies, r_airy_mas, edgecolor='k', facecolors='none', alpha=0.7)
	ax1.set_xscale('log')
	ax1.set_xlabel('Cut-off Frequency of the lowpass filter ($\\rho/\\rho_c$)', fontsize=18)
	ax1.set_ylabel('Centroid shift (mas)', fontsize=18)
	ax1.tick_params(axis='both', which='both', direction='in', labelsize=16, zorder=-1,)
	ax1.grid(True, alpha=0.3, which='major')
	ax1.grid(True, alpha=0.1, which='minor')
	ax1.set_xlim([0.3,550])
	plt.tight_layout()
	figtitle = 'new_lowpass_spatial_freq_distance' 
	if figtitle:
		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()
	# plt.savefig('./figures/wfe_frequency_analysis/lowpass_spatial_freq_distance.png', bbox_inches='tight', transparent=True)
	# plt.savefig('./figures/wfe_frequency_analysis/lowpass_spatial_freq_distance.pdf', bbox_inches='tight', transparent=True)





	################################
	### Bandpass filter analysis ###
	################################

	# Load in the data
	r_airy, r_err_airy = load_variable('./data/WFE frequency analysis/bandpass_frequency_filter_airy_distance.pkl')
	x0_airy, y0_airy, x0_airy_err, y0_airy_err = load_variable('./data/WFE frequency analysis/bandpass_frequency_filter_airy_locations.pkl')
	central_frequencies = load_variable('./data/WFE frequency analysis/bandpass_frequency_filter_central_frequencies.pkl')
	print(central_frequencies, len(central_frequencies), len(r_airy))
	r_airy_mas = r_airy * platescale
	r_err_airy_mas = r_err_airy * platescale


	# Fit a line in Logarithmic space
	log_cf = np.log10(central_frequencies)
	log_r  = np.log10(r_airy_mas)
	coefficients = np.polyfit(log_cf[2:], log_r[2:], 1)
	best_fit_fnc = np.poly1d(coefficients)
	best_fit_y = best_fit_fnc(log_cf)


	# Plotting the centroid shift for all bandpasses
	fig = plt.figure(figsize=(6.4,4.8))
	ax1 = fig.add_subplot(111)
	ax1.scatter(central_frequencies, r_airy_mas, edgecolor='k', facecolors='none', alpha=0.7)
	ax1.plot(central_frequencies[2:], 10**best_fit_y[2:], 'b-', alpha=0.5, label='Best fit: $y = 0.030 x^{-0.602}$')
	ax1.set_xscale('log')
	ax1.set_yscale('log')
	ax1.set_xlabel('Central Frequency of the bandpass filter ($\\rho/\\rho_c$)', fontsize=18)
	ax1.set_ylabel('Centroid shift (mas)', fontsize=18)
	ax1.tick_params(axis='both', which='both', direction='in', labelsize=16, zorder=-1,)
	ax1.grid(True, alpha=0.3, which='major')
	ax1.grid(True, alpha=0.1, which='minor')
	ax1.set_xlim([0.3,550])
	ax1.set_ylim([10**-6, 10**0])
	ax1.legend(fontsize=18)
	plt.tight_layout()
	figtitle = 'bandpass_spatial_freq_distance'
	if figtitle:
		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.savefig('./figures/wfe_frequency_analysis/bandpass_spatial_freq_distance.png', bbox_inches='tight', transparent=True)
	# plt.savefig('./figures/wfe_frequency_analysis/bandpass_spatial_freq_distance.pdf', bbox_inches='tight', transparent=True)
	plt.show()