from scipy.special import lambertw
import numpy as np
import matplotlib.pyplot as plt
import math as m

T = 5           # cas
dt = 0.001153   # delka kroku
N = int(T/dt)   # pocet casovych intervalu

# hodnoty soucastek z prik. 3.6, viz str. 60
C = 0.001
R = 5
R_Z = 15
L = 2
U_A = 1
f = 1
I_0 = 0.000001
n = 1
U_T = 0.0257


t = np.linspace(0, T, N+1) # casova promenna
U_out = np.zeros(N+1) # pole, do ktereho budeme ukladat vypoctene hodnoty

# vstupni zdroj
def u_in(t):
    return U_A*np.sin(2*m.pi*f*t)


for i in range(2, N + 1):
    """vypocet pomoci Lambertovy fce W, viz vztah (3.58) na str. 60"""
    a = (1/(2*n*U_T))*( abs(u_in(t[i])) + ( (L*dt + 2*L*C*R_Z + R_Z*R*C*dt)/(R_Z * dt**2) )*U_out[i-1] - ( (L*C)/(dt**2) )*U_out[i-2] )
    b = (1/(2*n*U_T))*( 1 + ( (L+R*dt) * (R_Z*C+dt) )/(R_Z * dt**2) )
    c = 1 - (C/(I_0*dt))*U_out[i-1]
    d = (C*R_Z + dt)/(R_Z*I_0*dt)
    z = (b/d) * np.exp(a + (c*b)/d)
    U_out[i] = -c/d + (1/b) * lambertw(z)

# graficke zobrazeni
plt.plot(t, u_in(t), label='U_in') # vstupni signal
plt.plot(t, U_out, label='U_out') # vystupni napeti
plt.ylabel('u [V]', rotation=0)
plt.xlabel('čas [s]')
plt.legend()
plt.show()