import numpy as np
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Analyses/Python/Analytical/Fourier Optics')
from fourierOptics import *
import scipy.special as spc
from pdb import set_trace
'''
def find_pixel_scale(x):
	""" Calculate the pixel size from the x, y plane """
	_dpix = np.mean(np.diff(x.real))
	if _dpix < 1e-99:
		_dpix = np.mean(np.diff((x.T).real))
	return _dpix

def seeing(wvl, r0):
	"""
		Seeing value in arcsec at wavelength wvl
		wvl in millimeter
		r0 in meter
	"""
	wvl *= 1e-3 # convert wvl from mm to meter
	return 3600*180/np.pi*0.976*wvl/r0

def fried_r0(wvl, seeing):
	"""
		Returns Fried's parameter value in meters
		wvl in millimeter
	"""
	wvl *= 1e-3 # convert wvl from mm to meter
	return 0.976*wvl/seeing*206264.8

def airmass(z):
	"""
		Converts zenith angle in degrees to airmass
	"""
	return 1/np.cos(np.deg2rad(z))

def spectrum(k, r0, L0=25):        
	"""SPECTRUM Phase power spectrum density computes the phase power
		spectrum density from the spatial frequency f and an
		atmosphere object
		
		Normal value for L0 = 25.
	"""                        
	cte = (24*spc.gamma(6/5)/5)**(5/6)*(spc.gamma(11/6)**2./(2*np.pi**(11/3)))
	return r0**(-5/3)*cte*(k**2 + 1/L0**2)**(-11/6)

def add_seeing(wvl, sng, field_tuple):
	"""
		Computes the seeing limited PSF for a supplied amplitude
		distribution.
		
		Parameters:
		wvl - wavelength in mm
		sng - seeing in arcseconds
		field_tuple - tuple containing (amplitude array, x array, y array)
		
		Returns:
		psf_seeing

	"""
	r01   = fried_r0(wvl, seeing=sng) * 1e3
	A_psf, x_psf, y_psf = field_tuple
	n,m = np.shape(x_psf)
	fx = x_psf / (3200*wvl)
	fy = y_psf / (3200*wvl)
	dk = find_pixel_scale(fx)

	# von-Karman spectrum of the atmosphere for the given frequencies in k-space
	psd   = spectrum(np.hypot(fx.real, fy.real), r0=r01, L0=25*1e3)
	psd[int(n/2),int(m/2)] = 0

	# Compute the atmospheric OTF
	# According to Roddier 1981, the atmospheric OTF can be derived from the phase structure function
	# assuming that the phase is stationary accross the pupil and behaves as a Gaussian variable
	Bphi = np.fft.fft2(np.fft.ifftshift(psd))*dk**2
	Dphi = np.fft.ifftshift((2 * (Bphi[0,0] - Bphi))) # ifftshift in Bphi en Dphi!
	otf_atm = np.exp(-0.5*Dphi)
	
	# Compute the OTF of the supplied PSF
	# We can derive the OTF from the inverse Fourier transform of the PSF
	psf = np.abs(A_psf)**2
	otf_psf = np.fft.fftshift(np.fft.ifft2(np.fft.fftshift(psf)))
	otf_psf = otf_psf/otf_psf.max()

	# Compute the seeing-limited PSF
	# Multiplication of OTF and then Fourier transform
	otf  = otf_psf * otf_atm
	psf_seeing = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(otf)))
	return psf_seeing
'''



if __name__ == '__main__':
	l  = 0.5e-3      # mm
	D_GT = 400          # mm
	dpix = 1.6
	n,m  = 1001, 1001
	x,y  = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))
	wfe = WavefrontError(n, m)

	# Create the Bahtinov mask and set it on the entrance pupil (ELT-M1)
	GT_MASK = Profiles().bahtinov_mask(D_GT/2, 186.53/2, n, m, dpix=1.6, N=30, oversample=1).astype('complex128')
	GT_M1 = ParabolicMirror(dx=0, dy=0, name='GT_Paraxial', CA=400, f=3200)
	GT_platescale = (3600*180/np.pi) / (GT_M1.CA * 8)

	# Set up
	Elist = []
	lList = np.linspace(0.398e-3, 0.727e-3, 25).tolist()

	# Dispersion over the band
	dispMax = 1.5 / 3600 		# maximum dispersion in degrees
	dispersion = np.linspace(0, dispMax, len(lList)).tolist()

	# Seeing conditions
	sng = 1.5


	# Iterate over the wavelengths
	for i in range(len(lList)):
		
		l = lList[i]
		print(l)
		# Generating the tilt to the input field.
		d = dispersion[i]
		field_tilt = wfe.generate_tilt(x,y,tx=0, ty=d)
		
		# Adding the tilt to the input field
		Ein = GT_MASK * np.exp(-1j * 2*np.pi/l * field_tilt) * GT_M1.phase_delay(x,y,l, paraxial=True)

		# Propagate the input field throug the paraxial model
		# Output diffraction limited PSF
		fou = FourierOptics(l)
		oC = [-0.6, 0.6, -0.6, 0.6]
		Eout, xout, yout = fou.fresnel_propagation(Ein, x, y, dz=GT_M1.f, show=False, 
			outputCoordinates=oC, apply_interpolation=True)
		psf_GT_DL = (Eout, xout, yout)

		psf_seeing = add_seeing(l, sng, psf_GT_DL)

		Elist.append(psf_seeing)
		# Elist.append(Eout)

	# Average over all the wavelengths to get the final field.
	Eca = np.asarray([np.abs(x) for x in Elist]).mean(axis=0) 	# Amplitude 												# Intensity
	Ecp = np.asarray([np.angle(x) for x in Elist]).mean(axis=0) # Phase

	Eca = Eca/Eca.max() * 64000

	# Adding background level and nosie, estimated from FITS data of the Gratama Telescope
	Eca = Eca + np.random.normal(2000, 13.5, np.shape(Eca))

	# Convert the scale to sky-angles using the plate scale
	x = xout * GT_platescale
	y = yout * GT_platescale
	
	# save_variable('GT_psf_with_dispersion.pkl', (Eca, x, y))

	# Plotting the results
	fig = plt.figure(figsize=(9,7))
	ax = fig.add_subplot(111)
	im = ax.imshow(Eca, cmap='hot', origin='lower', extent=[x.min(),x.max(),y.min(),y.max()], norm=LogNorm())
	ax.set_xlim([x.min(),x.max()])
	ax.set_ylim([y.min(),y.max()])
	ax.set_xlabel('x [arcsec]')
	ax.set_ylabel('y [arcsec]')
	fig.colorbar(im, ax=ax, label='Amplitude')
	plt.tight_layout()
	plt.show()

	plot_real_valued_results(np.sqrt(Eca), x, y, fig_size=(9,7))