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


'''
	Description:
	This script simulates the dispersed PSF of the small (40 cm)
	Gratama telescope at the University of Groningen under 1.5
	arcsecond 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))