import numpy as np
%matplotlib
import matplotlib.pyplot as plt
import control as ct
s = ct.tf('s')

def step(H):
    resp = ct.step_response(H)
    resp.plot()

# Traditional Bode Plot, Magnitude(dB) log-log, Phase(degrees) log-linear
def bode(H):
    resp = ct.bode_plot(H, dB=True, wrap_phase=False, deg=True)

def pole(H):
    print(ct.poles(H))

def feedback(H1,H2):
    G = ct.feedback(H1,H2)
    return(G)

# Unity Feedback
def fdbk(H):
    G = ct.feedback(H,1)
    return(G)

def shg():
    plt.show()

# Plot Magnitude and Phase(degrees) versus Freq(Hz) on log-linear scales
def bodeF(G):
    omega = np.logspace(-1, 4, 500)
    mag, phase, omega = ct.frequency_response(G, omega=omega)
    phase = phase * 180 / np.pi

    plt.subplot(2, 1, 1)
    plt.semilogx(omega / (2 * np.pi), mag.squeeze())
    plt.grid(True)
    plt.title("Magnitude and Phase versus frequency (hertz)")
    plt.ylabel("Magnitude (abs)")

    plt.subplot(2, 1, 2)
    plt.semilogx(omega / (2 * np.pi), phase.squeeze())
    plt.grid(True)
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Phase (Degrees)")


# Controller Gains and transfer function
Kp = 0.5
Kd = 0.2
Ki = 4

Ke2c = Kp + Kd * s + Ki / s  # Error to motor command

# Motor Command to angular acceleration
# Adjust Beta and Ga so that your arm's response to sines match
Ga = 5 # Change to match model
Beta = 5 # Change to match model
Hc2a = 1/(s+5) # FIX! 

# Angular acceleration to angle transfer function
Ha2t = 1000/(s*(s+1)) # FIX!

# Determine closed loop transfer function
Gs = feedback(Ke2c*Hc2a*Ha2t,1);  

# Plots Mag and Angle on log(Hz)-linear then in traditional
# Bode plot Mag and Angle on log(rad/s)-log(magnitude, log-linear(degrees)
plt.figure(1)
bodeF(Gs)
shg()

plt.figure(2)
bode(Gs)
shg()

# Closed loop step response
plt.figure(3)
step(Gs)
shg()

print(Ha2t)