import os, sys, math
import numpy as np
import control

import math
import matplotlib
import matplotlib.pyplot as plt
from warnings import warn

import numpy as np
import scipy as sp

# Frequencies of interest
omega = np.logspace(-1, 3, 1001)

# Laplace variable
s = control.tf('s')

# MIMO loop transfer gain for the spinning satellite example
# https://www.mathworks.com/help/robust/ug/mimo-stability-margins-for-spinning-satellite.html
P = control.ss([[0, 10],[-10, 0]], np.eye(2), [[1, 10], [-10, 1]], [[0, 0],[0, 0]]) # plant
C = control.ss([],[],[], [[1, -2], [0, 1]]) # controller
L = P*C # output loop gain

print(f"------------- Sensitivity function (S) -------------")
DM, GM, PM = control.disk_margins(L, omega, skew = 1.0, returnall = True) # S-based (S)
print(f"min(DM) = {min(DM)} (omega = {omega[np.argmin(DM)]})")
print(f"GM = {GM[np.argmin(DM)]} dB")
print(f"PM = {PM[np.argmin(DM)]} deg")
print(f"min(GM) = {min(GM)} dB")
print(f"min(PM) = {min(PM)} deg\n")

plt.figure(3)
plt.subplot(3,3,1)
plt.semilogx(omega, DM, label='$\\alpha$')
plt.legend()
plt.title('Disk Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 2])

plt.figure(3)
plt.subplot(3,3,4)
plt.semilogx(omega, GM, label='$\\gamma_{m}$')
plt.ylabel('Gain Margin (dB)')
plt.legend()
plt.title('Gain-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 40])

plt.figure(3)
plt.subplot(3,3,7)
plt.semilogx(omega, PM, label='$\\phi_{m}$')
plt.ylabel('Phase Margin (deg)')
plt.legend()
plt.title('Phase-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 90])

print(f"------------- Complementary sensitivity function (T) -------------")
DM, GM, PM = control.disk_margins(L, omega, skew = -1.0, returnall = True) # T-based (T)
print(f"min(DM) = {min(DM)} (omega = {omega[np.argmin(DM)]})")
print(f"GM = {GM[np.argmin(DM)]} dB")
print(f"PM = {PM[np.argmin(DM)]} deg")
print(f"min(GM) = {min(GM)} dB")
print(f"min(PM) = {min(PM)} deg\n")

plt.figure(3)
plt.subplot(3,3,2)
plt.semilogx(omega, DM, label='$\\alpha$')
plt.legend()
plt.title('Disk Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 2])

plt.figure(3)
plt.subplot(3,3,5)
plt.semilogx(omega, GM, label='$\\gamma_{m}$')
plt.ylabel('Gain Margin (dB)')
plt.legend()
plt.title('Gain-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 40])

plt.figure(3)
plt.subplot(3,3,8)
plt.semilogx(omega, PM, label='$\\phi_{m}$')
plt.ylabel('Phase Margin (deg)')
plt.legend()
plt.title('Phase-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 90])

print(f"------------- Balanced sensitivity function (S - T) -------------")
DM, GM, PM = control.disk_margins(L, omega, skew = 0.0, returnall = True) # balanced (S - T)
print(f"min(DM) = {min(DM)} (omega = {omega[np.argmin(DM)]})")
print(f"GM = {GM[np.argmin(DM)]} dB")
print(f"PM = {PM[np.argmin(DM)]} deg")
print(f"min(GM) = {min(GM)} dB")
print(f"min(PM) = {min(PM)} deg\n")

plt.figure(3)
plt.subplot(3,3,3)
plt.semilogx(omega, DM, label='$\\alpha$')
plt.legend()
plt.title('Disk Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 2])

plt.figure(3)
plt.subplot(3,3,6)
plt.semilogx(omega, GM, label='$\\gamma_{m}$')
plt.ylabel('Gain Margin (dB)')
plt.legend()
plt.title('Gain-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 40])

plt.figure(3)
plt.subplot(3,3,9)
plt.semilogx(omega, PM, label='$\\phi_{m}$')
plt.ylabel('Phase Margin (deg)')
plt.legend()
plt.title('Phase-Only Margin')
plt.grid()
plt.xlim([omega[0], omega[-1]])
plt.ylim([0, 90])