import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *

'''
	@Author:		Joost van den Born
	@Date:			14 July, 2022
	@Description:	I use this script to create a large array (1001x1001 and 3001x3001)
					that can be used for the phase screen in the WFE production.
					Because it takes very long to generate these at this high resolution,
					we will need to save them to a pickle file.

					What is saved:
					The array of coefficients
					The array of the full phase field
					
					To Do:
					- Normalize/Scale the amplitude in an appropriate manner.
'''


if __name__ == '__main__':
	# plt.rc('text', usetex=True)
	# plt.rc('font', family='serif')
	import time
	# Setup of simulation
	l = 1.49e-3
	# mcd = ParaxialMICADO(n=1001, m=1001, dpix=160)
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
	# mcd.plot_ray_trace()
	wfe = WavefrontError(mcd.n, mcd.m)


	# -------------------------------------------------------
	####### GENERATING THE PHASE WITH A FOURIER SERIES ######
    # -------------------------------------------------------
	# Get x3 and y3 (the coordinates to apply the ADC WFE to):
	# The focal plane area is about 200x200 mm, so we apply the wavefront error to an area
	# of 200x200 mm.
	xfp, yfp = np.meshgrid(np.linspace(-150, 150, mcd.n), np.linspace(-150, 150, mcd.m))
	# np.random.seed(2183)
	# coeff = np.random.uniform(-1,1,(int(mcd.n/3),int(mcd.m/3)))
	np.random.seed(2183)
	coeff = np.random.uniform(-1,1,(101,101))
	
	if False:
		print('Going in...')
		s1 = time.time()
		fs_array = fourier_series(xfp, yfp, coeff, Px=300, Py=300, xs=150, ys=150)
		s2 = time.time()
		print('Finished!')
		print('That took me {:.4f} s!'.format(s2-s1) )
		save_variable('wfe_fp_phase_array_coefficients_{}_{}.pkl'.format(mcd.n, mcd.m), coeff)
		save_variable('wfe_fp_phase_array_{}_{}.pkl'.format(mcd.n, mcd.m), fs_array)

	# -------------------------------------------------------
	# -------------------------------------------------------
	# -------------------------------------------------------

	if True:
		# coeff = load_variable('wfe_fp_phase_array_coefficients_{}_{}.pkl'.format(mcd.n, mcd.m))
		fs_array = load_variable('wfe_fp_phase_array_{}_{}.pkl'.format(mcd.n, mcd.m))

		wfe.phase = fs_array
		surf = wfe.generate_height_map(50e-6, 2, 1)
		rms = np.sqrt(np.mean(surf**2))
		print('RMS: {}'.format(rms))


	plt.figure()
	plt.imshow(fs_array, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], origin='lower')
	plt.colorbar()
	plt.show()


	plt.figure()
	plt.imshow(surf, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], origin='lower')
	plt.colorbar()
	plt.show()