#!/usr/bin/env Python3
import numpy as np
from pdb import set_trace
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from matplotlib.colors import LogNorm
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
from fourierOptics import *


'''
	@ Author: 	Joost van den Born
	@ Desc.:	Importing a grid sag file and scaling it appropriately
'''

class surface:
	def __init__(self,data):
		self.data 		= data
		self.resolution = np.shape(self.data)
		self.rms		= np.sqrt(np.mean(np.asarray(self.data)**2))
		self.ptv		= np.abs(np.max(self.data)) + np.abs(np.min(self.data))

	def scale_rms(self,factor):
		# Scale the surface data such that the RMS gets multiplied by factor
		_factor = factor ** 2
		self.calculate_2D_PSD()
		self.scaledPSD2D	= self.PSD2D * _factor
		# Fourier transform of the surface profile
		_in	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
		_in[:] 				= self.data
		_out	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
		_FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
		_FFTW_Obj()	
		_phase 				= np.angle(_out)

		_in	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
		_in[:] 				= np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.scaledPSD2D))
		_out	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
		_FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
		_FFTW_Obj()	
		self.scaledData 	= _out.real
		self.scaledrms 		= np.sqrt(np.mean(np.asarray(self.scaledData)**2))
		self.scaledptv		= np.abs(np.max(self.scaledData)) + np.abs(np.min(self.scaledData))

	def calculate_2D_PSD(self):
		''' Returns the 2D PSD map '''
		_in	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
		_in[:] 				= self.data
		_out	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
		_FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
		_FFTW_Obj()	
		self.D1 = np.fft.fftshift(_out)
		self.PSD2D = np.abs(self.D1)**2
		return self.PSD2D

	def calculate_1D_PSD(self, filteredDataFlag=False):
		''' Returns the azimuthally averaged PSD map '''
		try:
			if not filteredDataFlag:
				self.PSD1D = self.azimuthal_average()
			else:
				self.filteredPSD1D = self.azimuthal_average(filteredDataFlag=filteredDataFlag)
				return self.filteredPSD1D
		except AttributeError:
			self.calculate_2D_PSD()
			self.PSD1D = self.azimuthal_average()

		return self.PSD1D


	def set_lowpass_filter(self, r, dpix=1):
		''' Returns a map where all frequencies higher than r have been filtered out '''
		self.calculate_2D_PSD()
		# Masking out the parts that fall outside the lens clear aperture
		_dist_from_center = np.sqrt(x**2 + y**2) * dpix
		_mask = (_dist_from_center <= (r)).astype(int)
		_maskedPSD2D 	 		= self.PSD2D.copy()
		_maskedPSD2D 			*= _mask
		_phase 					= np.angle(np.fft.fft2(self.data))
		self.filteredPSD2D 		= _maskedPSD2D
		self.filteredData		= np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real

	def set_highpass_filter(self, r, dpix=1):
		'''Returns a map where all spatial frequencies lower than r have been filtered out'''
		self.calculate_2D_PSD()
		# Masking out the parts that fall outside the lens clear aperture
		x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2]
		_dist_from_center = np.sqrt(x**2 + y**2) * dpix
		_mask = (_dist_from_center <= (r)).astype(int)
		_maskedPSD2D 	 		= self.PSD2D.copy()
		_maskedPSD2D 			*= (1-_mask)
		_phase 					= np.angle(np.fft.fft2(self.data))
		self.filteredPSD2D 		= _maskedPSD2D
		self.filteredData		= np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real

	def set_bandpass_filter(self, rmin, rmax, dpix=1):
		'''Returns a map where all spatial frequencies lower than r2 and higher than r1 have been filtered out'''
		assert rmax > rmin, 'rmax < rmin, not physical!'
		self.calculate_2D_PSD()
		# Masking out the parts that fall outside the lens clear aperture
		x,y = np.ogrid[-self.resolution[0]/2:self.resolution[0]/2, -self.resolution[1]/2:self.resolution[1]/2]
		_dist_from_center = np.sqrt(x**2 + y**2) * dpix
		_mask_1 = (_dist_from_center <= (rmax)).astype(int)
		_mask_2 = (_dist_from_center >= (rmin)).astype(int)
		_maskedPSD2D 	 		= self.PSD2D.copy()
		_maskedPSD2D 			*= (_mask_1 & _mask_2)
		_phase 					= np.angle(np.fft.fft2(self.data))
		self.filteredPSD2D 		= _maskedPSD2D
		self.filteredData		= np.fft.ifft2(np.exp(1j*_phase) * np.sqrt(np.fft.fftshift(self.filteredPSD2D))).real

	def azimuthal_average(self, center=None, filteredDataFlag=False):
		"""
		Calculate the azimuthally averaged radial profile.

		image - The 2D image
		center - The [x,y] pixel coordinates used as the center. The default is 
		         None, which then uses the center of the image (including 
		         fracitonal pixels).

		"""

		# Calculate the indices from the image
		_y, _x = np.indices(self.resolution)
		if not center:
			center = np.array([(_x.max()-_x.min())/2.0, (_y.max()-_y.min())/2.0])

		_r = np.hypot(_x - center[0], _y - center[1])

		# Get sorted radii
		_ind = np.argsort(_r.flat)
		_r_sorted = _r.flat[_ind]
		if not filteredDataFlag:
			_i_sorted = self.PSD2D.flat[_ind]
		else:
			_i_sorted = self.filteredPSD2D.flat[_ind]

		# Get the integer part of the radii (bin size = 1)
		_r_int = _r_sorted.astype(int)

		# Find all pixels that fall within each radial bin.
		_deltar = _r_int[1:] - _r_int[:-1]  # Assumes all radii represented
		_rind 	= np.where(_deltar)[0]     # location of changed radius
		_nr 	= _rind[1:] - _rind[:-1]    # number of radius bin

		# Cumulative sum to figure out sums for each radius bin
		_csim = np.cumsum(_i_sorted, dtype=float)
		_tbin = _csim[_rind[1:]] - _csim[_rind[:-1]]
		_radial_prof = _tbin / _nr

		return _radial_prof



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

	# Pick a wavelength
	l = 1.35e-3

	# Initiate Fourier Optics
	fou = FourierOptics(l=l)
	
	# Pick sampling parameters
	dpix = 160
	n, m = 1001, 1001
	oC = [-.5, .5, -.5, .5]
	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.top_hat(38854/2, n, m, dpix=dpix, oversample=1)

	
	# 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))
	# oC = [x.min(), x.max(), y.min(), y.max()]

	# Add PSD wavefront error to the lens
	wfe = WavefrontError()
	h_data = wfe.generate_height_map(5e-5, 2, 1, n, m) # 50 nm RMS for standard

	PSDSurface = surface(h_data)
	# print(PSDSurface.rms)
	# PSDSurface.set_bandpass_filter(10, 20)
	# PSDSurface.calculate_1D_PSD(filteredDataFlag=False)
	# PSDSurface.calculate_1D_PSD(filteredDataFlag=True)

	# fig = plt.figure()
	# ax1 = fig.add_subplot(121)
	# ax2 = fig.add_subplot(122)
	# im1 = ax1.imshow(PSDSurface.filteredData, origin='lower', cmap='jet')
	# ax2.loglog(PSDSurface.PSD1D)
	# ax2.loglog(PSDSurface.filteredPSD1D)
	# ax2.grid(alpha=0.3, which='major')
	# ax2.grid(alpha=0.1, which='minor')
	# fig.colorbar(im1, ax=ax1)
	# plt.tight_layout()
	# plt.show()

	popt_airy_list, popt_gauss_list, popt_moffat_list, popt_cust_list = [], [], [], []
	pcov_airy_list, pcov_gauss_list, pcov_moffat_list, custom_var_list = [], [], [], []
	strehl_list = []
	Exyzlist = []
	# Scaling the surface profile by a factor f
	



	'''
	

	M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945)
	t_l = M1.phase_delay(x,y,l)

	# The input distribution gets a curved wavefront after the lens
	U_in = A * t_l

	Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)

	# Preparing some figures
	fig1 = plt.figure()
	ax = fig1.add_subplot(111)
	im = ax.imshow(np.abs(Exyz_og)**2, cmap=plt.cm.jet, origin='bottom',
	    extent=oC)
	fig1.colorbar(im)

	fig2 = plt.figure()
	ax21 = fig2.add_subplot(121)
	ax21.semilogy((np.abs(Exyz_og)**2)[500], 'k-', label='Original')
	ax21.grid(True, alpha=0.3)
	

	ax22 = fig2.add_subplot(122)
	ax22.grid(True, alpha=0.3)
	

	fig3 = plt.figure()
	ax31 = fig3.add_subplot(131)
	ax32 = fig3.add_subplot(132)
	ax33 = fig3.add_subplot(133)

	# Start of fitting routine
	psfFit = PSFFitter(xp.real,yp.real,np.abs(Exyz_og)**2)

	gaussianFit = moffatFit = airyFit = True

	# Fit a Gaussian profile
	if gaussianFit:
		gaussianFitted, popt_gauss, pcov_gauss = psfFit.fit_gaussian_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0))
		A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = popt_gauss
		ax.contour(xp.real, yp.real, gaussianFitted, 8, colors='w', linestyles='solid')
		ax.plot(x0_gauss,y0_gauss, 'wx')
		ax21.semilogy(gaussianFitted[500], 'r-', label='Gaussian fit')
		ax22.semilogy((np.abs(Exyz_og)**2)[500] - gaussianFitted[500], 'r-', label='Gaussian differential')
		im31 = ax31.imshow(np.abs(Exyz_og)**2-gaussianFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
		fig3.colorbar(im31, ax=ax31)

	if moffatFit:
		moffatFitted, popt_moffat, pcov_moffat = psfFit.fit_moffat_PSF(initialGuess=(0.03**2, 0, 0, 0.02, 0.02, 0, 0, 4))
		A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = popt_moffat
		ax.contour(xp.real, yp.real, moffatFitted, 8, colors='y', linestyles='dashed')
		ax.plot(x0_moffat,y0_moffat, 'yx')
		ax21.semilogy(moffatFitted[500], 'g-', label='Moffat fit')
		ax22.semilogy((np.abs(Exyz_og)**2)[500] - moffatFitted[500], 'g-', label='Moffat differential')
		im32 = ax32.imshow(np.abs(Exyz_og)**2-moffatFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
		fig3.colorbar(im32, ax=ax32)

	if airyFit:
		airyFitted, popt_airy, pcov_airy = psfFit.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
		amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = popt_airy
		ax.contour(xp.real, yp.real, airyFitted, 8, colors='c', linestyles='dashdot')
		ax.plot(x0_airy,y0_airy, 'cx')
		ax21.semilogy(airyFitted[500], 'b-', label='Airy fit')
		ax22.semilogy((np.abs(Exyz_og)**2)[500] - airyFitted[500], 'b-', label='Airy differential')
		im33 = ax33.imshow(np.abs(Exyz_og)**2-airyFitted, cmap=plt.cm.jet, origin='bottom', extent=oC)
		fig3.colorbar(im33, ax=ax33)
	
	ax21.legend()
	ax22.legend()
	plt.tight_layout()
	plt.show()




	'''
	

	# Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
	M1 = ParabolicMirror(dx=0, dy=0, name='ELT-M1', CA=38854/2, f=679945)
	t_l = M1.phase_delay(x,y,l)
	U_in = A * t_l
	Exyz_og,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)

	eef2, r2 = fou.encircled_energy_function(np.abs(Exyz_og)**2, xp, yp)
	plt.plot(r2, eef2); plt.grid();plt.show()
	set_trace()
	RMSlist = []
	flist = [0.01, 0.1, 0.2, 0.5, 0.75, 1, 1.5, 2, 3, 4, 5, 6]
	for f in flist:
		PSDSurface.scale_rms(f)	
		scal = PSDSurface.scaledData
		print(f, '{:.1f} nm'.format(PSDSurface.scaledrms*1e6))
		RMSlist.append(PSDSurface.scaledrms*1e6)
		M1.add_wavefront_error(scal)
		t_l = M1.phase_delay(x,y,l)

		# The input distribution gets a curved wavefront after the lens
		U_in = A * t_l

		# Performing the propagation using Fresnel approximation to the focal plane.
		Exyz_psd,xp,yp = fou.fresnel_propagation(U_in, x, y, dz=M1.f, show=False, outputCoordinates=oC)
		
		# Set up fitting object
		psfFit_psd = PSFFitter(xp.real,yp.real,np.abs(Exyz_psd)**2)

		# Fit airy pattern, Gaussian and Moffat functions
		airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.03**2, 0.0, 0.0, 0.0, 0.03))
		gaussianFitted, popt_gauss, pcov_gauss = psfFit_psd.fit_gaussian_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0))
		# moffatFitted, popt_moffat, pcov_moffat = psfFit_psd.fit_moffat_PSF(initialGuess=(0.04, 0, 0, 0.01, 0.01, 0, 0, 4))
		custFitted, popt_cust, cust_var, cust_coupling = psfFit_psd.fit_custom_field(Exyz_og, Exyz_psd, return_coupling=True)


		popt_airy_list.append(popt_airy)
		popt_gauss_list.append(popt_gauss)
		# popt_moffat_list.append(popt_moffat)
		popt_cust_list.append(popt_cust)

		pcov_airy_list.append([ np.sqrt(pcov_airy[i][i]) for i in range(len(pcov_airy)) ])
		pcov_gauss_list.append([ np.sqrt(pcov_gauss[i][i]) for i in range(len(pcov_gauss)) ])
		# pcov_moffat_list.append([ np.sqrt(pcov_moffat[i][i]) for i in range(len(pcov_moffat)) ])
		custom_var_list.append(np.sqrt(cust_var))
		strehl_list.append(np.abs(cust_coupling)**2)
		Exyzlist.append(Exyz_psd)



	# Optimum fit parameters for all scaled surface errors
	amplitude_airy, x0_airy, y0_airy, R_airy, offset_airy = np.asarray(popt_airy_list).T
	A_gauss, x0_gauss, y0_gauss, sigma_x_gauss, sigma_y_gauss, theta_gauss, offset_gauss = np.asarray(popt_gauss_list).T
	# A_moffat, x0_moffat, y0_moffat, sigma_x_moffat, sigma_y_moffat, theta_moffat, offset_moffat, beta_moffat = np.asarray(popt_moffat_list).T
	x0_cust, y0_cust = np.asarray(popt_cust_list).T # Check if this is correct!

	# Errors in the optimum fit parameters for all scaled surface errors
	amplitude_airy_err, x0_airy_err, y0_airy_err, R_airy_err, offset_airy_err = np.asarray(pcov_airy_list).T
	A_gauss_err, x0_gauss_err, y0_gauss_err, sigma_x_gauss_err, sigma_y_gauss_err, theta_gauss_err, offset_gauss_err = np.asarray(pcov_gauss_list).T
	# A_moffat_err, x0_moffat_err, y0_moffat_err, sigma_x_moffat_err, sigma_y_moffat_err, theta_moffat_err, offset_moffat_err, beta_moffat_err = np.asarray(pcov_moffat_list).T
	x0_cust_err, y0_cust_err = np.asarray(custom_var_list).T


	# Distance from original centroid at (0,0)
	r_airy = np.sqrt(x0_airy**2+y0_airy**2)
	r_err_airy = np.sqrt( (x0_airy**2*x0_airy_err**2+y0_airy**2*y0_airy_err**2) / (x0_airy**2+y0_airy**2))
	
	r_gauss = np.sqrt(x0_gauss**2+y0_gauss**2)
	r_err_gauss = np.sqrt( (x0_gauss**2*x0_gauss_err**2+y0_gauss**2*y0_gauss_err**2) / (x0_gauss**2+y0_gauss**2))

	# r_moffat = np.sqrt(x0_moffat**2+y0_moffat**2)
	# r_err_moffat = np.sqrt( (x0_moffat**2*x0_moffat_err**2+y0_moffat**2*y0_moffat_err**2) / (x0_moffat**2+y0_moffat**2))

	r_cust = np.sqrt(x0_cust**2 +  y0_cust**2)
	r_err_cust = np.sqrt( (x0_cust**2*x0_cust_err**2+y0_cust**2*y0_cust_err**2) / (x0_cust**2+y0_cust**2))

	fig = plt.figure()
	ax1 = fig.add_subplot(121)
	ax2 = fig.add_subplot(122)
	ax1.errorbar(RMSlist, r_gauss, yerr=r_err_gauss, color='green', marker='o', linestyle='--', label='Gaussian', capsize=3, markeredgewidth=1)
	ax1.errorbar(RMSlist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1)
	# ax1.errorbar(RMSlist, r_moffat, yerr=r_err_moffat, color='green', marker='o', linestyle='--', label='Moffat',  capsize=3, markeredgewidth=1)
	ax1.errorbar(RMSlist, r_cust, yerr=r_err_cust, color='orange', marker='o', linestyle='--', label='Custom', capsize=3, markeredgewidth=1)
	ax1.axhline(1.875e-4, color='red', linestyle='-', label='50 $\mu$as @ 4 mas/pix')
	ax1.set_xscale('log')
	# ax1.set_yscale('log')
	ax1.set_xlabel('RMS WFE (nm)')
	ax1.set_ylabel('Distance')
	ax1.legend()
	ax1.grid(True, alpha=0.3, which='major')
	ax1.grid(True, alpha=0.1, which='minor')

	ax2.errorbar(x0_airy, y0_airy, xerr=x0_airy_err, yerr=y0_airy_err, color='k', marker='o', linestyle='--', label='Airy')
	ax2.errorbar(x0_cust, y0_cust, xerr=x0_cust_err, yerr=y0_cust_err, color='r', marker='o', linestyle='--', label='Custom')
	ax2.set_aspect('equal')
	ax2.grid(True, alpha=0.3)
	plt.tight_layout()
	plt.show(block=False)

	fig2 = plt.figure()
	for i in range(len(flist)):
		axi = fig2.add_subplot(4,3,i+1)
		imi = axi.imshow(np.abs(Exyzlist[i])**2, extent=oC, cmap='jet')
		# fig2.colorbar(imi, ax=axi)
		axi.title.set_text('RMS: {:.1f} nm (S={:.2f})'.format(RMSlist[i], strehl_list[i]))
		axi.set_xlim([-0.2,0.2])
		axi.set_ylim([-0.2,0.2])
	plt.tight_layout()
	plt.show()