mag_ax.grid(visible=True, which='both', axis='both')
phase_ax.grid(visible=True, which='both', axis='both')

import control as ctrl
import matplotlib.pyplot as plt
import numpy as np

# Define your transfer function (example)
sys = ctrl.tf([300000], [1,360, 0])

# Add delay
# (num,den) = ct.pade(0.25,3)
# sys = ct.tf(num,den)*sys
print(sys)

# Plot the Bode diagram first to get the axes handles
# Use plot=False if you only want the data, but here we want the plot object
# ctrl.bode(sys, dB=True, deg=True, color='cyan', display_margins=True)
ctrl.bode_plot(sys, dB=True, deg=True, color='cyan', display_margins=True)

# Get the current figure and axes
fig = plt.gcf()
mag_ax, phase_ax = fig.axes

# Enable logarithmic frequency lines when 'display_margins=True'
mag_ax.grid(visible=True, which='both', axis='both')
phase_ax.grid(visible=True, which='both', axis='both')

# Calculate the margins and frequencies
Gm, Pm, Wcg, Wcp = ctrl.margin(sys)
GmdB = 20 * np.log10(Gm)

# Add custom lines and labels for Gain Margin
mag_ax.plot([Wcg, Wcg], [GmdB, 0], 'r--', linewidth=2) # Vertical line at phase crossover
mag_ax.annotate(f'    GM: {GmdB:.2f} dB', xy=(Wcg, GmdB/2), xytext=(Wcg, GmdB/2), color='red')

# Add custom lines and labels for Phase Margin
phase_ax.plot([Wcp, Wcp], [-180, -180 + Pm], 'g--', linewidth=2) # Vertical line at gain crossover
phase_ax.annotate(f'    PM: {Pm:.2f} deg', xy=(Wcp, -180 + Pm/2), xytext=(Wcp, -180 + Pm/2), color='red')

plt.show()