#!/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(self,factor):
		# Scale the surface data such that the RMS gets multiplied by factor
		_factor = factor ** 2
		self.PSD2Dmethod()
		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 PSD2Dmethod(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


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

	# Pick a wavelength
	l = 30e-3

	# Initiate Fourier Optics
	fou = FourierOptics(l=l)
	
	# Pick sampling parameters
	dpix = 1/50
	n, m = 1001, 1001
	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(2.5, 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(1e-2, 2, 1, n, m)

	PSDSurface = surface(h_data)
	print(PSDSurface.rms)


	popt_airy_list, popt_gauss_list, popt_moffat_list = [], [], []
	pcov_airy_list, pcov_gauss_list, pcov_moffat_list = [], [], []

	Exyzlist = []
	# Scaling the surface profile by a factor f
	flist = [0.01, 0.1, 0.2, 0.5, 0.75, 1, 1.5, 2, 2.5]
	for f in flist:
		PSDSurface.scale(f)	
		scal = PSDSurface.scaledData
		print(f, PSDSurface.scaledrms)


		# Create a lens with a focal length ~100 mm, with a clear aperture radius of 4 mm and thickness of 1 mm.
		L1 = LensObject('N-BK7', 1, 4, name='L1')
		L1.set_first_surface('EVEN ASPHERE', R=11.65, k=-1.10, a=[0.0, 3.69067205091E-5, -1.285461192617E-8, -1.400167659754E-10, 
					-2.51311664812E-13, 5.01789877647E-16, 5.855871513042E-18, -1.127794407419E-20])
		L1.set_second_surface('SPHERICAL', R=1e99)
		L1.add_wavefront_error(scal)
		t_l = L1.phase_delay(x,y,l, paraxial=False)
		# print('Focal length of lens {} is {:.1f} mm.'.format(L1.name, L1.f))

		# 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.
		savename = None
		Exyz_psd,_,_ = fou.fresnel_propagation(U_in, x, y, dz=L1.f, show=False, save=savename, outputCoordinates=oC)

		# Fit airy pattern
		psfFit_psd = PSFFitter(x,y,np.abs(Exyz_psd)**2)
		airyFitted, popt_airy, pcov_airy = psfFit_psd.fit_airy_PSF(initialGuess=(0.01, 0.0, 0.0, 0.0, 0.5))
		gaussianFitted, popt_gauss, pcov_gauss = psfFit_psd.fit_gaussian_PSF(initialGuess=(0.01, 0, 0, 0.07, 0.07, 0, 0))
		moffatFitted, popt_moffat, pcov_moffat = psfFit_psd.fit_moffat_PSF(initialGuess=(0, 0, 0, 0.01, 0.01, 0, 0, 1))

		popt_airy_list.append(popt_airy)
		popt_gauss_list.append(popt_gauss)
		popt_moffat_list.append(popt_moffat)

		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)) ])
		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
	
	# 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

	# 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))

	set_trace()

	fig = plt.figure()
	ax1 = fig.add_subplot(121)
	ax2 = fig.add_subplot(122)
	ax1.errorbar(flist, r_airy, yerr=r_err_airy, color='blue', marker='o', linestyle='--', label='Airy', capsize=3, markeredgewidth=1)
	ax1.errorbar(flist, r_gauss, yerr=r_err_gauss, color='red', marker='o', linestyle='--', label='Gaussian', capsize=3, markeredgewidth=1)
	ax1.errorbar(flist, r_moffat, yerr=r_err_moffat, color='green', marker='o', linestyle='--', label='Moffat',  capsize=3, markeredgewidth=1)
	ax1.set_xscale('log')
	ax1.set_xlabel('Scaling factor')
	ax1.set_ylabel('Distance')
	ax1.legend()
	ax1.grid(True, alpha=0.3)

	ax2.errorbar(x0_airy, y0_airy, xerr=x0_airy_err, yerr=y0_airy_err, color='k', marker='o', linestyle='--', label='Airy')
	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('Scaling factor: {}'.format(flist[i]))
		axi.set_xlim([-2.5,2.5])
		axi.set_ylim([-2.5,2.5])
	plt.show()