import sympy as sm
import numpy as np
import matplotlib.pyplot as plt

sm.init_printing()

print("PROBLEM 1:")
print("part a:")
s, z, t = sm.symbols("s z t")

J, B, RHO, RADIUS, L, RESIST, T = sm.symbols(
    "J, B, rho,  r,  L,  R, T", Real=True, Positive=True
)
K_T, K_B = sm.symbols("K_T,  K_b")

# loop gain
Loop_Gain = K_T * 1 / (J * s + B)

# transfer function from motor current command to linear displacement (ignoring feedback)
X_over_I = Loop_Gain * RHO / s * RADIUS
sm.pprint(X_over_I)
X_over_I = X_over_I.expand().simplify()
sm.pprint(X_over_I)

print("part b:")
# recall ZOH_eq
# G(z) = (1-z**-1) Z{L**-1{G(s)/s}}
print("Finding inverse laplace of G/s")
G_s = X_over_I
sm.pprint(G_s / s)
g_t = sm.inverse_laplace_transform(G_s / s, s, t)
sm.pprint(g_t.expand())

# make z substitution subs

G_z = (1 - z**-1) * (
    K_T * T * RADIUS * RHO / B * (z / (z - 1) ** 2)
    - J * K_T * RADIUS * RHO / B**2
    + J * K_T * RADIUS * RHO / B**2 * (z / (z - sm.exp(-B / J * T)))
)

sm.pprint(G_z)

# part c was on the board
print("part c")
F = sm.Matrix([[0, RHO], [0, -B / J]])
G = sm.Matrix([0, K_T / J])

# part d
print("part d")
A = sm.exp(F * T)
B = sm.integrate(A, (T, 0, T))
sm.pprint(F)
sm.pprint(G)
sm.pprint(A)
sm.pprint(B)

########################################################a
print("PROBLEM 4:")
print("bilinear transform")
T = 0.2  # s

G_s = 10 / (s + 10)
G_z = G_s.subs(s, 2 * (z - 1) / T * (z + 1))
sm.pprint(G_z)
sm.pprint(G_z.simplify())

print("bilinear with warping")
omega_star = 10  # rad/s
G_z_warped = G_s.subs(s, (omega_star / np.tan(omega_star * T / 2)) * (z - 1) / (z + 1))
sm.pprint(G_z_warped)
sm.pprint(G_z_warped.simplify())

# Create a frequency vector from 0 to 10 rad/s (avoid exactly 0 to prevent log(0) issues)
omega = np.linspace(5, 15, 500)

# Continuous-time frequency response (evaluate at s = j*omega)
G_s_response = 10 / (10 + 1j * omega)

# For the discrete transfer functions, substitute z = exp(j*omega*T)
z_vals = np.exp(1j * omega * T)

# Convert symbolic expressions into functions that accept numpy arrays
G_z_func = sm.lambdify(z, G_z, "numpy")
G_z_warped_func = sm.lambdify(z, G_z_warped, "numpy")

G_z_response = G_z_func(z_vals)
G_z_warped_response = G_z_warped_func(z_vals)


# Helper function to compute magnitude (in dB) and phase (in degrees)
def compute_mag_phase(response):
    mag = 20 * np.log10(np.abs(response))
    phase = np.angle(response, deg=True)
    return mag, phase


# Compute responses
mag_s, phase_s = compute_mag_phase(G_s_response)
mag_z, phase_z = compute_mag_phase(G_z_response)
mag_z_warped, phase_z_warped = compute_mag_phase(G_z_warped_response)

# Plot the Bode plots
plt.figure(figsize=(10, 8))

# Magnitude plot
plt.subplot(2, 1, 1)
plt.semilogx(omega, mag_s, label="Continuous: G(s)")
plt.semilogx(omega, mag_z, label="Bilinear: G(z)")
plt.semilogx(omega, mag_z_warped, label="Warped: G(z)_warped")
plt.title("Bode Magnitude Plot")
plt.ylabel("Magnitude (dB)")
plt.legend()
plt.grid(True)

# Phase plot
plt.subplot(2, 1, 2)
plt.semilogx(omega, phase_s, label="Continuous: G(s)")
plt.semilogx(omega, phase_z, label="Bilinear: G(z)")
plt.semilogx(omega, phase_z_warped, label="Warped: G(z)_warped")
plt.title("Bode Phase Plot")
plt.xlabel("Frequency (rad/s)")
plt.ylabel("Phase (degrees)")
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()