AutoControlCourse/app.py

import os
import gradio as gr
import time
from functools import partial

# 从各个模块导入所需的功能
import config
from analysis_functions import (
    display_transfer_function,
    time_domain_analysis,
    frequency_domain_analysis,
    root_locus_analysis
)
# ===== 算例演示模块函数导入（四阶段设计）=====
from case_demo_functions import run_distillation_demo, run_gpr_training, run_engine_design, run_motor_design, run_hybrid_demo
from chatbot import chat_with_ai
from user_stats import get_online_status_html, update_user_activity
from ui_components import (
    create_header,
    create_time_domain_tab,
    create_frequency_domain_tab,
    create_root_locus_tab,
    create_case_demo_tab,
    create_chatbot_tab
)


# ===== 系统资源监控 =====
def get_system_monitor_html():
    """获取 CPU / 内存 / GPU 使用率的 HTML 小组件"""
    try:
        import psutil
        cpu_pct = psutil.cpu_percent(interval=0)
        mem = psutil.virtual_memory()
        mem_pct = mem.percent
        mem_used_gb = mem.used / (1024 ** 3)
        mem_total_gb = mem.total / (1024 ** 3)
    except ImportError:
        return "<div style='text-align:center;color:#999;font-size:0.8em;'>psutil 未安装，无法监控系统资源</div>"

    # GPU 信息 — 优先用 nvidia-smi（不依赖 PyTorch CUDA 版本），再用 torch.cuda 兜底
    gpu_html = ""
    try:
        import subprocess as _sp
        _r = _sp.run(
            ['nvidia-smi', '--query-gpu=name,memory.used,memory.total,utilization.gpu',
             '--format=csv,noheader,nounits'],
            capture_output=True, text=True, timeout=3
        )
        if _r.returncode == 0 and _r.stdout.strip():
            _parts = [p.strip() for p in _r.stdout.strip().split('\n')[0].split(',')]
            _gpu_mem_used = float(_parts[1]) / 1024   # MiB → GiB
            _gpu_mem_total = float(_parts[2]) / 1024
            _gpu_util = float(_parts[3])
            _gc = '#ff6b6b' if _gpu_util > 80 else '#ffd93d' if _gpu_util > 50 else '#6bcb77'
            gpu_html = (
                f"<div style='display:inline-flex;align-items:center;gap:6px;'>"
                f"<span>🎮 GPU</span>"
                f"<div style='width:90px;height:8px;background:#444;border-radius:4px;overflow:hidden;'>"
                f"<div style='width:{min(_gpu_util, 100):.0f}%;height:100%;background:{_gc};'></div>"
                f"</div>"
                f"<span>{_gpu_util:.0f}% &nbsp; {_gpu_mem_used:.1f}/{_gpu_mem_total:.0f} GB</span>"
                f"</div>"
            )
        else:
            raise RuntimeError("nvidia-smi no output")
    except Exception:
        try:
            import torch as _torch
            if _torch.cuda.is_available():
                _mem_alloc = _torch.cuda.memory_allocated(0) / (1024 ** 3)
                _mem_total = _torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
                _util = _mem_alloc / max(_mem_total, 0.01) * 100
                _gc = '#ff6b6b' if _util > 80 else '#ffd93d' if _util > 50 else '#6bcb77'
                gpu_html = (
                    f"<div style='display:inline-flex;align-items:center;gap:6px;'>"
                    f"<span>🎮 GPU</span>"
                    f"<div style='width:90px;height:8px;background:#444;border-radius:4px;overflow:hidden;'>"
                    f"<div style='width:{min(_util, 100):.0f}%;height:100%;background:{_gc};'></div>"
                    f"</div>"
                    f"<span>{_mem_alloc:.1f}/{_mem_total:.0f} GB</span>"
                    f"</div>"
                )
            else:
                gpu_html = "<div style='display:inline-flex;align-items:center;gap:4px;'><span>🎮 GPU N/A</span></div>"
        except Exception:
            gpu_html = "<div style='display:inline-flex;align-items:center;gap:4px;'><span>🎮 GPU N/A</span></div>"

    cpu_color = '#ff6b6b' if cpu_pct > 80 else '#ffd93d' if cpu_pct > 50 else '#6bcb77'
    mem_color = '#ff6b6b' if mem_pct > 80 else '#ffd93d' if mem_pct > 50 else '#6bcb77'

    html = (
        f"<div style='display:flex;justify-content:center;gap:20px;flex-wrap:wrap;"
        f"font-size:0.82em;color:#ddd;padding:4px 10px;'>"
        # CPU
        f"<div style='display:inline-flex;align-items:center;gap:6px;'>"
        f"<span>🖥️ CPU</span>"
        f"<div style='width:90px;height:8px;background:#444;border-radius:4px;overflow:hidden;'>"
        f"<div style='width:{min(cpu_pct, 100):.0f}%;height:100%;background:{cpu_color};'></div>"
        f"</div>"
        f"<span>{cpu_pct:.0f}%</span>"
        f"</div>"
        # Memory
        f"<div style='display:inline-flex;align-items:center;gap:6px;'>"
        f"<span>💾 RAM</span>"
        f"<div style='width:90px;height:8px;background:#444;border-radius:4px;overflow:hidden;'>"
        f"<div style='width:{min(mem_pct, 100):.0f}%;height:100%;background:{mem_color};'></div>"
        f"</div>"
        f"<span>{mem_used_gb:.1f}/{mem_total_gb:.0f} GB ({mem_pct:.0f}%)</span>"
        f"</div>"
        # GPU
        f"{gpu_html}"
        f"</div>"
    )
    return html

