import numpy as np
from scipy import signal
import plotly.graph_objects as go
import sympy as sp
from plotly.subplots import make_subplots
import matplotlib.pyplot as plt
import control

# Parameters
Lafe_val = 3.1e-3
Lg_val = 2e-3
Cf_val = 3.3e-6

# Symbolic Definition
s = sp.symbols('s')
Lafe, Lg, Cf = sp.symbols('Lafe Lg Cf')
G_g = 1 / (Lafe * Lg * Cf * s**3 + (Lafe + Lg) * s)
G_c = (1 + Lg * Cf * s**2) / (Lafe * Lg * Cf * s**3 + (Lafe + Lg) * s)

# Display symbolic transfer functions
print("Symbolic transfer function Gg(s):")
sp.pprint(G_g)
print("Symbolic transfer function Gc(s):")
sp.pprint(G_c)

## Transfer Function 1

# Calculate poles and zeros
# Extract numerator and denominator
num_g_sym, den_g_sym = sp.fraction(G_g)

zeros_g = sp.solve(num_g_sym, s)
poles_g = sp.solve(den_g_sym, s)

# Substitute parameter values
G_g_num = G_g.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val})
zeros_g_num = [z.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val}) for z in zeros_g]
poles_g_num = [p.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val}) for p in poles_g]

# Display zeros
print("\nZeros of the numerator:")
if zeros_g:
    for z_sym, z_num in zip(zeros_g, zeros_g_num):
        print(f"  s = {z_sym}   ≈ {sp.N(z_num, 4)}")
else:
    print("  No zeros (numerator is constant).")

# Display poles
print("\nPoles of the denominator:")
if poles_g:
    for p_sym, p_num in zip(poles_g, poles_g_num):
        print(f"  s = {p_sym}   ≈ {sp.N(p_num, 4)}")
else:
    print("  No poles (denominator is constant).")

# Extract numerator and denominator
num_g, den_g = sp.fraction(G_g_num)

# Calculate coefficients
num_coeffs_g = sp.Poly(num_g, s).all_coeffs()
den_coeffs_g = sp.Poly(den_g, s).all_coeffs()

# Convert to float 
num_coeffs_g = [float(c) for c in num_coeffs_g]
den_coeffs_g = [float(c) for c in den_coeffs_g]

# Transfer function
system_g = control.TransferFunction(num_coeffs_g, den_coeffs_g)

## Transfer Function 2

# Calculate poles and zeros
# Extract numerator and denominator
num_c_sym, den_c_sym = sp.fraction(G_c)

zeros_c = sp.solve(num_c_sym, s)
poles_c = sp.solve(den_c_sym, s)

# Substitute parameter values
G_c_num = G_c.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val})
zeros_c_num = [z.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val}) for z in zeros_c]
poles_c_num = [p.subs({Lafe: Lafe_val, Lg: Lg_val, Cf: Cf_val}) for p in poles_c]

# Display zeros
print("\nZeros of the numerator:")
if zeros_c:
    for z_sym, z_num in zip(zeros_c, zeros_c_num):
        print(f"  s = {z_sym}   ≈ {sp.N(z_num, 4)}")
else:
    print("  No zeros (numerator is constant).")

# Display poles
print("\nPoles of the denominator:")
if poles_c:
    for p_sym, p_num in zip(poles_c, poles_c_num):
        print(f"  s = {p_sym}   ≈ {sp.N(p_num, 4)}")
else:
    print("  No poles (denominator is constant).")

# Extract numerator and denominator
num_c, den_c = sp.fraction(G_c_num)

# Calculate coefficients
num_coeffs_c = sp.Poly(num_c, s).all_coeffs()
den_coeffs_c = sp.Poly(den_c, s).all_coeffs()

# Convert to float 
num_coeffs_c = [float(c) for c in num_coeffs_c]
den_coeffs_c = [float(c) for c in den_coeffs_c]

# Transfer function
system_c = control.TransferFunction(num_coeffs_c, den_coeffs_c)

# Frequency range
frequencies_rad_s = np.logspace(2, 5, num=500)
frequencies_pole = np.linspace(0.99 * float(sp.Abs(poles_c_num[1])), 1.01 * float(sp.Abs(poles_c_num[1])), 100)
frequencies_zero = np.linspace(0.99 * float(sp.Abs(zeros_c_num[1])), 1.01 * float(sp.Abs(zeros_c_num[1])), 100)
combined_frequencies = np.sort(np.unique(np.concatenate((
    frequencies_rad_s,
    frequencies_pole,
    frequencies_zero,
    [float(sp.Abs(poles_c_num[1]))],
    [float(sp.Abs(zeros_c_num[1]))]
))))

# Plot Bode diagram
control.bode_plot([system_g, system_c], combined_frequencies, dB=True, Hz=False, deg=True)
plt.show()