from matplotlib import pyplot as plt
import control as ct

# Controller gain and sampling period
k = 5.4
Ts = 0.003

# Define the transfer functions ins the 's' domain
Gp = ct.tf([1740], [0.25, 1, 0]) # Plant
Gc = ct.tf([k, 250], [1, k*50])  # Phase lead compensator
G_series = ct.series(Gp, Gc)     # Couple them in series  # Gp x Gc

# Unity feedback - discrete
T_s = ct.feedback(Gp * Gc)

# Convert closed loop transfer function to discrete time
T_z = ct.c2d(T_s, Ts = 0.003, method='zoh')
print(T_z)

plt.style.use('dark_background')

# Plot this response
response = ct.step_response(T_z)

# Plot using a discrete stem plot
plt.style.use('dark_background')
plt.figure()
plt.stem(response.time, response.outputs)
plt.title("Discrete-Time Step Response (dt=True)")
plt.xlabel(f"Sample Index (k), Ts = {Ts} s")
plt.ylabel("Amplitude")
plt.xlim(-0.01, xmax=0.1)
plt.ylim(-0.05, ymax=1.5)
plt.grid(True)
plt.show()