#!/usr/bin/env Python3
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from matplotlib.colors import LogNorm
import time
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *

'''
	Description:
	This script looks how the PSF behaves for an annular lens.
	Specifically at 30 microns. After plotting the PSF at the 
	focal length, a cross section of the intensity profile
	at different distances from the lens is shown.
'''

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

	# Start time counter
	start = time.time()

	# Pick a wavelength
	l = 30e-3

	# Initiate Fourier Optics
	fou = FourierOptics(l=l)
	
	# Pick sampling parameters
	dpix = 1/50
	n, m = 2001, 2001
	print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))
	
	# Initiate the profiles class and pick a profile (tophat, obscured tophat or spider)
	prof = Profiles()
	A = prof.obscured_top_hat(2.5, 2.3, n, m, dpix=dpix)
	plt.imshow(A); plt.show()

	
	# Create an input grid
	x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))
	
	# Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
	L1 = SphericalLens('N-BK7', R1=50, R2=1e99, thickness=1, CA=8, name='L1')
	print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f))

	# Calculate the phase delay introduced by the lens
	t_l = L1.phase_delay(x,y,l, paraxial=True)
	U_in = A * t_l

	# Performing the propagation using Fresnel approximation to the focal plane.
	oC = [x.min(), x.max(), y.min(), y.max()]
	savename = False
	Exyz_og,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=True, save=savename, outputCoordinates=oC)

	# Plotting the Modulation Transfer Function of this lens system
	E1, fx1, fy1 = fou.modulation_transfer_function(np.abs(Exyz_og)**2, x,y,dpix)
	plt.figure()
	plt.plot(fx1[1000:], E1[1000, 1000:],'b-')
	plt.xlim(0,1.67)
	plt.ylim(0,1)
	plt.xlabel('Spatial frequency (cycles/aperture)')
	plt.ylabel('MTF')
	plt.grid()
	plt.show()
	



	# The following produces a cross-section of the propagation for various z, so that you can see how the beam evolves.
	if True:
		Ics 	= []
		Ics_x 	= []
		Isum 	= []
		xmax = []
		zList = np.linspace(10,150,20).tolist()
		xin, yin = x, y
		
		for i in zList:
			title = f"{zList.index(i)+1:03d}"
			print (' Calculating {0}/{1}'.format(title, f"{len(zList):03d}"), end="\r")

			# Propagate the wavefront
			# Exyz, x, y = fou.planeWavePropagationFFTW(A, dz=i*l, dpix=dpix, oversample=1, show=False, save=False)#, scaleInputArray=(n,m))
			Exyz, x, y = fou.fresnel_propagation(U_in, xin, yin, dz=i, show=False, save=False, outputCoordinates=oC)
			
			# Process the results by taking a slice of the data
			Isum.append(np.sum(np.abs(Exyz)**2))
			Ics.append((np.abs(Exyz[int(Exyz.shape[0]/2)])**2))
			Ics_x.append(x.real[0])	
			y_ext = [y.real[0,0],y.real[-1,0]]
			
			# Reset the data
			x = None
			y = None
		
		# Plotting all slices
		fig = plt.figure(figsize=(8,4))
		ax  = fig.add_subplot(111)
		ax.tick_params(axis ='both', direction='in') 

		# Plotting the square root of the amplitude for better contrast
		im = ax.imshow(np.sqrt(np.asarray(Ics)).T, cmap='hot', aspect='auto', interpolation='bicubic', extent=[zList[0],zList[-1],y_ext[0],y_ext[1]])
		ax.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
		
		'''
		# Plotting the intensity relative to the input
		ax2 = fig.add_subplot(312)
		ax2.plot(zList, np.asarray(Ics)[:,int(np.asarray(Ics).shape[1]/2)], 'k-')
		ax2.yaxis.set_major_formatter(FormatStrFormatter('%.1E'))
		ax2.set_xlabel('Propagation distance (mm)')
		ax2.tick_params(axis ='both', direction='in') 
		ax2.set_ylabel('Amplitude')
		ax2.set_xlim([np.min(zList), np.max(zList)])
		ax2.grid(True, alpha=0.3)

		ax3 = fig.add_subplot(313)
		ax3.grid(True, alpha=0.3)
		ax3.set_xlabel('Propagation distance (mm)')
		ax3.tick_params(axis ='both', direction='in') 
		ax3.set_ylabel('Normalized power')
		ax3.set_xlim([np.min(zList), np.max(zList)])
		ax3.plot(zList, np.asarray(Isum))
		'''
		# Saving the figure? Then show.
		plt.tight_layout()
		# plt.savefig('fresnel_cross.pdf', bbox_inches='tight')
		# plt.savefig('fresnel_cross.png', bbox_inches='tight')
		# plt.savefig('simple_lens_cross.pdf', bbox_inches='tight', transparent=True)
		plt.show()


	# Closing down the script and stopping the timer.
	end = time.time()
	print('The program took {:.3f} seconds!'.format((end-start)))