import numpy as np
import matplotlib.pyplot as plt
from numpy.polynomial import polynomial
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *
from pdb import set_trace

def polyfit2d(x, y, z, deg):
	x = np.asarray(x)
	y = np.asarray(y)
	z = np.asarray(z)
	deg = np.asarray(deg)
	vander = polynomial.polyvander2d(x, y, deg)
	vander = vander.reshape((-1,vander.shape[-1])) # Wat gebeurt hier? reshape(-1) lijkt niets te doen?
	z = z.reshape((vander.shape[0],))
	c = np.linalg.lstsq(vander, z, rcond=None)[0]
	return c.reshape(deg+1)



if __name__ == '__main__':
	plt.rc('text', usetex=True)
	plt.rc('font', family='serif')

	# Initializing MICADO paraxial model
	mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
	wfe = WavefrontError(mcd.n, mcd.m)
	
	# Setup
	l = 1.49e-3

	# Set field points
	hxList = [0]
	hyList = [0]
	
	wfe_p1 = load_variable('../data/compressed items/ADC_tWFE_p1_zern_20nm_psd_15nm.pkl')
	wfe_p2 = load_variable('../data/compressed items/ADC_tWFE_p2_zern_20nm_psd_15nm.pkl')
	wfe_p3 = load_variable('../data/compressed items/ADC_tWFE_p3_zern_20nm_psd_15nm.pkl')
	wfe_p4 = load_variable('../data/compressed items/ADC_tWFE_p4_zern_20nm_psd_15nm.pkl')
	
	# wfe_p1 = load_variable('../data/compressed items/ADC_tWFE_p1_psd_25nm.pkl')
	# wfe_p2 = load_variable('../data/compressed items/ADC_tWFE_p2_psd_25nm.pkl')
	# wfe_p3 = load_variable('../data/compressed items/ADC_tWFE_p3_psd_25nm.pkl')
	# wfe_p4 = load_variable('../data/compressed items/ADC_tWFE_p4_psd_25nm.pkl')



	prop_surf = 4
	E, x, y = mcd.propagate(l=l,surface=prop_surf, zoom=False, show=False, 
							print_progress=True, hx=0, hy=0)
	# Save the resulting 'perfect' PSF for later
	tx_combined = 0
	ty_combined = 0

	print(mcd.ELT_ADC.p1.CA)

	for wfe_surf in [wfe_p1, wfe_p2, wfe_p3, wfe_p4]:
		
		# Don't show anything outside the clear aperture.
		mask = np.sqrt(x**2 + y**2) > 40#mcd.ELT_ADC.p1.CA/2
		# Plot P1
		phasescreenMasked = wfe_surf.copy()
		wfe_p1[mask] = np.nan
		phasescreenMasked[mask] = 0

		# Bepaling van de gemiddelde tilt voor deze coordinaten, i.e. een lineaire fit.
		pf1 = polyfit2d(x.real[~mask], y.real[~mask], phasescreenMasked[~mask], [1,1])        

		# Coefficienten van de fit vertaalt naar graden tilt
		o1x = np.arctan(pf1[1,0])
		o1y = np.arctan(pf1[0,1])
		print(o1x, o1y)
		o1x_deg = np.degrees(o1x)
		
		o1y_deg = np.degrees(o1y)
		tx_combined += o1x
		ty_combined += o1y
		print('tx: {} rad ({} deg = {:.2f} arcsec)'.format(o1x, o1x_deg, o1x_deg*3600))
		print('ty: {} rad ({} deg = {:.2f} arcsec)'.format(o1y, o1y_deg, o1y_deg*3600))

	print('Combined results:\n-----------------')
	print('tx: {} rad --> distortion of {} μas'.format(tx_combined, 1601.688*tx_combined*0.267*1e6))
	print('ty: {} rad --> distortion of {} μas'.format(ty_combined, 1601.688*ty_combined*0.267*1e6))
	tr_combined = np.sqrt(tx_combined**2 + ty_combined**2)
	print('tr: {} rad --> distortion of {} μas'.format(tr_combined, 1601.688*tr_combined*0.267*1e6))