# 加载外部CSS文件
with open(os.path.join(os.path.dirname(__file__), "assets", "styles.css"), "r", encoding="utf-8") as f:
    custom_css = f.read()

# --- 主应用界面 ---
with gr.Blocks(title="自动控制理论学习网站 - AI+数智平台", css=custom_css) as demo:
    # 1. 创建UI组件
    # 用户会话ID（隐藏组件）
    session_id = gr.State(value=lambda: str(time.time()) + "_" + str(hash(time.time())))

    # 创建头部信息和在线计数器
    online_counter = create_header()

    # 系统资源监控（始终可见）
    system_monitor = gr.HTML(value=get_system_monitor_html, elem_id="system-monitor")

    # 创建功能选项卡
    with gr.Tabs() as tabs:
        with gr.TabItem("⏱️ 时域分析 (Time Domain)", id=0):
            time_domain_ui = create_time_domain_tab()
        with gr.TabItem("📊 频域分析 (Frequency Domain)", id=1):
            freq_domain_ui = create_frequency_domain_tab()
        with gr.TabItem("🎯 根轨迹 (Root Locus)", id=2):
            root_locus_ui = create_root_locus_tab()
        # ===== 新增：算例演示 Tab（位于根轨迹与智能问答之间） =====
        with gr.TabItem("🧪 算例演示 (Case Demo)", id=3):
            case_demo_ui = create_case_demo_tab()
        with gr.TabItem("🤖 智能问答 (Q&A)", id=4):
            chatbot_ui = create_chatbot_tab()

    # 2. 绑定事件逻辑
    # --- 通用函数 ---
    # 每次操作前更新用户活跃状态
    def wrap_with_activity_update(fn, sid):
        update_user_activity(sid)
        # 使用 partial 将 session_id 绑定到函数上
        # 这样Gradio调用时就不需要显式传递session_id了
        return partial(fn, session_id=sid)

    # --- 时域分析事件 ---
    time_domain_ui["confirm_button"].click(
        fn=display_transfer_function,
        inputs=[time_domain_ui["num_input"], time_domain_ui["den_input"]],
        outputs=[time_domain_ui["tf_display"]]
    ).then(lambda: get_online_status_html(), outputs=online_counter)

    time_domain_ui["analyze_button"].click(
        fn=time_domain_analysis,
        inputs=[time_domain_ui["num_input"], time_domain_ui["den_input"]],
        outputs=[time_domain_ui["output_plot"], time_domain_ui["output_metrics"]]
    ).then(lambda: get_online_status_html(), outputs=online_counter)

    # --- 频域分析事件 ---
    def update_frequency_analysis_wrapper(num, den, log_k):
        k = 10**log_k
        fig, metrics, tf_latex, stability = frequency_domain_analysis(num, den, k)
        return fig, metrics, tf_latex, stability, k, get_online_status_html()

    freq_inputs = [freq_domain_ui["num_input"], freq_domain_ui["den_input"], freq_domain_ui["log_k_slider"]]
    freq_outputs = [
        freq_domain_ui["plot_output"],
        freq_domain_ui["metrics_display"],
        freq_domain_ui["tf_display"],
        freq_domain_ui["stability_display"],
        freq_domain_ui["k_number_display"],
        online_counter
    ]
    freq_domain_ui["log_k_slider"].release(
        fn=update_frequency_analysis_wrapper,
        inputs=freq_inputs,
        outputs=freq_outputs
    )

    # --- 根轨迹分析事件 ---
    def update_rl_view_wrapper(log_k, num, den):
        fig, poles, k_val = root_locus_analysis(num, den, log_k)
        return fig, poles, k_val, get_online_status_html()

    rl_inputs = [root_locus_ui["log_k_slider"], root_locus_ui["num_input"], root_locus_ui["den_input"]]
    rl_outputs = [
        root_locus_ui["plot_output"],
        root_locus_ui["poles_display"],
        root_locus_ui["k_number_display"],
        online_counter
    ]
    root_locus_ui["log_k_slider"].release(
        fn=update_rl_view_wrapper,
        inputs=rl_inputs,
        outputs=rl_outputs
    )
    
    # 频域：当传递函数输入框变化时自动更新
    freq_domain_ui["num_input"].change(
        fn=update_frequency_analysis_wrapper,
        inputs=freq_inputs, outputs=freq_outputs
    )
    freq_domain_ui["den_input"].change(
        fn=update_frequency_analysis_wrapper,
        inputs=freq_inputs, outputs=freq_outputs
    )

    # 根轨迹：当传递函数输入框变化时自动更新
    root_locus_ui["num_input"].change(
        fn=update_rl_view_wrapper,
        inputs=rl_inputs, outputs=rl_outputs
    )
    root_locus_ui["den_input"].change(
        fn=update_rl_view_wrapper,
        inputs=rl_inputs, outputs=rl_outputs
    )

    # ===== 算例演示事件绑定（四阶段）=====

    # --- 阶段零-A：GPR 模型训练 ---
    def run_gpr_wrapper(mode, sid, progress=gr.Progress(track_tqdm=True)):
        update_user_activity(sid)
        fig, summary = run_gpr_training(mode=mode, progress=progress)
        return fig, summary, get_online_status_html()

    case_demo_ui["gpr_run_button"].click(
        fn=run_gpr_wrapper,
        inputs=[case_demo_ui["gpr_mode"], session_id],
        outputs=[case_demo_ui["gpr_plot"], case_demo_ui["gpr_summary"], online_counter]
    )

    # --- 阶段零-B：NN 模型训练（蒸馏）---
    def run_distillation_wrapper(epochs, lr, hidden, sid, progress=gr.Progress(track_tqdm=True)):
        update_user_activity(sid)
        fig, summary = run_distillation_demo(epochs, lr, hidden, progress=progress)
        return fig, summary, get_online_status_html()

    case_demo_ui["distill_run_button"].click(
        fn=run_distillation_wrapper,
        inputs=[
            case_demo_ui["distill_epochs"], case_demo_ui["distill_lr"],
            case_demo_ui["distill_hidden"],
            session_id
        ],
        outputs=[case_demo_ui["distill_plot"], case_demo_ui["distill_summary"], online_counter]
    )

    # --- 阶段一：发动机控制器设计 ---
    def run_engine_design_wrapper(sim_time, dt, init_power, target_power,
                                   controller_type,
                                   kp, ki, kd, tau_fuel, K_inertia,
                                   mpc_horizon, mpc_W_power, mpc_W_dcost, mpc_overshoot,
                                   sid, progress=gr.Progress(track_tqdm=True)):
        update_user_activity(sid)
        fig, summary = run_engine_design(
            sim_time, dt, init_power, target_power,
            controller_type,
            kp, ki, kd, tau_fuel, K_inertia,
            mpc_horizon, mpc_W_power, mpc_W_dcost, mpc_overshoot / 100.0,
            progress=progress
        )
        return fig, summary, get_online_status_html()

    case_demo_ui["eng_run_button"].click(
        fn=run_engine_design_wrapper,
        inputs=[
            case_demo_ui["eng_sim_time"], case_demo_ui["eng_dt"],
            case_demo_ui["eng_init_power"], case_demo_ui["eng_target_power"],
            case_demo_ui["eng_controller_type"],
            case_demo_ui["eng_kp"], case_demo_ui["eng_ki"], case_demo_ui["eng_kd"],
            case_demo_ui["eng_tau_fuel"], case_demo_ui["eng_K_inertia"],
            case_demo_ui["eng_mpc_horizon"], case_demo_ui["eng_mpc_W_power"],
            case_demo_ui["eng_mpc_W_dcost"], case_demo_ui["eng_mpc_overshoot"],
            session_id
        ],
        outputs=[case_demo_ui["eng_plot"], case_demo_ui["eng_summary"], online_counter]
    )

    # --- 阶段二：电机控制器设计 ---
    def run_motor_design_wrapper(sim_time, dt, target_rpm, load_torque,
                                  controller_type,
                                  kp, ki, kd, J,
                                  mpc_W_speed, mpc_W_dcost, mpc_overshoot,
                                  sid, progress=gr.Progress(track_tqdm=True)):
        update_user_activity(sid)
        fig, summary = run_motor_design(
            sim_time, dt, target_rpm, load_torque,
            controller_type,
            kp, ki, kd, J,
            mpc_W_speed, mpc_W_dcost, mpc_overshoot / 100.0,
            progress=progress
        )
        return fig, summary, get_online_status_html()

    case_demo_ui["mot_run_button"].click(
        fn=run_motor_design_wrapper,
        inputs=[
            case_demo_ui["mot_sim_time"], case_demo_ui["mot_dt"],
            case_demo_ui["mot_target_rpm"], case_demo_ui["mot_load_torque"],
            case_demo_ui["mot_controller_type"],
            case_demo_ui["mot_kp"], case_demo_ui["mot_ki"], case_demo_ui["mot_kd"],
            case_demo_ui["mot_J"],
            case_demo_ui["mot_mpc_W_speed"], case_demo_ui["mot_mpc_W_dcost"],
            case_demo_ui["mot_mpc_overshoot"],
            session_id
        ],
        outputs=[case_demo_ui["mot_plot"], case_demo_ui["mot_summary"], online_counter]
    )

    # --- 阶段三：能量管理策略设计（自动引用前两阶段控制器参数）---
    def run_hybrid_demo_wrapper(sim_time, dt, initial_soc, initial_engine_power,
                                 profile,
                                 eng_ctrl_type, eng_kp, eng_ki, eng_kd,
                                 eng_mpc_horizon, eng_mpc_W_power, eng_mpc_W_dcost, eng_mpc_overshoot,
                                 mot_ctrl_type, mot_kp, mot_ki, mot_kd, mot_J,
                                 mot_mpc_W_speed, mot_mpc_W_dcost, mot_mpc_overshoot,
                                 soc_target, soc_low, soc_high,
                                 p_eng_min, p_eng_max, p_charge, k_soc,
                                 power_reserve, battery_capacity, sid,
                                 progress=gr.Progress(track_tqdm=True)):
        update_user_activity(sid)
        fig, summary, table_data = run_hybrid_demo(
            sim_time, dt, initial_soc, initial_engine_power, profile,
            eng_ctrl_type, eng_kp, eng_ki, eng_kd,
            eng_mpc_horizon, eng_mpc_W_power, eng_mpc_W_dcost, eng_mpc_overshoot / 100.0,
            mot_ctrl_type, mot_kp, mot_ki, mot_kd, mot_J,
            mot_mpc_W_speed, mot_mpc_W_dcost, mot_mpc_overshoot / 100.0,
            soc_target, soc_low, soc_high,
            p_eng_min, p_eng_max, p_charge, k_soc,
            power_reserve, battery_capacity,
            progress=progress
        )
        return fig, summary, table_data, get_online_status_html()

    case_demo_ui["hybrid_run_button"].click(
        fn=run_hybrid_demo_wrapper,
        inputs=[
            case_demo_ui["sim_time"], case_demo_ui["dt"],
            case_demo_ui["initial_soc"], case_demo_ui["initial_engine_power"],
            case_demo_ui["profile"],
            # 发动机控制器参数
            case_demo_ui["eng_controller_type"],
            case_demo_ui["eng_kp"], case_demo_ui["eng_ki"], case_demo_ui["eng_kd"],
            case_demo_ui["eng_mpc_horizon"], case_demo_ui["eng_mpc_W_power"],
            case_demo_ui["eng_mpc_W_dcost"], case_demo_ui["eng_mpc_overshoot"],
            # 电机控制器参数
            case_demo_ui["mot_controller_type"],
            case_demo_ui["mot_kp"], case_demo_ui["mot_ki"], case_demo_ui["mot_kd"],
            case_demo_ui["mot_J"],
            case_demo_ui["mot_mpc_W_speed"], case_demo_ui["mot_mpc_W_dcost"],
            case_demo_ui["mot_mpc_overshoot"],
            # 能量管理策略参数
            case_demo_ui["soc_target"], case_demo_ui["soc_low"], case_demo_ui["soc_high"],
            case_demo_ui["p_eng_min"], case_demo_ui["p_eng_max"],
            case_demo_ui["p_charge"], case_demo_ui["k_soc"],
            case_demo_ui["power_reserve"], case_demo_ui["battery_capacity"],
            session_id
        ],
        outputs=[
            case_demo_ui["hybrid_plot"], case_demo_ui["hybrid_summary"],
            case_demo_ui["hybrid_table"], online_counter
        ]
    )

    # --- 聊天机器人事件 ---
    async def chat_wrapper(message, history, sid):
        update_user_activity(sid)
        # chat_with_ai 是一个生成器，Gradio可以直接处理
        async for response in chat_with_ai(message, history):
            yield response

    chatbot_ui["send_button"].click(
        fn=chat_wrapper,
        inputs=[chatbot_ui["chat_input"], chatbot_ui["chatbot"], session_id],
        outputs=chatbot_ui["chatbot"]
    ).then(lambda: ("", get_online_status_html()), outputs=[chatbot_ui["chat_input"], online_counter])

    chatbot_ui["chat_input"].submit(
        fn=chat_wrapper,
        inputs=[chatbot_ui["chat_input"], chatbot_ui["chatbot"], session_id],
        outputs=chatbot_ui["chatbot"]
    ).then(lambda: ("", get_online_status_html()), outputs=[chatbot_ui["chat_input"], online_counter])
    
    def clear_chat_wrapper(sid):
        update_user_activity(sid)
        return [], get_online_status_html()

    chatbot_ui["clear_button"].click(
        fn=clear_chat_wrapper,
        inputs=[session_id],
        outputs=[chatbot_ui["chatbot"], online_counter]
    )
    
    # --- 页面加载和定时器事件 ---
    demo.load(fn=lambda: get_online_status_html(), outputs=[online_counter])

    gr.Timer(10).tick(fn=get_online_status_html, outputs=online_counter)

    # 系统资源监控定时刷新（每 3 秒）
    gr.Timer(3).tick(fn=get_system_monitor_html, outputs=system_monitor)


if __name__ == "__main__":
    demo.queue().launch(
        server_name=config.SERVER_NAME,
        server_port=config.SERVER_PORT,
        share=config.SHARE
    )