import numpy as np
import matplotlib.pyplot as plt
import control
from control import tf, step_response, impulse_response, bode
import os

# Создаём папку для сохранения изображений
if not os.path.exists("img"):
    os.makedirs("img")
    
###############################################################################
# Функция проверки строгой каузальности системы
###############################################################################
def is_strictly_proper(sys_obj):
    """
    Проверяет, является ли система строго каузальной,
    то есть степень знаменателя не меньше степени числителя.
    """
    num = sys_obj.num[0][0]
    den = sys_obj.den[0][0]
    return (len(den) - 1) >= (len(num) - 1)

###############################################################################
# Общие базовые данные
###############################################################################
# Для систем с одним параметром:
K0_1 = 1.0
K_list_1 = [round(0.85 * K0_1, 2), round(K0_1, 2), round(1.15 * K0_1, 2)]

# Для систем с двумя параметрами (на основе методических указаний):
K0_2 = 3.0
T0_2 = 1.0
K_list_2 = [round(0.9 * K0_2, 2), round(K0_2, 2), round(1.1 * K0_2, 2)]  # [2.70, 3.00, 3.30]
T_list_2 = [round(0.85 * T0_2, 2), round(T0_2, 2), round(1.15 * T0_2, 2)]  # [0.85, 1.00, 1.15]

s = control.tf('s')  # Оператор s

###############################################################################
# Универсальная функция построения 4 графиков в одном окне
###############################################################################
def plot_4_in_1(title_str, systems, legend_labels):
    """
    Для списка систем (systems) и меток (legend_labels) строит:
      1) Переходную характеристику (step response)
      2) Импульсную характеристику (impulse response) – если система строго каузальна
      3) ЛАЧХ (Bode magnitude, дБ)
      4) ЛФЧХ (Bode phase, град)
    Если система не строго каузальна, временные характеристики не строятся,
    а в легенде добавляется пометка "(N/A)".
    """
    fig = plt.figure(figsize=(10, 8))
    ax_step       = fig.add_subplot(2, 2, 1)
    ax_impulse    = fig.add_subplot(2, 2, 2)
    ax_bode_mag   = fig.add_subplot(2, 2, 3)
    ax_bode_phase = fig.add_subplot(2, 2, 4)

    # Построение временных характеристик
    for i, W in enumerate(systems):
        if is_strictly_proper(W):
            t_step, y_step = step_response(W)
            ax_step.plot(t_step, y_step, label=legend_labels[i])
            
            t_imp, y_imp = impulse_response(W)
            ax_impulse.plot(t_imp, y_imp, label=legend_labels[i])
        else:
            ax_step.plot([], [], label=legend_labels[i] + " (N/A)")
            ax_impulse.plot([], [], label=legend_labels[i] + " (N/A)")
            print(f"Система {legend_labels[i]} не строго каузальна. Временные характеристики не строятся.")

    ax_step.set_title("Переходная (step)")
    ax_step.set_xlabel("t, с")
    ax_step.set_ylabel("y(t)")
    ax_step.grid(True)
    ax_step.legend(loc="best")

    ax_impulse.set_title("Импульсная (impulse)")
    ax_impulse.set_xlabel("t, с")
    ax_impulse.set_ylabel("y(t)")
    ax_impulse.grid(True)
    ax_impulse.legend(loc="best")

    # Построение частотных характеристик (Bode)
    omega = np.logspace(-2, 2, 200)
    for i, W in enumerate(systems):
        mag, phase, w = bode(W, omega, plot=False)
        mag_dB = 20 * np.log10(mag)
        phase_deg = np.rad2deg(phase)
        ax_bode_mag.semilogx(w, mag_dB, label=legend_labels[i])
        ax_bode_phase.semilogx(w, phase_deg, label=legend_labels[i])
        
    ax_bode_mag.set_title("ЛАЧХ (magnitude, дБ)")
    ax_bode_mag.set_xlabel(r"$\omega$, рад/с")
    ax_bode_mag.set_ylabel("Амплитуда, дБ")
    ax_bode_mag.grid(True, which="both")
    ax_bode_mag.legend(loc="best")

    ax_bode_phase.set_title("ЛФЧХ (phase, град)")
    ax_bode_phase.set_xlabel(r"$\omega$, рад/с")
    ax_bode_phase.set_ylabel("Фаза, град")
    ax_bode_phase.grid(True, which="both")
    ax_bode_phase.legend(loc="best")

    fig.suptitle(title_str, fontsize=14)
    fig.tight_layout()

    existing_imgs = [f for f in os.listdir("img") if f.endswith(".png")]
    file_num = len(existing_imgs) + 1
    plt.savefig(f"img/{file_num}.png")

###############################################################################
# Эксперименты для систем с одним параметром
###############################################################################
# 1) Интегрирующее звено: W(s) = K/s
def example_integrator():
    systems = []
    legends = []
    for K in K_list_1:
        W = tf([K], [1, 0])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Интегрирующее звено (K варьируется)", systems, legends)

