import numpy as np
import matplotlib.pyplot as plt
import gaussianoptics as go


lens_coupler = 12e-3


Radius = 0.05
alpha = 15
R_mirr_t = -Radius*np.cos(alpha*np.pi/180)
R_mirr_s = -Radius/np.cos(alpha*np.pi/180)
pos_colTele2= 497.5e-3

opticalbench_tangential_trans = [['path',1,0,'f'],
                             ['thinlens',12e-3,lens_coupler,'r'],
                             ['path',1,0,'f'],
                             ['thicklens',[1,1.52,13.1e-3,np.inf,11.7e-3],20e-3,'r'],
                             ['path',1,0,'f'],
                             ['thinlens',50e-3,50e-3,'r'],
                             ['path',1,0,'f'],
                             ['flatmirror',0,50e-3,'r'],
                             ['path',1,0,'f'],
                             ['curvedmirror',[0.05,alpha,0],50e-3,'r'],
                             ['path',1,0,'f'],
                             ['flatrefraction',[1,1.52],50e-3,'r'],
                             ['path',1,0,'f'],
                             ['curvedrefraction',[1.52,1,R_mirr_t],15e-3,'r'],
                             ['path',1,0,'f'],
                             ['thinlens',100e-3,50e-3,'r'],
                             ['path',1,0.3,'r']
                        ]

opticalbench_sagital_trans = [['path',1,0,'f'],
                             ['thinlens',12e-3,lens_coupler,'r'],
                             ['path',1,0,'f'],
                             ['thicklens',[1,1.52,13.1e-3,np.inf,11.7e-3],20e-3,'r'],
                             ['path',1,0,'f'],
                             ['thinlens',50e-3,50e-3,'r'],
                             ['path',1,0,'f'],
                             ['flatmirror',0,50e-3,'r'],
                             ['path',1,0,'f'],
                             ['curvedmirror',[0.05,alpha,1],50e-3,'r'],
                             ['path',1,0,'f'],
                             ['flatrefraction',[1,1.52],50e-3,'r'],
                             ['path',1,0,'f'],
                             ['curvedrefraction',[1.52,1,R_mirr_s],15e-3,'r'],
                             ['path',1,0,'f'],
                             ['thinlens',100e-3,50e-3,'r'],
                             ['path',1,0.3,'r']]
'''
opticalbench = [['path',1,0,'f'],
                ['thinlens',12e-3,lens_coupler,'r'],
                ['path',1,0,'f'],
                ['thinlens',f1_tele,pos_f1,'r'],
                ['path',1,0,'f'],
                ['thinlens',f2_tele,pos_f2,'r'],
                ['path',1,0,'f'],
                ['thinlens',f_halo1,pos_halo1,'r'],
                ['path',1,0,'f'],
                ['thinlens',f_halo2,pos_halo2,'r'],
                ['path',1,0,'f'],
                ['thinlens',f1_colTele,pos_colTele1,'r'],
                ['path',1,0,'f'],
                ['thinlens',f2_colTele,pos_colTele2,'r'],
                ['path',1,0.3,'r']
]
'''
lam = 854e-9

N =5001

#Beamwaist inside fiber 780HP
w0 = 2.5e-6
z0 = np.pi*w0**2/lam
q0 = np.complex(0,z0)


fig, ax = plt.subplots()

[x,w,R,OB,M_tot] = go.gaussianoptics(opticalbench_tangential_trans,lam,q0,0,N)
ax.plot(x*1000, w*1000, color='red')
ax.plot(x*1000, -w*1000, color='red')

[x,w,R,OB,M_tot] = go.gaussianoptics(opticalbench_sagital_trans,lam,q0,0,N)
ax.plot(x*1000, w*1000, color='blue')
ax.plot(x*1000, -w*1000, color='blue')

ax.set(xlabel='position (mm)', ylabel='beam radius (mm)', title='Testbench')

plt.grid(True,which="both")

Ymin = -10
Ymax = 10
elements = len(OB)
for i in range(elements):
    if OB[i][0] == 'flatmirror':
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color='blue',linewidth=1)
    elif OB[i][0] == 'path':
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color=[0.5,0.5,0.5],linewidth=1)
    elif OB[i][0] == 'flatrefraction':
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color='red',linewidth=1)
    elif OB[i][0] == 'thicklens':
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color='green',linewidth=1)
    elif OB[i][0] == 'thinlens':
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color=[0.75,0,1],linewidth=1)
    else:
        ax.plot([OB[i][3]*1000, OB[i][3]*1000],[Ymin/1.1, Ymax/1.1],color='black',linewidth=1)


plt.show()