import pypho
from pypho import functions as pf

# from pypho_functions import *
# from pypho import functions as pf
import numpy as np
import copy
import matplotlib.pyplot as plt
from scipy.interpolate import UnivariateSpline

# Define network elements
gp       = pypho.setup(nos =2**4, sps = 256, symbolrate = 10.0e9)


symbolsrc   = pypho.symbols(glova = gp, nos = gp.nos, pattern = 'debruijn')
esigsrc  = pypho.signalsrc(glova = gp, pulseshape = 'rect' , fwhm = 0.85)
sig_1550 = pypho.lasmod(glova = gp, power = 0, Df = 0, theta = 0)
SSMF     = pypho.fiber(glova = gp, l = 60.0e3,  D = 17.0,   S = 0, alpha = 0.2e-12, gamma = 1.4e-12, phi_max = 10.0)

# Simulation
bits = symbolsrc()
esig = esigsrc(bitsequence = bits)

E_Tx = sig_1550(esig = esig)                                                      



# Define your parameters here
T_0 = 25.0e-12
z = SSMF.l
D = 17.0
beta_2, beta_3 = pf.DS_to_beta(17.0, 0, gp.lambda0)


# Create a single pulse with gaussian shape (not power!)
E_Tx[0]['E'][0] = E_Tx[0]['E'][0]*0 + np.exp(-(gp.timeax()-gp.timeax()[-1]/2)**2 / (2.0*T_0**2) )

E = copy.deepcopy(E_Tx)


# Fiber transmission
E = SSMF(E = E, D = D, l = z)
#sys.exit()

# Get FWHM of the input signal E_Tx
spline_0 = UnivariateSpline(gp.timeax()*1.0e12, np.abs(E_Tx[0]['E'][0])-1*np.max(np.abs(E_Tx[0]['E'][0]))/2, s=0)
r1_0, r2_0 = spline_0.roots() # find the roots
# Get FWHM of the output signal E
spline_1 = UnivariateSpline(gp.timeax()*1.0e12, np.abs(E[0]['E'][0])-1*np.max(np.abs(E[0]['E'][0]))/2, s=0)
r1_1, r2_1 = spline_1.roots() # find the roots

T_FWHM_0 = (r2_0-r1_0) * 1e-12
T_0_plot = T_FWHM_0 / 2.35482
T_FWHM_1 = (r2_1-r1_1) * 1e-12
L_D = (T_0_plot)**2  / np.abs(beta_2)

# Plot Input and Output signal 
plt.figure(1)
plt.plot(gp.timeax()*1.0e12, np.abs(E_Tx[0]['E'][0]), 'r', label='$E(0, t)$')
plt.plot(gp.timeax()*1.0e12, np.abs(E[0]['E'][0]),    'g', label=f'$E(z={SSMF.l/1000}km, t$)')

plt.annotate(text='', xy=(r1_0,np.max(np.abs(E_Tx[0]['E'][0]))/2), xytext=(r2_0,np.max(np.abs(E_Tx[0]['E'][0]))/2), arrowprops=dict(arrowstyle='<->'))
plt.text(np.max((r2_0,r2_1))+10, np.max(np.abs(E_Tx[0]['E'][0]))/2, f'$T_{{FWHM,0}}$ = {r2_0-r1_0:.2f} ps', fontsize=12, horizontalalignment='left', verticalalignment='center')
plt.annotate(text='', xy=(r1_1,np.max(np.abs(E[0]['E'][0]))/2), xytext=(r2_1,np.max(np.abs(E[0]['E'][0]))/2), arrowprops=dict(arrowstyle='<->'))
plt.text(np.max((r2_0,r2_1))+10, np.max(np.abs(E[0]['E'][0]))/2, f'$T_{{FWHM,1}}$ = {r2_1-r1_1:.2f} ps', fontsize=12, horizontalalignment='left', verticalalignment='center')

plt.ylabel('$|E|$ a.u.')
plt.xlabel('Time $t$ [ps]')
plt.grid()
legend = plt.legend(loc='upper right')
# Print the results
print(f'Input signal 1/e-pulse width by definition: T_0 = {T_0*1e12:.6f} ps')
print(f'Input signal 1/e-pulse width from plot: T_0 = {T_0_plot*1e12:.6f} ps')
print(f'Input signal FWHM-pulse width from plot: T_FWHM,0 = {T_FWHM_0*1e12:.6f} ps')
print(f'Output signal FWHM-pulse width from plot: T_FWHM,1 = {T_FWHM_1*1e12:.6f} ps')
print(f'Calculated output FWHM-pulse width: T_FWHM,1 = {T_FWHM_0 * np.sqrt(1 + (z/L_D)**2)*1e12:.6f} ps')

plt.show()