#!/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
from scipy import interpolate
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *

'''
	Description:
	This early script simulates a dispersed PSF. By rotating the 
	ADC prisms we can inspect the change in the dispersion. 
	The theoretical best position of the ADC did not result in 
	an undispersed PSF, for some reason.
'''


def perform_padding(A, xs, ys):
	assert A.shape[0] <= xs and A.shape[1] <= ys, 'New size is smaller than array size.'
	x_margin = int((xs - A.shape[0])/2)
	y_margin = int((ys - A.shape[1])/2)
	return np.pad(A, (x_margin, y_margin))

def find_nearest(array, value):
	array = np.asarray(array)
	diff = (np.abs(array - value)).argmin()
	idx = np.unravel_index(diff, array.shape)
	return idx

def find_maximum(array):
	idx = np.unravel_index(array.argmax(), array.shape)
	return idx

def select_subsection(x, y , outputCoordinates):
	select_1 = find_nearest(x, outputCoordinates[0])[1]
	select_2 = find_nearest(x, outputCoordinates[1])[1]
	select_3 = find_nearest(y, outputCoordinates[2])[0]
	select_4 = find_nearest(y, outputCoordinates[3])[0]
	return [select_1, select_2, select_3, select_4]

def perform_extrapolation(x, y, n, m):
	# Determine properties of x,y meshgrid
	xshape = x.shape[0]
	yshape = y.shape[0]
	dpix = (np.abs(np.min(x)) + np.abs(np.max(x))) / (xshape-1)

	# Extrapolate from, and include, the old x,y meshgrid
	xnew = np.linspace(-(n-1)/2*dpix, (n-1)/2*dpix, n)
	ynew = np.linspace(-(m-1)/2*dpix, (m-1)/2*dpix, m)
	xnew, ynew = np.meshgrid(xnew, ynew)
	return xnew, ynew

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


	# Dispersion levels:
	lList = np.linspace(1.49e-3, 1.78e-3, 2).tolist()
	dispMax = 17.9e-3 / 3600 # dispersion in degrees
	dispersion = np.linspace(0, dispMax, len(lList)).tolist()
	Exyz_list = []

	# Input parameters:
	dpix = 50
	n, m = 3001, 3001
	x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))
	print ('Your resolution is {}x{} with a pixel size of {}'.format(n,m,dpix))
	print('This means your grid is {} mm by {} mm, with each pixel {} mm.'.format(x.max()-x.min(), y.max()-y.min(), dpix))

	###########################################################
	#################### SETTING UP MICADO ####################
	###########################################################

	# Define the ELT mask (the spiders and central obscuration)
	prof 		= Profiles()
	# ELT_MASK 	= prof.top_hat(19427, n,m,dpix)#.astype('float64')
	# ELT_MASK = prof.obscured_top_hat(4000, 500, n, m, dpix)
	ELT_MASK 	= prof.ELT_pupil(n,m,dpix)
	
	# plt.imshow(ELT_MASK);plt.colorbar();plt.show()



	

	

	###########################################################

	for l in lList:
		i = lList.index(l)
		print('Wavelength is {:.6f} mm. ({:02d}/{}).'.format(l,lList.index(l)+1, len(lList)))
		x, y = np.meshgrid(np.linspace(-n/2*dpix, n/2*dpix, n), np.linspace(-m/2*dpix, m/2*dpix, m))

		fou = FourierOptics(l=l)
		wfe = WavefrontError(n,m)
		dispersion_tilt = wfe.generate_tilt(x, y, tx=0, ty=dispersion[i])

		# ELT Mirror system combined into a single lens
		ELT_MIRROR 	= ParabolicMirror(dx=0, dy=0, name='ELT-paraxial', CA=38542, f=6.84022e5)
		
		ELT_FP		= OpticalSurface(name='MCD-IFP')

		ELT_COL		= ParaxialLens(CA=154.54, f=1418.94, name='MCD-COL')

		ELT_ADC 	= ADCObject()

		optAngle = ELT_ADC.find_optimum_angle(17.9, lList[0], lList[-1])
		print(optAngle)
		print('ADC angle now set to {:.4f}'.format(ELT_ADC.theta))
		# ELT_ADC.set_orientation(180)
		ELT_ADC.set_orientation(-optAngle+90)

		ELT_CAM		= ParaxialLens(CA=89.2, f=4278.32701, name='MCD-CAM')

		ELT_DET		= OpticalSurface(name='MCD-DET')

		MCD 		= OpticalTrain()
		MCD.add_object(ELT_MIRROR, 	dist=0)
		MCD.add_object(ELT_FP, 		dist=ELT_MIRROR.f)
		MCD.add_object(ELT_COL, 	dist=ELT_COL.f)
		ELT_ADC.add_to_optical_train(MCD, 1209.4718)
		MCD.add_object(ELT_CAM, 	dist=50)
		MCD.add_object(ELT_DET, 	dist=ELT_CAM.f)

		# Print lens data for verification
		# MCD.print_lens_data()

		###########################################################
		###################### PROPAGATION ########################
		###########################################################
		
		coordinatesList = [
							[-1.5, 1.5, -1.5, 1.5],
							[-200, 200, -200, 200],
							[-120, 120, -120, 120],
							[-120, 120, -120, 120],
							[-120, 120, -120, 120],
							[-120, 120, -120, 120],
							[-120, 120, -120, 120],
							[-1.5, 1.5, -1.5, 1.5]
							]

		xin, yin = x, y
		
		print('Starting propagation...')
		
		# From ELT-M1 to ELT-FP
		print('Propagating from ELT-M1 to ELT-IFP', end='\r')
		wf = ELT_MASK.astype('complex128')
		ELT_MIRROR.add_wavefront_error(dispersion_tilt)
		wf *= ELT_MIRROR.phase_delay(x,y,l, paraxial=True)
		Exyz, x, y = fou.fresnel_propagation(wf, xin, yin, dz=ELT_MIRROR.f, oversample=None, 
			show=False, save=False)#, outputCoordinates=coordinatesList[0], apply_interpolation=False)
		# save_variable('./data/ELT_FP_{}.pkl'.format(l*1000), [Exyz, x, y])

		# From ELT-FP to MCD-COL
		print('Propagating from ELT-IFP to MCD-COL', end='\r')
		Exyz2, x2, y2 = fou.fresnel_propagation(Exyz, x, y, dz=ELT_COL.f, oversample=None, 
			show=False, save=False, outputCoordinates=coordinatesList[1], apply_interpolation=False)

		# From MCD-COL to MCD-ADC-P1
		print('Propagating from MCD-COL to MCD-ADC-P1', end='\r')
		Exyz2 *= ELT_COL.phase_delay(x2,y2,l, paraxial=True)
		Exyz3, x3, y3 = fou.plane_wave_propagation(Exyz2, x2, y2, dz=MCD.sequentialDistance[3], oversample=None,
			show=False, save=False, outputCoordinates=None)

		# From MCD-ADC-P1 to MCD-ADC-P2
		print('Propagating from MCD-ADC-P1 to MCD-ADC-P2', end='\r')
		Exyz3 *= ELT_ADC.get_phase_delay(x3,y3,l, theta=None, paraxial=True, prism=1)
		Exyz4, x4, y4 = fou.plane_wave_propagation(Exyz3, x3, y3, dz=MCD.sequentialDistance[4], oversample=None,
			show=False, save=False, outputCoordinates=None)	

		# From MCD-ADC-P2 to MCD-ADC-P3
		print('Propagating from MCD-ADC-P2 to MCD-ADC-P3', end='\r')
		Exyz4 *= ELT_ADC.get_phase_delay(x4,y4,l, theta=None, paraxial=True, prism=2)
		Exyz5, x5, y5 = fou.plane_wave_propagation(Exyz4, x4, y4, dz=MCD.sequentialDistance[5], oversample=None,
			show=False, save=False, outputCoordinates=None)

		# From MCD-ADC-P3 to MCD-ADC-P4
		print('Propagating from MCD-ADC-P3 to MCD-ADC-P4', end='\r')
		Exyz5 *= ELT_ADC.get_phase_delay(x5,y5,l, theta=None, paraxial=True, prism=3)
		Exyz6, x6, y6 = fou.plane_wave_propagation(Exyz5, x5, y5, dz=MCD.sequentialDistance[6], oversample=None,
			show=False, save=False, outputCoordinates=None)

		# From MCD-ADC-P4 to MCD-CAM
		print('Propagating from MCD-ADC-P4 to MCD-CAM', end='\r')
		Exyz6 *= ELT_ADC.get_phase_delay(x6,y6,l, theta=None, paraxial=True, prism=4)
		Exyz7, x7, y7 = fou.plane_wave_propagation(Exyz6, x6, y6, dz=MCD.sequentialDistance[7], oversample=None,
			show=False, save=False, outputCoordinates=None)
		
		# From MCD-CAM to MCD-DET
		print('Propagating from MCD-CAM to MCD-DET', end='\r')
		Exyz7 *= ELT_CAM.phase_delay(x7,y7,l, paraxial=True)
		Exyz8, x8, y8 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None, 
			show=False, save=False, outputCoordinates=None, apply_interpolation=False)

		# The result here shifts the PSF to the side, due to the phase tilts of the ADC
		# surfaces. We find the locations of the maximum and redo the the fresnel propagation
		# around these coordinates.
		# print('Zooming in on the PSF location...', end='\r')
		# imax_1, imax_2 = find_maximum(np.abs(Exyz8))
		# x8max = x8[imax_1, imax_2]
		# y8max = y8[imax_1, imax_2]
		# newCoordinates = np.asarray(coordinatesList[7]) + np.array([x8max, x8max, y8max, y8max])
		# newCoordinates = newCoordinates.real.tolist()

		# Exyz9, x9, y9 = fou.fresnel_propagation(Exyz7, x7, y7, dz=ELT_CAM.f, oversample=None, 
		# 	show=False, save=False, outputCoordinates=newCoordinates, apply_interpolation=True)
		print('Propagation finished!')


		Exyz_list.append(np.abs(Exyz8)**2)
		fou = None
		del fou
		wfe = None
		del wfe
		MCD = None
		del MCD
		wf = None
		del wf
		ELT_MIRROR = None
		del ELT_MIRROR
		ELT_COL = None
		del ELT_COL
		ELT_ADC = None
		del ELT_ADC
		ELT_CAM = None
		del ELT_CAM
		ELT_DET = None
		del ELT_DET
		###########################################################


	Exyz_total = np.asarray(Exyz_list).sum(axis=0)

	_xip = x8
	_yip = y8
	# save = 'ELT-MCD-Hband_dispersion_corrected.pdf'
	_fig = plt.figure()
	_ax = _fig.add_subplot(111)
	_im = _ax.imshow(Exyz_total.T, cmap='hot', 
		extent=[_xip.min(), _xip.max(), _yip.min(), _yip.max()], norm=LogNorm(vmin=2e-7))
	_ax.set_xlim([_xip.min(), _xip.max()])
	_ax.set_ylim([_yip.min(), _yip.max()])
	_ax.set_xlabel('x [mm]')
	_ax.set_ylabel('y [mm]')
	_fig.colorbar(_im, ax=_ax, label='Intensity')
	plt.tight_layout()
	# plt.savefig(save, bbox_inches='tight', transparent=True)
	plt.show()