import numpy as np
import control as ct
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import re

# --- 辅助函数：将系数数组转换为LaTeX数学公式字符串 ---
def coeffs_to_latex(coeffs, var='s'):
    """
    一个更健壮的函数，用于将系数数组转换为美观的LaTeX字符串。
    """
    coeffs = np.atleast_1d(coeffs)
    if coeffs.size == 1:
        return f"{coeffs[0]:g}"

    latex_str = ""
    degree = len(coeffs) - 1
    for i, coeff in enumerate(coeffs):
        if np.isclose(coeff, 0):
            continue
        
        sign = ""
        if i > 0 and latex_str:
            sign = " + " if coeff > 0 else " - "
        elif coeff < 0:
            sign = "-"
        
        coeff_abs = abs(coeff)
        coeff_str = f"{coeff_abs:g}" if not np.isclose(coeff_abs, 1) or degree - i == 0 else ""
        
        power = degree - i
        power_str = ""
        if power > 0:
            power_str = var
            if power > 1:
                power_str += f"^{{{power}}}"
        
        latex_str += f"{sign}{coeff_str}{power_str}"

    return latex_str if latex_str else "0"

def _parse_coeffs(s_val):
    """从字符串解析系数，返回Numpy数组"""
    cleaned_str = re.sub(r'[^0-9,\-.]', '', s_val)
    if not cleaned_str:
        return np.array([])
    return np.array([float(n) for n in cleaned_str.split(',') if n])

# --- 功能函数1：显示传递函数 ---
def display_transfer_function(num_str, den_str):
    try:
        num_coeffs = _parse_coeffs(num_str)
        den_coeffs = _parse_coeffs(den_str)
        if num_coeffs.size == 0 or den_coeffs.size == 0:
            return "分子或分母不能为空"
        num_latex = coeffs_to_latex(num_coeffs)
        den_latex = coeffs_to_latex(den_coeffs)
        return f"$$ G(s) = \\frac{{{num_latex}}}{{{den_latex}}} $$"
    except Exception as e:
        return f"输入格式错误: {e}"

# --- 功能函数2：执行时域分析和绘图 ---
def time_domain_analysis(num_str, den_str):
    try:
        num = _parse_coeffs(num_str)
        den = _parse_coeffs(den_str)
        if num.size == 0 or den.size == 0:
            return None, "错误：分子或分母系数不能为空。"
        
        system = ct.TransferFunction(num, den)
        t = np.linspace(0, 15, 1000)
        T_step, yout_step = ct.step_response(system, T=t)
        T_impulse, yout_impulse = ct.impulse_response(system, T=t)
        
        metrics_text = "无法计算性能指标。"
        try:
            info = ct.step_info(system)
            if isinstance(info, dict):
                metrics_text = (
                    f"上升时间: {info.get('RiseTime', float('nan')):.2f} s\n"
                    f"峰值时间: {info.get('PeakTime', float('inf')):.2f} s\n"
                    f"峰值: {info.get('Peak', float('inf')):.2f}\n"
                    f"超调量: {info.get('Overshoot', float('nan')):.1f} %\n"
                    f"调节时间: {info.get('SettlingTime', float('nan')):.2f} s\n"
                    f"稳态值: {info.get('SteadyStateValue', float('nan')):.2f}"
                )
        except (RuntimeError, Exception):
            metrics_text = "系统可能不稳定，无法计算阶跃响应指标。"

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        ax1.plot(T_step, yout_step); ax1.set_title("Unit Step Response"); ax1.set_xlabel("Time (s)"); ax1.set_ylabel("Amplitude"); ax1.grid(True)
        ax2.plot(T_impulse, yout_impulse); ax2.set_title("Unit Impulse Response"); ax2.set_xlabel("Time (s)"); ax2.set_ylabel("Amplitude"); ax2.grid(True)
        plt.tight_layout()
        
        return fig, metrics_text
    except Exception as e:
        return None, f"错误: {e}\n请检查您的输入。"

