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 compare the Strehl ratio of 50 different instances
					of a generated WFE. Higher spatial frequency perturbations
					should follow the Marechal approximation more closely than 
					lower order shape errors only. We verify this using Power 
					Spectral Density and Zernike modes (66 coefficients).
					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)
		A2 	= np.append(A, A_o)
	else:
		A2 = A
	return A2


if __name__ == '__main__':
	# Setting some input parameters
	dpix 		= 160
	n, m 		= 1001, 1001
	l 			= 1.49e-3						# wavelength
	zern_rms 	= 25e-5							# RMS of the zernike component of the WFE
	midspat_rms = 25e-5							# RMS of the PSD component of the WFE
	rmsList 	= np.linspace(1e-6, 1e-3, 25) 	# 1 to 1000 nm WFE RMS in mm
	number_of_iterations = 50
	input_field_location = './data/ADC_input_field_1001x1001_1.49_onaxis_obscured.pkl'


	# Initiate ELT-MICADO and wavefront error objects
	mcd = ParaxialMICADO(l=l, n=n, m=m, dpix=dpix, entrance_pupil='open')
	wfe = WavefrontError(n,m)

	print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))

	# Propagate unperturbed field
	E, x, y = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
								print_progress=False, hx=0, hy=0,
								load_up_to_adc=input_field_location)

	# Initiate psf fitter object to calculate the strehl ratios
	psf = PSFFitter(x, y, E)

	# Create some empty lists for later
	strehl_zern_complete, strehl_psd_complete = [], []
	rms_zern_list_complete, rms_psd_list_complete = [], []


	# Repeat the analysis 50 times
	for i in range(number_of_iterations):
		
		print('Iteration {}/{}'.format(i+1,number_of_iterations))
		
		# Generate a PSD surface
		psd_surface = wfe.generate_height_map(midspat_rms, 2, 1)

		# Fit the PSD surface using zernike terms
		# Generating a normalized grid for the Zernike evaluation
		yz,xz 		= np.mgrid[-1:1:wfe.n*1j, -1:1:wfe.m*1j]
		_r 			= np.sqrt(xz**2 + yz**2)
		_mask		= _r > 1
		coeffic 	= wfe.fit_zernike_coefficients(psd_surface, modes=66)
		
		StrehlList_zern, StrehlList_psd = [], []
		rms_zern_list, rms_psd_list = [], []


		# Compare results at different RMS values
		for rms in rmsList:

			########################################
			# WFE described by Zernike polynomials #
			########################################
			# Make a new perturbed wavefront error with the required RMS
			coeffic = wfe.scale_zernike_coefficients(coeffic, rms)
			hzern 	= wfe.zernike_cartesian(coeffic, xz, yz)

			# Add the new wavefront error
			mcd.ELT_ADC.p1.add_wavefront_error(hzern)

			# Performing the propagation using Fresnel approximation to the focal plane.
			Exyz_zern,xzern,yzern = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
											print_progress=False, hx=0, hy=0,
											load_up_to_adc=input_field_location)
			Strehl_zern = np.abs(psf.get_coupling(E, Exyz_zern))**2
			StrehlList_zern.append(Strehl_zern)
			########################################


			########################################
			# WFE described by PSD profile         #
			########################################
			factor = rms / midspat_rms
			h_psd = wfe.scale_rms(factor, data=psd_surface)
			
			# Add the new wavefront error
			mcd.ELT_ADC.p1.add_wavefront_error(h_psd)

			# Performing the propagation using Fresnel approximation to the focal plane.
			Exyz_psd,xpsd,ypsd = mcd.propagate(l=l,surface=-1, zoom=False, show=False, 
											print_progress=False, hx=0, hy=0,
											load_up_to_adc=input_field_location)
			Strehl_psd = np.abs(psf.get_coupling(E, Exyz_psd))**2
			StrehlList_psd.append(Strehl_psd)
 			########################################

 			# Find the RMS value of the wavefront errros
			rms_zern 	= np.sqrt(np.mean(np.asarray(hzern[~_mask])**2))
			rms_psd 	= np.sqrt(np.mean(np.asarray(h_psd)**2))
			rms_zern_list.append(rms_zern)
			rms_psd_list.append(rms_psd)

			# Print results
			print('RMS: {:.3f}, S_zern: {:.3f}, S_psd: {:.3f}, \
				RMS_zern: {:.3f}, RMS_psd: {:.3f}'.format(
					rms*1e6, Strehl_zern, Strehl_psd, rms_zern*1e6, rms_psd*1e6))
			

		# Put the results in the correct list
		strehl_zern_complete.append(StrehlList_zern)
		strehl_psd_complete.append(StrehlList_psd)
		rms_zern_list_complete.append(rms_zern_list)
		rms_psd_list_complete.append(rms_psd_list)		

	rms_zern_list_complete = np.asarray(rms_zern_list_complete)
	rms_psd_list_complete = np.asarray(rms_psd_list_complete)

	# Import old results if they are available and append them to the results.
	rmsList_present 		= rmsList.copy()
	rmsList 				= append_to_existing_data(rmsList, './data/SR_RMS_values_obscured.pkl')
	strehl_zern_complete 	= append_to_existing_data(strehl_zern_complete, './data/SR_strehl_zern_obscured.pkl')
	strehl_psd_complete 	= append_to_existing_data(strehl_psd_complete, './data/SR_strehl_psd_obscured.pkl')
	rms_zern_list_complete 	= append_to_existing_data(rms_zern_list_complete, './data/SR_rms_zern_complete_obscured.pkl')
	rms_psd_list_complete 	= append_to_existing_data(rms_psd_list_complete, './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)

	save_variable('./data/SR_RMS_values_obscured.pkl', rmsList)
	save_variable('./data/SR_strehl_zern_obscured.pkl', strehl_zern_complete)
	save_variable('./data/SR_strehl_psd_obscured.pkl', strehl_psd_complete)
	save_variable('./data/SR_rms_zern_complete_obscured.pkl', rms_zern_list_complete)
	save_variable('./data/SR_rms_psd_complete_obscured.pkl', rms_psd_list_complete)