#%%
import finesse.thermal.hello_vinet as hv
import matplotlib.pyplot as plt
import numpy as np
import finesse
from finesse.knm import Map
from finesse.materials import Material,FusedSilica
from QCL import *
from finesse.utilities.maps import circular_aperture

#define functions to get full thermal lens from TE and TR
def Topl (w,a = 25.4e-3/2,h = 2e-3, material  = FusedSilica):

    #output TR +TE OPL change from HV formulation 

    r = np.linspace(-a, a, 1000)

    U_s_bulk = hv.surface_deformation_substrate_heating_HG00(r, a, h, w, material,s_max=200) *(material.nr-1)  # last part turns distance into OPL 
    _, Z_Bulk_per_W, = hv.thermal_lenses_HG00(r, a, h, w, material,s_max=200)
    return r, U_s_bulk+Z_Bulk_per_W
def Topl_coating (w,a = 25.4e-3/2,h = 5e-3, material  = FusedSilica):

    #output TR +TE OPL change from HV formulation 

    r = np.linspace(-a, a, 1000)

    U_s_coat = hv.surface_deformation_coating_heating_HG00(r, a, h, w, material,s_max=200) *(material.nr-1)  # last part turns distance into OPL 
    Z_coat_per_W, _, = hv.thermal_lenses_HG00(r, a, h, w, material,s_max=200)
    return r, U_s_coat+Z_coat_per_W

# %%
#parameters 
waist = 300e-6
rad_dim =25.4e-3/2
width_dim =1e-3 #thickness of window 
P_ab =2 # absorbed pwoer (assuming all)

## calculate thermal lens, assuming bulk absorption of 2W
r,OPD = Topl(waist,a = rad_dim, h=width_dim,material =FusedSilica)
#fit RC usinf finesse
rx,ry = np.meshgrid(r,r)
oplxy = np.interp(np.sqrt(rx.flatten()**2+ry.flatten()**2),r,OPD).reshape((len(r),len(r)))
test_map =  Map(r,r,opd = oplxy*P_ab, amplitude=circular_aperture(r, r, rad_dim))
Rc = test_map.get_radius_of_curvature_reflection(waist)
nr = 1.3628 #fused silica nr at 4.6um
f = 1/(-2*Rc[0]*(nr-1))

print(f' finesse {Rc} m')
plt.figure(figsize=[6,4])
plt.plot(r*1e3,(OPD-OPD.max())*1e6*2,c= 'cornflowerblue', label="hv OPD")  #everything is normalised for 1 W of absorbed power, HV assumes P_tot>> P_abs so its linear 
plt.plot(r*1e3,r**2*1/(2*Rc[0])*1e6,'--',c='darkblue', label ='fitted Rc')
plt.xlabel("r [mm]")
plt.ylabel("OPD [um]")
plt.grid(alpha = 0.3)
plt.legend()
plt.ylim(-10,2)
plt.title(f'Hello Vinet modelled thermal lens in Fused Silica \n f ={f:.4f} m \n Rc = {Rc[0]:.4f} m ')
plt.savefig('Thermal_len_in_SiO2_bulk', dpi=150, bbox_inches='tight')

M_lens = np.matrix([[1, 0],[-1/f,1]])
# M_lens = M_lens@M_lens
#%%
#propogate thorugh system to ITM 
QCLs = ['0918', '0920', '0919', '0851']
for qcln in QCLs:
    print(qcln)
    print('------------')
    QCL = QCLProfile(qcln, no_plot=True)
    model,_ = QCL.load_model()
    qinx =  model.QCLBeam.qx
    qiny =  model.QCLBeam.qy
    ITM_sol = model.propagate_beam_astig(
                from_node = model.QCL.fr.o, 
                to_node = model.ITM.fr1.i,
            )
    print('current system')
    print('q ;',model.QCLBeam.qx.q ,model.QCLBeam.qy.q )
    ITM_w =np.array([ITM_sol.qx('ITM_HR.p1.i').w ,ITM_sol.qy('ITM_HR.p1.i').w])
    ITM_q =np.array([ITM_sol.qx('ITM_HR.p1.i').q ,ITM_sol.qy('ITM_HR.p1.i').q])

    print('w at ITM', ITM_w, 'q at ITM', ITM_q)
    #add lens
    qoutx= (M_lens[0,0]*qinx+M_lens[0,1])/(M_lens[1,0]*qinx+M_lens[1,1])
    qouty= (M_lens[0,0]*qiny+M_lens[0,1])/(M_lens[1,0]*qiny+M_lens[1,1])

    modelnew_q,_ = QCL.load_model()
    modelnew_q.QCLBeam.qx = qoutx
    modelnew_q.QCLBeam.qy = qouty

    ITM_sol_new_q = modelnew_q.propagate_beam_astig(
                from_node = modelnew_q.QCL.fr.o, 
                to_node = modelnew_q.ITM.fr1.i,
            )
    print('with lens')
    print('q ;',modelnew_q.QCLBeam.qx.q ,modelnew_q.QCLBeam.qy.q )
    ITM_w_new_q =np.array([ITM_sol_new_q.qx('ITM_HR.p1.i').w ,ITM_sol_new_q.qy('ITM_HR.p1.i').w])
    ITM_q_new_q =np.array([ITM_sol_new_q.qx('ITM_HR.p1.i').q ,ITM_sol_new_q.qy('ITM_HR.p1.i').q])
    print('w at ITM', ITM_w_new_q, 'q at ITM', ITM_q_new_q)

# %% checked power at each point 
# from thorlabs 2% transmission at 4.6 through fused silica windows. 
#converting to absorption coefficent/ m  
# transmission from https://www.crystran.com/optical-materials/fused-silica-sio2/
fusedsilicadb_km  = np.log10(0.02)*10/6e-3*1e3
fusedsilicappm_cm = fusedsilicadb_km*2.3
fusedsilicam = fusedsilicappm_cm*1e-4
thickness = 1e-3
print( 'Power after first window', np.exp(fusedsilicam*thickness/np.cos(np.pi/4)))

print( 'Power after second window', np.exp(fusedsilicam*thickness/np.cos(np.pi/4))**2)
print( 'Power reflect back through both windows ', np.exp(fusedsilicam*thickness/np.cos(np.pi/4))**4)





# %%