# --- 功能函数3：执行频域分析和绘图 ---
def frequency_domain_analysis(num_str, den_str, k):
    try:
        num_coeffs = _parse_coeffs(num_str)
        den_coeffs = _parse_coeffs(den_str)

        if num_coeffs.size == 0 or den_coeffs.size == 0:
            return None, "错误：系数不能为空。", "", ""

        system_with_gain = ct.TransferFunction(k * num_coeffs, den_coeffs)
        tf_latex_with_gain = f"$$ G_{{open}}(s) = K \\times \\frac{{{coeffs_to_latex(num_coeffs)}}}{{{coeffs_to_latex(den_coeffs)}}} \\quad \\text{{where }} K = {k:.2f} $$"

        fig = plt.figure(figsize=(12, 6)); gs = fig.add_gridspec(2, 2)
        ax_mag = fig.add_subplot(gs[0, 0]); ax_phase = fig.add_subplot(gs[1, 0], sharex=ax_mag)
        omega_range = np.logspace(-2, 3, 1000)
        mag, phase, omega = ct.frequency_response(system_with_gain, omega_range)
        
        ax_mag.semilogx(omega, 20 * np.log10(mag)); ax_mag.grid(True, which='both'); ax_mag.set_ylabel("Magnitude (dB)"); ax_mag.set_title("Bode Plot")
        ax_phase.semilogx(omega, np.rad2deg(phase)); ax_phase.grid(True, which='both'); ax_phase.set_ylabel("Phase (deg)"); ax_phase.set_xlabel("Frequency (rad/s)"); ax_phase.axhline(y=-180, color='r', linestyle='--', linewidth=0.8)

        ax_nyquist = fig.add_subplot(gs[:, 1])
        ct.nyquist_plot(system_with_gain, ax=ax_nyquist); ax_nyquist.set_title("Nyquist Plot"); ax_nyquist.grid(True)
        plt.tight_layout()

        metrics_text, stability_text = "无法计算稳定裕度。", "**评估**: <font color='orange'>**无法判断**</font>"
        try:
            gm, pm, _, _ = ct.margin(system_with_gain)
            gm_db = 20 * np.log10(gm) if gm > 0 and np.isfinite(gm) else float('inf')
            
            is_stable = gm_db > 0 and pm > 0
            stability_text = f"**评估**: <font color='{'green' if is_stable else 'red'}'>**{'系统稳定' if is_stable else '系统不稳定'}**</font>"

            metrics_text = ""
            if np.isinf(gm_db): metrics_text += f"增益裕度 (GM): inf dB\n(注: 相位未穿过-180°线)\n"
            else: metrics_text += f"增益裕度 (GM): {gm_db:.2f} dB\n"
            if np.isinf(pm): metrics_text += f"相位裕度 (PM): 未定义"
            else: metrics_text += f"相位裕度 (PM): {pm:.2f} deg"
        except Exception:
            pass
            
        return fig, metrics_text, tf_latex_with_gain, stability_text
    except Exception as e:
        return None, f"错误: {e}", "", ""

# --- 功能函数4：执行根轨迹分析 ---
def root_locus_analysis(num_str, den_str, log_k):
    try:
        k = 10**log_k
        num_coeffs = _parse_coeffs(num_str)
        den_coeffs = _parse_coeffs(den_str)

        if num_coeffs.size == 0 or den_coeffs.size == 0:
            return None, "错误：系数不能为空。", k
        
        open_loop_system = ct.TransferFunction(num_coeffs, den_coeffs)
        fig, ax = plt.subplots(figsize=(8, 6))
        
        closed_loop_system = ct.feedback(k * open_loop_system, 1)
        current_poles = ct.poles(closed_loop_system)
        
        # 绘制根轨迹
        rlist, klist = ct.root_locus(open_loop_system, plot=False, grid=False)
        for i in range(rlist.shape[1]): 
            ax.plot(np.real(rlist[:, i]), np.imag(rlist[:, i]), 'b-', linewidth=1.5, alpha=0.7)
        
        # 只标记当前增益下的闭环极点
        ax.plot(np.real(current_poles), np.imag(current_poles), 'rx', 
                markersize=12, markeredgewidth=2.5, label=f'Poles at K={k:.2f}', zorder=5)
        
        # 动态计算坐标轴范围 - 只基于当前闭环极点位置
        pole_real = np.real(current_poles)
        pole_imag = np.imag(current_poles)
        
        # 找到极点的最大绝对值（实部和虚部）
        max_real = np.max(np.abs(pole_real))
        max_imag = np.max(np.abs(pole_imag))
        
        # 取实部和虚部的最大值作为基准
        max_val = max(max_real, max_imag)
        
        # 确保最小显示范围
        if max_val < 1.0:
            max_val = 1.0
        
        # 添加 30% 的边距，使图像更美观
        margin = 1.3
        limit = max_val * margin
        
        # 设置对称的坐标轴范围（以原点为中心）
        ax.set_xlim(-limit, limit)
        ax.set_ylim(-limit, limit)
        
        # 绘制虚轴（稳定性边界）和实轴
        ax.axvline(x=0, color='k', linestyle='--', linewidth=0.8, alpha=0.5)
        ax.axhline(y=0, color='k', linestyle='-', linewidth=0.5, alpha=0.3)
        
        ax.set_xlabel("Real Axis")
        ax.set_ylabel("Imaginary Axis")
        ax.set_title("Root Locus")
        ax.grid(True, alpha=0.3)
        ax.legend(loc='best', fontsize=9)
        ax.set_aspect('equal', adjustable='box')

        poles_text = "Closed-Loop Poles:\n"
        for p in current_poles: 
            poles_text += f"{p.real:.3f} {'+' if p.imag >= 0 else '-'} {abs(p.imag):.3f}j\n"
        return fig, poles_text, k
    except Exception as e:
        return None, f"错误: {e}", 10**log_k