# 2) Дифференцирующее звено: W(s) = K*s
def example_differentiator():
    systems = []
    legends = []
    for K in K_list_1:
        W = tf([K, 0], [1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Дифференцирующее звено (K варьируется)", systems, legends)

# 3) Усилительное звено: W(s) = K
def example_amplifier():
    systems = []
    legends = []
    for K in K_list_1:
        W = tf([K], [1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Усилительное звено (K варьируется)", systems, legends)

###############################################################################
# Эксперименты для систем с двумя параметрами (вариация по K и по T)
###############################################################################
# 4) Апериодическое звено 1-го порядка: W(s) = K/(T*s + 1)
def example_first_order_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K], [T0_2, 1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Апериод. звено 1-го порядка (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_first_order_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2], [T, 1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Апериод. звено 1-го порядка (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 5) Апериодическое звено 2-го порядка: W(s) = K/(T^2*s^2 + a*T*s + 1), a=3.0
def example_second_order_vary_K():
    systems = []
    legends = []
    a = 3.0
    for K in K_list_2:
        W = tf([K], [T0_2**2, a * T0_2, 1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Апериод. звено 2-го порядка (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_second_order_vary_T():
    systems = []
    legends = []
    a = 3.0
    for T in T_list_2:
        W = tf([K0_2], [T**2, a * T, 1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Апериод. звено 2-го порядка (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 6) Колебательное звено: W(s) = K/(T^2*s^2 + 2*ξ*T*s + 1), ξ=0.7
def example_oscillatory_vary_K():
    systems = []
    legends = []
    xi = 0.7
    for K in K_list_2:
        W = tf([K], [T0_2**2, 2 * xi * T0_2, 1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Колебательное звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_oscillatory_vary_T():
    systems = []
    legends = []
    xi = 0.7
    for T in T_list_2:
        W = tf([K0_2], [T**2, 2 * xi * T, 1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Колебательное звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 7) Консервативное звено: W(s) = K/(T*s^2 + 1)
def example_conservative_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K], [T0_2, 0, 1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Консервативное звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_conservative_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2], [T, 0, 1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Консервативное звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 8) Реальное интегрирующее звено: W(s) = K/(s*(T*s + 1))
def example_real_integrator_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K], [T0_2, 1, 0])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Реальное интегрирующее звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_real_integrator_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2], [T, 1, 0])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Реальное интегрирующее звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 9) Реальное дифференцирующее звено: W(s) = K*s/(T*s + 1)
def example_real_differentiator_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K, 0], [T0_2, 1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Реальное дифференцирующее звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_real_differentiator_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2, 0], [T, 1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Реальное дифференцирующее звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 10) Форсирующее звено: W(s) = K*(T*s + 1)
def example_forcing_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K * T0_2, K], [1])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Форсирующее звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_forcing_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2 * T, K0_2], [1])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Форсирующее звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

# 11) Изодромное звено: W(s) = K*(T*s + 1)/s
def example_isodromic_vary_K():
    systems = []
    legends = []
    for K in K_list_2:
        W = tf([K * T0_2, K], [1, 0])
        systems.append(W)
        legends.append(f"K={K:.2f}")
    plot_4_in_1("Изодромное звено (T const = " + f"{T0_2:.2f}" + ", K варьируется)", systems, legends)

def example_isodromic_vary_T():
    systems = []
    legends = []
    for T in T_list_2:
        W = tf([K0_2 * T, K0_2], [1, 0])
        systems.append(W)
        legends.append(f"T={T:.2f}")
    plot_4_in_1("Изодромное звено (K const = " + f"{K0_2:.2f}" + ", T варьируется)", systems, legends)

###############################################################################
# Запуск экспериментов
###############################################################################
if __name__ == "__main__":
    # Системы с одним параметром
    example_integrator()
    example_differentiator()
    example_amplifier()
    
    # Апериодическое звено 1-го порядка (2 эксперимента)
    example_first_order_vary_K()
    example_first_order_vary_T()
    
    # Апериодическое звено 2-го порядка (2 эксперимента)
    example_second_order_vary_K()
    example_second_order_vary_T()
    
    # Колебательное звено (2 эксперимента)
    example_oscillatory_vary_K()
    example_oscillatory_vary_T()
    
    # Консервативное звено (2 эксперимента)
    example_conservative_vary_K()
    example_conservative_vary_T()
    
    # Реальное интегрирующее звено (2 эксперимента)
    example_real_integrator_vary_K()
    example_real_integrator_vary_T()
    
    # Реальное дифференцирующее звено (2 эксперимента)
    example_real_differentiator_vary_K()
    example_real_differentiator_vary_T()
    
    # Форсирующее звено (2 эксперимента)
    example_forcing_vary_K()
    example_forcing_vary_T()
    
    # Изодромное звено (2 эксперимента)
    example_isodromic_vary_K()
    example_isodromic_vary_T()