"""
plot_wave.py
============
波形与能量动态图：读取 display.txt，绘制原子位移波形
（纵波 + 2 个横波）和系统能量/输入功率随时间变化的二维动画。

用法：
    python plot_wave.py                          # 使用 dynamics 根目录下 output/
    python plot_wave.py examples/case05/output    # 指定案例输出目录
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import os
import sys
import json

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import compute


def load_disp_data(output_dir):
    """加载 display.npz（优先）或 display.txt。"""
    npz_path = os.path.join(output_dir, "display.npz")
    txt_path = os.path.join(output_dir, "display.txt")
    if os.path.exists(npz_path):
        return compute.load_display_npz(npz_path)
    if os.path.exists(txt_path):
        return compute.load_display_txt(txt_path)
    raise FileNotFoundError(f"找不到 display.npz 或 display.txt in {output_dir}")


def _header_json(header_fields, key, default):
    raw = header_fields.get(key, "")
    if not raw:
        return default
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return default


def _load_wave_dataset(output_dir):
    """Load wave/energy plotting data from display metadata or sibling input files."""
    disp_data = load_disp_data(output_dir)
    header = disp_data["header_fields"]

    x = disp_data["frames_x"]
    y = disp_data["frames_y"]
    z = disp_data["frames_z"]
    vx = disp_data["frames_vx"]
    vy = disp_data["frames_vy"]
    vz = disp_data["frames_vz"]
    atom_ids = np.array(disp_data["atom_ids"], dtype=np.int64)
    n_frames = x.shape[0]
    dt = float(header.get("DT", 0.001))
    nstep = int(header.get("NSTEP", 1))
    t = np.arange(n_frames, dtype=np.float64) * dt * nstep

    masses = np.array(_header_json(header, "atom_masses", []), dtype=np.float64)
    pos_0 = np.array(_header_json(header, "atom_positions", []), dtype=np.float64)
    bond_pairs = np.array(_header_json(header, "bond_pairs", []), dtype=np.int64)
    bond_stiffness = np.array(_header_json(header, "bond_stiffness", []), dtype=np.float64)
    bond_rest_lengths = np.array(_header_json(header, "bond_rest_lengths", []), dtype=np.float64)
    gravity_vec = _header_json(header, "G", [0.0, 0.0, 0.0])

    # Backward-compatible fallback for older display.txt outputs.
    if masses.size == 0 or pos_0.size == 0:
        input_dir = os.path.join(os.path.dirname(output_dir), "input")
        coord_path = os.path.join(input_dir, "coord.txt")
        if os.path.exists(coord_path):
            (_, masses_fb, _, positions_fb, _, _) = compute.load_coord_file(coord_path)
            masses = np.array(masses_fb, dtype=np.float64)
            pos_0 = np.array(positions_fb, dtype=np.float64)
        else:
            raise ValueError("display.txt 缺少 atom_masses/atom_positions 元数据，且未找到 input/coord.txt")

    if bond_pairs.size == 0:
        input_dir = os.path.join(os.path.dirname(output_dir), "input")
        connection_path = os.path.join(input_dir, "connection.txt")
        bond_path = os.path.join(input_dir, "bond.txt")
        if os.path.exists(connection_path) and os.path.exists(bond_path):
            bond_map = compute.load_bond_parameters(bond_path)
            pairs_fb, _, stiffness_fb, rest_lengths_fb = compute.load_bond_connections(
                connection_path, atom_ids, pos_0, bond_map)
            bond_pairs = np.array(pairs_fb, dtype=np.int64)
            bond_stiffness = np.array(stiffness_fb, dtype=np.float64)
            bond_rest_lengths = np.array(rest_lengths_fb, dtype=np.float64)

    return {
        "n_frames": n_frames,
        "t": t,
        "x": x,
        "y": y,
        "z": z,
        "vx": vx,
        "vy": vy,
        "vz": vz,
        "pos_0": pos_0,
        "masses": masses,
        "atom_ids": atom_ids,
        "bond_pairs": bond_pairs,
        "bond_stiffness": bond_stiffness,
        "bond_rest_lengths": bond_rest_lengths,
        "gravity_field": int(header.get("gravity_field", 0)),
        "gravity_interaction": int(header.get("gravity_interaction", 0)),
        "gravity_strength": float(header.get("gravity_strength", 1.0)),
        "G": gravity_vec,
        "driving_force": int(header.get("driving_force", 0)),
        "display_amp": _parse_display_amp(header.get("display_amp", "")),
    }


def _parse_display_amp(raw):
    if not raw or not str(raw).strip():
        return np.ones(3)
    try:
        import ast as _ast
        v = np.array(_ast.literal_eval(str(raw).strip()), dtype=np.float64)
        return v if v.shape == (3,) else np.ones(3)
    except Exception:
        return np.ones(3)


def compute_energy(x, y, z, vx, vy, vz, masses, mass_arr,
                   bond_pairs, bond_stiffness, bond_rest_lengths,
                   gravity_field, G, gravity_interaction, gravity_strength):
    """计算系统各能量分量。

    Returns:
        ek_sys:  系统动能 (n_frames,)
        us_sys:  系统弹性势能 (n_frames,)
        ug_sys:  系统重力势能 (n_frames,)
        ugr_sys: 系统万有引力势能 (n_frames,)
    """
    n_frames = x.shape[0]
    masses_2d = masses[np.newaxis, :]  # (1, n_atoms)

    # 动能 Ek = ½ m v²
    ek = 0.5 * masses_2d * (vx**2 + vy**2 + vz**2)
    ek_sys = np.sum(ek, axis=1)

    # 弹性势能 Us = ½ k (d - d₀)²
    us_sys = np.zeros(n_frames)
    if bond_pairs is not None and len(bond_pairs) > 0:
        for b_idx in range(len(bond_pairs)):
            i, j = bond_pairs[b_idx]
            dx = x[:, j] - x[:, i]
            dy = y[:, j] - y[:, i]
            dz = z[:, j] - z[:, i]
            dist = np.sqrt(dx**2 + dy**2 + dz**2)
            stretch = dist - bond_rest_lengths[b_idx]
            us_sys += 0.5 * bond_stiffness[b_idx] * stretch**2

    # 均匀重力场势能 Ug = -m G·r
    ug_sys = np.zeros(n_frames)
    if gravity_field:
        G_vec = np.array(G)
        ug_sys = -masses_2d * (G_vec[0] * x + G_vec[1] * y + G_vec[2] * z)
        ug_sys = np.sum(ug_sys, axis=1)

    # 万有引力势能 Ug_grav = -G_grav Σ m_i m_j / r
    ugr_sys = np.zeros(n_frames)
    if gravity_interaction:
        n_atoms = len(masses)
        # 为避免巨大计算量，仅当原子数较少时计算
        if n_atoms <= 200:
            for i in range(n_atoms):
                for j in range(i + 1, n_atoms):
                    dx = x[:, j] - x[:, i]
                    dy = y[:, j] - y[:, i]
                    dz = z[:, j] - z[:, i]
                    dist = np.sqrt(dx**2 + dy**2 + dz**2)
                    dist = np.maximum(dist, 1e-12)
                    pair_pe = -gravity_strength * masses[i] * masses[j] / dist
                    ugr_sys += pair_pe

    return ek_sys, us_sys, ug_sys, ugr_sys


def _load_driver_info(output_dir):
    """从 input/driver.txt 读取驱动原子的 atom_id, amp, freq（仅用于能量计算）。
    返回 dict: {atom_id: {'amp': [ax,ay,az], 'freq': [fx,fy,fz]}}，失败返回 {}。
    """
    input_dir = os.path.join(os.path.dirname(output_dir), "input")
    path = os.path.join(input_dir, "driver.txt")
    if not os.path.exists(path):
        return {}
    drivers = {}
    try:
        with open(path, encoding="utf-8") as f:
            f.readline()   # skip header
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                parts = line.split()
                if len(parts) < 10:
                    continue
                n = int(parts[0])
                amp  = np.array([float(parts[1]), float(parts[2]), float(parts[3])])
                freq = np.array([float(parts[4]), float(parts[5]), float(parts[6])])
                drivers[n] = {"amp": amp, "freq": freq}
    except Exception:
        pass
    return drivers


def compute_per_atom_energy(x, y, z, vx, vy, vz, masses,
                            bond_pairs, bond_stiffness, bond_rest_lengths,
                            atom_ids, driver_info):
    """计算每帧每个粒子的动能、势能、总能。

    势能分配规则：
    - 非驱动粒子之间的键：势能各分一半
    - 键的一端是驱动粒子：势能全部归非驱动端
    - 驱动粒子自身：按简谐振子 PE = ½ m (2πf)² A² cos²(2πft+φ) 计算
      （此处用 KE_driven = ½m·v² 的互补式：PE_driven = E_total_sho - KE_driven）

    Returns:
        ek_atom: (n_frames, n_atoms)  每原子动能
        pe_atom: (n_frames, n_atoms)  每原子势能
        et_atom: (n_frames, n_atoms)  每原子总能
    """
    n_frames, n_atoms = x.shape

    # 动能 per atom
    masses_2d = masses[np.newaxis, :]
    ek_atom = 0.5 * masses_2d * (vx**2 + vy**2 + vz**2)  # (n_frames, n_atoms)

    # 构建 atom_id → index 映射，以及驱动原子 index 集合
    id_to_idx = {int(aid): i for i, aid in enumerate(atom_ids)}
    driven_idx = set()
    for aid in driver_info:
        if aid in id_to_idx:
            driven_idx.add(id_to_idx[aid])

    # 势能 per atom（弹簧键）
    pe_atom = np.zeros((n_frames, n_atoms))
    if bond_pairs is not None and len(bond_pairs) > 0:
        for b in range(len(bond_pairs)):
            i, j = int(bond_pairs[b, 0]), int(bond_pairs[b, 1])
            ddx = x[:, j] - x[:, i]
            ddy = y[:, j] - y[:, i]
            ddz = z[:, j] - z[:, i]
            dist = np.sqrt(ddx**2 + ddy**2 + ddz**2)
            bond_pe = 0.5 * bond_stiffness[b] * (dist - bond_rest_lengths[b])**2
            i_driven = i in driven_idx
            j_driven = j in driven_idx
            if i_driven and not j_driven:
                pe_atom[:, j] += bond_pe   # 全归非驱动端
            elif j_driven and not i_driven:
                pe_atom[:, i] += bond_pe   # 全归非驱动端
            else:
                pe_atom[:, i] += bond_pe * 0.5
                pe_atom[:, j] += bond_pe * 0.5

    # 驱动粒子：用简谐振子总能 E_sho = ½m(2πf)²A²，PE_sho = E_sho - KE
    TWO_PI = 2.0 * np.pi
    for aid, info in driver_info.items():
        if aid not in id_to_idx:
            continue
        idx = id_to_idx[aid]
        m = masses[idx]
        amp  = info["amp"]
        freq = info["freq"]
        e_sho = 0.5 * m * np.sum((TWO_PI * freq)**2 * amp**2)
        pe_sho = e_sho - ek_atom[:, idx]
        pe_atom[:, idx] = np.maximum(pe_sho, 0.0)   # SHO 势能非负

    et_atom = ek_atom + pe_atom
    return ek_atom, pe_atom, et_atom


def compute_energy_flux(x, y, z, vx, vy, vz,
                        bond_pairs, bond_stiffness, bond_rest_lengths):
    """计算每帧每根键的能流密度（Hardy 公式）。

    对键 b 连接原子 i, j（j > i，方向 i→j）：

        J_b = ½ · F_{b,i} · (v_i + v_j)

    其中 F_{b,i} = k(d - r₀) * (r_j - r_i)/d 是键对原子 i 的弹力矢量，
    点乘取两端速度均值。

    - J > 0：能量从 i 流向 j（沿键方向正流）
    - J = 0：驻波，能量不流动
    - 沿链从左到右，J 的分布揭示能量传播方向

    Returns:
        flux: (n_frames, n_bonds)  每帧每键的能流（标量）
        bond_xpos: (n_bonds,)  各键中点的初始 x 坐标（用于绘图横轴）
    """
    if bond_pairs is None or len(bond_pairs) == 0:
        n_frames = x.shape[0]
        return np.zeros((n_frames, 0)), np.zeros(0)

    n_bonds = len(bond_pairs)
    n_frames = x.shape[0]
    flux = np.zeros((n_frames, n_bonds))

    # 键中点初始 x 坐标（用于横轴定位）
    bond_xpos = np.array([
        0.5 * (x[0, bond_pairs[b, 0]] + x[0, bond_pairs[b, 1]])
        for b in range(n_bonds)
    ])

    for b in range(n_bonds):
        i, j = int(bond_pairs[b, 0]), int(bond_pairs[b, 1])
        k  = bond_stiffness[b]
        r0 = bond_rest_lengths[b]

        # 键矢量与长度
        dx_ = x[:, j] - x[:, i]
        dy_ = y[:, j] - y[:, i]
        dz_ = z[:, j] - z[:, i]
        d   = np.sqrt(dx_**2 + dy_**2 + dz_**2)
        d   = np.maximum(d, 1e-12)

        # 弹力矢量（作用于原子 i，指向 j 方向）
        fac = k * (d - r0) / d        # 标量因子
        fx  = fac * dx_
        fy  = fac * dy_
        fz  = fac * dz_

        # 能流 = F_i · (v_i + v_j) / 2
        flux[:, b] = 0.5 * (fx * (vx[:, i] + vx[:, j])
                           + fy * (vy[:, i] + vy[:, j])
                           + fz * (vz[:, i] + vz[:, j]))

    return flux, bond_xpos


def compute_driver_work_power(x, y, z, vx, vy, vz,
                              bond_pairs, bond_stiffness, bond_rest_lengths,
                              atom_ids, driver_info):
    """计算每个驱动原子通过键对系统（非驱动原子）做功的功率。

    对于驱动原子 d 与系统原子 j 之间的键：
        P_{d→j} = F_{d→j} · v_j
    其中 F_{d→j} 是键对系统原子 j 的弹簧力。

    Returns:
        drv_powers: dict {atom_id: (n_frames,)}  每个驱动原子的瞬时功率
        total_power: (n_frames,)                 所有驱动原子功率之和
    """
    if bond_pairs is None or len(bond_pairs) == 0:
        n_frames = x.shape[0]
        return {}, np.zeros(n_frames)

    id_to_idx = {int(aid): i for i, aid in enumerate(atom_ids)}
    driven_idx = {id_to_idx[aid] for aid in driver_info if aid in id_to_idx}
    n_frames = x.shape[0]

    drv_powers = {}

    for b in range(len(bond_pairs)):
        ii, jj = int(bond_pairs[b, 0]), int(bond_pairs[b, 1])
        i_drv = ii in driven_idx
        j_drv = jj in driven_idx
        if i_drv == j_drv:   # 两端同为驱动或同为自由，跳过
            continue

        drv_loc = ii if i_drv else jj   # 驱动端 index
        sys_loc = jj if i_drv else ii   # 系统端 index
        drv_aid = int(atom_ids[drv_loc])

        dx_ = x[:, jj] - x[:, ii]
        dy_ = y[:, jj] - y[:, ii]
        dz_ = z[:, jj] - z[:, ii]
        dist = np.sqrt(dx_**2 + dy_**2 + dz_**2)
        dist = np.maximum(dist, 1e-12)

        k  = bond_stiffness[b]
        r0 = bond_rest_lengths[b]
        fac = k * (dist - r0) / dist   # 标量弹力因子

        # 作用于系统原子的弹簧力：指向驱动原子方向
        if i_drv:   # drv=i, sys=j: 力方向 j→i，即 -(dx_, dy_, dz_)
            fx = -fac * dx_
            fy = -fac * dy_
            fz = -fac * dz_
        else:       # drv=j, sys=i: 力方向 i→j，即 +(dx_, dy_, dz_)
            fx = fac * dx_
            fy = fac * dy_
            fz = fac * dz_

        power = fx * vx[:, sys_loc] + fy * vy[:, sys_loc] + fz * vz[:, sys_loc]

        if drv_aid not in drv_powers:
            drv_powers[drv_aid] = np.zeros(n_frames)
        drv_powers[drv_aid] += power

    total_power = sum(drv_powers.values()) if drv_powers else np.zeros(n_frames)
    return drv_powers, total_power


def plot_wave(output_dir, save_gif=False, save_mp4=False, show=True):
    """主绘图函数：读取 display.txt 并生成波形+能量动画。

    布局（4行×1列，纵向排列）：
      行0：x/y/z 位移波形叠加在同一子图（vs 原子序号）
      行1：每粒子动能、势能、总能叠加在同一子图
      行2：键能流密度 J（Hardy 公式，vs 键中点位置）
      行3：系统总能量随时间变化

    Args:
        output_dir: 输出目录（含 display.npz 或 display.txt）
        save_gif:   是否保存 GIF
        save_mp4:   是否保存 MP4
        show:       是否弹出交互窗口
    """
    data = _load_wave_dataset(output_dir)

    n_frames = int(data["n_frames"])
    t = np.array(data["t"])

    x  = np.array(data["x"])
    y  = np.array(data["y"])
    z  = np.array(data["z"])
    vx = np.array(data["vx"])
    vy = np.array(data["vy"])
    vz = np.array(data["vz"])

    pos_0    = np.array(data["pos_0"])
    masses   = np.array(data["masses"])
    atom_ids = np.array(data["atom_ids"])
    n_atoms  = len(atom_ids)

    bond_pairs       = np.array(data.get("bond_pairs",       []), dtype=np.int64)
    bond_stiffness   = np.array(data.get("bond_stiffness",   []), dtype=np.float64)
    bond_rest_lengths= np.array(data.get("bond_rest_lengths",[]), dtype=np.float64)

    gravity_field       = int(data.get("gravity_field", 0))
    gravity_interaction = int(data.get("gravity_interaction", 0))
    G                   = data.get("G", [0, 0, 0])
    gravity_strength    = float(data.get("gravity_strength", 1.0))
    driving_force       = int(data.get("driving_force", 0))

    # 驱动原子信息（用于势能计算）
    driver_info = _load_driver_info(output_dir) if driving_force else {}

    # ── 位移 ──
    dx = x - pos_0[np.newaxis, :, 0]
    dy = y - pos_0[np.newaxis, :, 1]
    dz = z - pos_0[np.newaxis, :, 2]

    # ── 每粒子能量（图2 与图3 共用同一套计算）──
    ek_atom, pe_atom, et_atom = compute_per_atom_energy(
        x, y, z, vx, vy, vz, masses,
        bond_pairs, bond_stiffness, bond_rest_lengths,
        atom_ids, driver_info)

    # ── 系统总能量 = 各粒子求和（与图2 完全一致）──
    ek_sys  = np.sum(ek_atom, axis=1)
    us_sys  = np.sum(pe_atom, axis=1)
    e_total = np.sum(et_atom, axis=1)
    power   = np.gradient(e_total, t)
    # 重力势能：仍用原有函数提供（若启用重力场）
    _, _, ug_sys, ugr_sys = compute_energy(
        x, y, z, vx, vy, vz, masses, masses,
        bond_pairs, bond_stiffness, bond_rest_lengths,
        gravity_field, G, gravity_interaction, gravity_strength)
    if gravity_field or gravity_interaction:
        e_total = e_total + ug_sys + ugr_sys
        power   = np.gradient(e_total, t)

    # ── 能流密度 ──
    flux, bond_xpos = compute_energy_flux(
        x, y, z, vx, vy, vz,
        bond_pairs, bond_stiffness, bond_rest_lengths)

    # ── 驱动做功功率 ──
    drv_powers, total_drv_power = compute_driver_work_power(
        x, y, z, vx, vy, vz,
        bond_pairs, bond_stiffness, bond_rest_lengths,
        atom_ids, driver_info)

    # ── 原子可视化预计算 ──
    display_amp = np.array(data.get("display_amp", [1.0, 1.0, 1.0]), dtype=np.float64)
    eq_x_vis = pos_0[:, 0]
    eq_z_vis = pos_0[:, 2]
    # 视觉坐标 = 平衡位置 + 放大的位移
    x_vis = eq_x_vis + (x - eq_x_vis) * display_amp[0]   # (n_frames, n_atoms)
    z_vis = eq_z_vis + (z - eq_z_vis) * display_amp[2]

    # 找边界原子（与驱动原子成键的系统原子）及对应键
    id_to_idx_vis = {int(aid): i for i, aid in enumerate(atom_ids)}
    driven_set_vis = {id_to_idx_vis[aid] for aid in driver_info if aid in id_to_idx_vis}
    bond_boundary_list = []   # (drv_idx, sys_idx, bond_b)
    for _b in range(len(bond_pairs)):
        _ii, _jj = int(bond_pairs[_b, 0]), int(bond_pairs[_b, 1])
        if (_ii in driven_set_vis) ^ (_jj in driven_set_vis):
            _drv = _ii if _ii in driven_set_vis else _jj
            _sys = _jj if _ii in driven_set_vis else _ii
            bond_boundary_list.append((_drv, _sys, _b))

    # 唯一边界原子索引列表
    _bnd_set = {}
    for _drv, _sys, _b in bond_boundary_list:
        if _sys not in _bnd_set:
            _bnd_set[_sys] = len(_bnd_set)
    boundary_atom_idx = np.array(list(_bnd_set.keys()), dtype=int)
    n_boundary = len(boundary_atom_idx)

    _lat = (eq_x_vis[-1] - eq_x_vis[0]) / max(n_atoms - 1, 1)
    _z_all = z_vis.reshape(-1)
    _z_min, _z_max = np.min(_z_all), np.max(_z_all)
    _z_mg = max((_z_max - _z_min) * 0.2, _lat * 2)
    _z_range = max((_z_max - _z_min) + 2 * _z_mg, _lat * 4)
    _arrow_len = _z_range * 0.50   # 箭头最大显示长度 = 纵坐标范围的 50%

    # 预计算边界原子受到的驱动力
    # 方向：沿显示坐标下的键方向（消除坐标轴比例失真）；大小：胡克力模 k|d-r0|
    bnd_fx_scaled = np.zeros((n_frames, max(n_boundary, 1)))
    bnd_fz_scaled = np.zeros((n_frames, max(n_boundary, 1)))
    _f_mag_all = []
    for _drv, _sys, _b in bond_boundary_list:
        _bi = _bnd_set[_sys]
        # 物理键长
        _dx3 = x[:, _drv] - x[:, _sys]
        _dy3 = y[:, _drv] - y[:, _sys]
        _dz3 = z[:, _drv] - z[:, _sys]
        _dist = np.maximum(np.sqrt(_dx3**2 + _dy3**2 + _dz3**2), 1e-12)
        # 有符号力大小（正 = 拉向驱动原子，负 = 推离）
        _f_signed = bond_stiffness[_b] * (_dist - bond_rest_lengths[_b])
        # 显示坐标下的键方向（x-z 平面）
        _dx_d = x_vis[:, _drv] - x_vis[:, _sys]
        _dz_d = z_vis[:, _drv] - z_vis[:, _sys]
        _disp_len = np.maximum(np.sqrt(_dx_d**2 + _dz_d**2), 1e-12)
        bnd_fx_scaled[:, _bi] += _f_signed * _dx_d / _disp_len
        bnd_fz_scaled[:, _bi] += _f_signed * _dz_d / _disp_len
        _f_mag_all.append(np.abs(_f_signed))
    _f_max = np.max(_f_mag_all) if _f_mag_all else 1.0
    _f_max = _f_max if _f_max > 1e-20 else 1.0
    bnd_fx_scaled = bnd_fx_scaled / _f_max * _arrow_len
    bnd_fz_scaled = bnd_fz_scaled / _f_max * _arrow_len

    # 边界原子速度（方向沿实际速度，大小归一化）
    bnd_vx_raw = vx[:, boundary_atom_idx] if n_boundary > 0 else np.zeros((n_frames, 1))
    bnd_vz_raw = vz[:, boundary_atom_idx] if n_boundary > 0 else np.zeros((n_frames, 1))
    _v_max = np.max(np.sqrt(bnd_vx_raw**2 + bnd_vz_raw**2)) if n_boundary > 0 else 1.0
    _v_max = _v_max if _v_max > 1e-20 else 1.0
    bnd_vx_scaled = bnd_vx_raw / _v_max * _arrow_len
    bnd_vz_scaled = bnd_vz_raw / _v_max * _arrow_len

    # y 轴范围 ──
    def get_ylim(arr):
        vmax = np.max(np.abs(arr))
        if vmax < 1e-10:
            return -1.0, 1.0
        m = vmax * 0.2
        return -vmax - m, vmax + m

    def get_ylim_pos(arr):
        vmax = np.max(arr)
        if vmax < 1e-12:
            return 0.0, 1.0
        return 0.0, vmax * 1.2

    # 共用 y 轴范围：位移图取三个方向最大值统一
    disp_vmax = max(np.max(np.abs(dx)), np.max(np.abs(dy)), np.max(np.abs(dz)))
    disp_vmax = disp_vmax if disp_vmax > 1e-10 else 1.0
    disp_ylim = (-disp_vmax * 1.2, disp_vmax * 1.2)

    # 每粒子能量：取三者最大值统一 y 轴
    energy_vmax = max(np.max(ek_atom), np.max(pe_atom), np.max(et_atom))
    energy_vmax = energy_vmax if energy_vmax > 1e-12 else 1.0
    energy_ylim = (0.0, energy_vmax * 1.2)

    e_max = max(np.max(e_total), 1e-12)
    e_min = min(np.min(e_total), 0.0)
    p_max = max(np.percentile(np.abs(power), 95) if len(power) > 0 else 0.0, 0.0)

    # 能流 y 轴范围（对称，正负各半）
    if flux.size > 0:
        flux_vmax = np.max(np.abs(flux))
        flux_vmax = flux_vmax if flux_vmax > 1e-12 else 1.0
        flux_ylim = (-flux_vmax * 1.2, flux_vmax * 1.2)
    else:
        flux_ylim = (-1.0, 1.0)

    atom_idx = np.arange(n_atoms)

    # ── 驱动/非驱动粒子能量（右下图）──
    id_to_idx  = {int(aid): i for i, aid in enumerate(atom_ids)}
    driven_idx = np.array([id_to_idx[aid] for aid in driver_info if aid in id_to_idx], dtype=int)
    free_idx   = np.setdiff1d(np.arange(n_atoms), driven_idx)
    has_driver = len(driven_idx) > 0

    if has_driver:
        ek_drv  = np.sum(ek_atom[:, driven_idx], axis=1)
        ep_drv  = np.sum(pe_atom[:, driven_idx], axis=1)
        ek_free = np.sum(ek_atom[:, free_idx],   axis=1)
        ep_free = np.sum(pe_atom[:, free_idx],   axis=1)

    # ── 图形布局：左3行、右3行（subplot_mosaic）──
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False

    from matplotlib.collections import LineCollection as _LC

    fig, axes = plt.subplot_mosaic(
        [['atoms',  'ep'],
         ['wave',   'drv'],
         ['energy', 'pwr']],
        figsize=(20, 15))
    ax_atoms  = axes['atoms']
    ax_wave   = axes['wave']
    ax_ep     = axes['ep']
    ax_energy = axes['energy']
    ax_drv    = axes['drv']
    ax_pwr    = axes['pwr']
    fig.subplots_adjust(hspace=0.45, wspace=0.32, top=0.97)

    # ── 左上：原子位置 + 键 + 力/速度箭头 ──
    _x_min, _x_max = eq_x_vis[0], eq_x_vis[-1]
    ax_atoms.set_xlim(_x_min - _lat, _x_max + _lat)
    ax_atoms.set_ylim(_z_min - _z_mg, _z_max + _z_mg)
    ax_atoms.set_xlabel("位置 $x$")
    ax_atoms.set_ylabel("位移 $z$（放大 {:.0f}×）".format(display_amp[2]))
    ax_atoms.set_title("原子运动（红=驱动，箭头：红=驱动力，蓝=边界速度）")
    ax_atoms.set_aspect('auto')
    ax_atoms.grid(True, alpha=0.2)

    # 键线段（LineCollection，初始帧）
    def _make_bond_segs(frame_idx):
        segs = []
        for _b in range(len(bond_pairs)):
            _ii, _jj = int(bond_pairs[_b, 0]), int(bond_pairs[_b, 1])
            segs.append([(x_vis[frame_idx, _ii], z_vis[frame_idx, _ii]),
                         (x_vis[frame_idx, _jj], z_vis[frame_idx, _jj])])
        return segs

    _bond_lc = _LC(_make_bond_segs(0), colors='#888888', linewidths=0.8, zorder=1)
    ax_atoms.add_collection(_bond_lc)

    # 散点：自由原子（黑）
    _free_mask = np.array([i not in driven_set_vis for i in range(n_atoms)])
    _scat_free, = ax_atoms.plot(
        x_vis[0, _free_mask], z_vis[0, _free_mask],
        'o', color='black', ms=4, zorder=3)

    # 散点：驱动原子（红）
    _drv_mask = ~_free_mask
    _scat_drv, = ax_atoms.plot(
        x_vis[0, _drv_mask], z_vis[0, _drv_mask],
        'o', color='red', ms=6, zorder=4)

    # 力箭头（红，边界原子）
    _q_force = ax_atoms.quiver(
        x_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
        z_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
        bnd_fx_scaled[0] if n_boundary > 0 else [],
        bnd_fz_scaled[0] if n_boundary > 0 else [],
        color='red', angles='xy', scale_units='xy', scale=1,
        width=0.007, headwidth=5, headlength=5, zorder=5)

    # 速度箭头（蓝，边界原子）
    _q_vel = ax_atoms.quiver(
        x_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
        z_vis[0, boundary_atom_idx] if n_boundary > 0 else [],
        bnd_vx_scaled[0] if n_boundary > 0 else [],
        bnd_vz_scaled[0] if n_boundary > 0 else [],
        color='blue', angles='xy', scale_units='xy', scale=1,
        width=0.007, headwidth=5, headlength=5, zorder=5)

    # ── 左上：x/y/z 位移波形 ──
    ax_wave.set_xlim(0, n_atoms - 1)
    ax_wave.set_ylim(disp_ylim)
    ax_wave.set_xlabel("原子序号")
    ax_wave.set_ylabel("位移 $u$")
    ax_wave.set_title("粒子位移（$x$ / $y$ / $z$ 方向）")
    ax_wave.grid(True, alpha=0.3)

    wave_disps  = [dx, dy, dz]
    wave_labels = ["$u_x$（纵波）", "$u_y$（横波）", "$u_z$（横波）"]
    wave_colors = ["#2563eb", "#ea580c", "#16a34a"]
    wave_lines  = []
    for label, color in zip(wave_labels, wave_colors):
        ln, = ax_wave.plot([], [], color=color, linewidth=1.5, label=label)
        wave_lines.append(ln)
    ax_wave.legend(loc="upper right", fontsize=9)
    _dt_frame = (t[1] - t[0]) if len(t) > 1 else 0.0
    _t_total_str = f"{t[-1] + _dt_frame:.2f} s"
    _time_axes = [ax_atoms, ax_wave, ax_energy, ax_ep, ax_drv, ax_pwr]
    time_texts = [
        ax.text(0.02, 0.97, "", transform=ax.transAxes,
                fontsize=9, verticalalignment="top",
                bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.7))
        for ax in _time_axes
    ]
    time_text = time_texts[1]   # 保留旧名兼容下面的代码

    # ── 左下：每粒子能量（左轴）+ 能流密度（右轴）──
    ax_energy.set_xlim(0, n_atoms - 1)
    ax_energy.set_ylim(energy_ylim)
    ax_energy.set_xlabel("原子序号 / 键位置")
    ax_energy.set_ylabel("能量 $E$")
    ax_energy.set_title(
        r"每粒子能量（$E_k$/$E_p$/$E_{tot}$）与能流密度 $J$"
    )
    ax_energy.grid(True, alpha=0.3)

    energy_arrays = [ek_atom, pe_atom, et_atom]
    energy_labels = ["$E_k$（动能）", "$E_p$（势能）", "$E_{tot}$（总能）"]
    energy_colors = ["#16a34a", "#b45309", "#7c3aed"]
    energy_lines  = []
    for label, color in zip(energy_labels, energy_colors):
        ln, = ax_energy.plot([], [], color=color, linewidth=1.5, label=label)
        energy_lines.append(ln)

    ax_flux = ax_energy.twinx()
    ax_flux.set_ylim(flux_ylim)
    ax_flux.set_ylabel("能流密度 $J$", color="#dc2626")
    ax_flux.tick_params(axis='y', labelcolor="#dc2626")
    ax_flux.axhline(0, color="gray", linewidth=0.8, linestyle="--")
    flux_line, = ax_flux.plot([], [], color="#dc2626", linewidth=1.5,
                              label="$J$（能流密度）")
    handles_e, labels_e = ax_energy.get_legend_handles_labels()
    handles_f, labels_f = ax_flux.get_legend_handles_labels()
    ax_energy.legend(handles_e + handles_f, labels_e + labels_f,
                     loc="upper right", fontsize=9)

    # ── 右上：系统总能量随时间 ──
    ax_ep.set_xlim(t[0], t[-1])
    ep_margin = (e_max - e_min) * 0.15 if e_max > e_min else e_max * 0.15
    ax_ep.set_ylim(e_min - ep_margin, e_max + ep_margin)
    ax_ep.set_clip_on(True)
    ax_ep.set_xlabel("时间 $t$ (s)")
    ax_ep.set_ylabel("能量 $E$ / 功率 $P$")
    ax_ep.set_title("系统能量与输入功率")
    ax_ep.grid(True, alpha=0.3)

    ln_ek, = ax_ep.plot([], [], "b-",     lw=1.5, label="$E_k$（动能）")
    ln_us, = ax_ep.plot([], [], "orange", lw=1.5, label="$E_s$（弹性势能）")
    ln_et, = ax_ep.plot([], [], "r--",    lw=1.5, label="$E_{tot}$（总能量）")
    ln_pw, = ax_ep.plot([], [], "g-",     lw=1.5, alpha=0.7, label=r"$P_{in}=dE/dt$")
    ln_ug  = None
    ln_ugr = None
    if gravity_field:
        ln_ug, = ax_ep.plot([], [], "purple", lw=1.0, alpha=0.5, label="$E_g$（重力势能）")
    if gravity_interaction and n_atoms <= 200:
        ln_ugr, = ax_ep.plot([], [], "brown", lw=1.0, alpha=0.5, label="$E_{gr}$（万有引力势能）")
    ax_ep.legend(loc="upper right", fontsize=9)

    # ── 右下：驱动/非驱动粒子能量随时间 ──
    ax_drv.set_xlim(t[0], t[-1])
    ax_drv.set_xlabel("时间 $t$ (s)")
    ax_drv.set_ylabel("能量 $E$")
    ax_drv.grid(True, alpha=0.3)
    ln_ek_drv = ln_ep_drv = ln_ek_free = ln_ep_free = None
    ln_et_drv = ln_et_free = None
    if has_driver:
        et_drv  = ek_drv  + ep_drv
        et_free = ek_free + ep_free
        drv_ids   = sorted(driver_info.keys())
        ax_drv.set_title(f"驱动粒子（序号 {drv_ids}）向系统做功")
        ln_ek_drv,  = ax_drv.plot([], [], color="#dc2626", lw=1.2, linestyle="--",
                                  label=r"$E_k^{drv}$（驱动动能）")
        ln_ep_drv,  = ax_drv.plot([], [], color="#f97316", lw=1.2, linestyle="--",
                                  label=r"$E_p^{drv}$（驱动势能）")
        ln_et_drv,  = ax_drv.plot([], [], color="#7f1d1d", lw=2.0,
                                  label=r"$E_{tot}^{drv}$（驱动总能）")
        ln_ek_free, = ax_drv.plot([], [], color="#2563eb", lw=1.2, linestyle="--",
                                  label=r"$E_k^{sys}$（系统动能）")
        ln_ep_free, = ax_drv.plot([], [], color="#16a34a", lw=1.2, linestyle="--",
                                  label=r"$E_p^{sys}$（系统势能）")
        ln_et_free, = ax_drv.plot([], [], color="#1e3a5f", lw=2.0,
                                  label=r"$E_{tot}^{sys}$（系统总能）")
        ax_drv.legend(loc="upper right", fontsize=8)
        # y 轴一次定好
        _drv_all = np.concatenate([ek_drv, ep_drv, et_drv, ek_free, ep_free, et_free])
        _dy_max = np.max(_drv_all)
        _dy_min = np.min(_drv_all)
        _dy_mg  = (_dy_max - _dy_min) * 0.15 if _dy_max > _dy_min else abs(_dy_max) * 0.15 + 1e-12
        ax_drv.set_ylim(_dy_min - _dy_mg, _dy_max + _dy_mg)
    else:
        ax_drv.set_title("驱动粒子能量（无驱动力）")
        ax_drv.text(0.5, 0.5, "无驱动力", transform=ax_drv.transAxes,
                    ha="center", va="center", fontsize=12, color="gray")

    # ── 右下：驱动做功功率 ──
    ax_pwr.set_xlim(t[0], t[-1])
    ax_pwr.set_xlabel("时间 $t$ (s)")
    ax_pwr.set_ylabel("功率 $P$")
    ax_pwr.set_title("驱动原子对系统做功的功率 $P = \\mathbf{F}_{bond}\\cdot\\mathbf{v}_{sys}$")
    ax_pwr.axhline(0, color="gray", linewidth=0.8, linestyle="--")
    ax_pwr.grid(True, alpha=0.3)

    pwr_colors = ["#dc2626", "#2563eb", "#16a34a", "#f97316", "#7c3aed"]
    ln_pwr_each = {}   # aid -> Line2D
    if drv_powers:
        for idx_d, (aid, _) in enumerate(sorted(drv_powers.items())):
            color = pwr_colors[idx_d % len(pwr_colors)]
            ln, = ax_pwr.plot([], [], color=color, lw=1.2, linestyle="--",
                              label=f"$P_{{drv,{aid}}}$（原子 {aid}）")
            ln_pwr_each[aid] = ln
    ln_pwr_total, = ax_pwr.plot([], [], color="black", lw=2.0,
                                label=r"$P_{total}$（总功率）")
    ax_pwr.legend(loc="upper right", fontsize=9)

    # y 轴一次定好
    if drv_powers:
        _pw_all = np.concatenate(list(drv_powers.values()) + [total_drv_power])
        _pw_max = np.max(_pw_all)
        _pw_min = np.min(_pw_all)
        _pw_mg  = (_pw_max - _pw_min) * 0.15 if _pw_max > _pw_min else abs(_pw_max) * 0.15 + 1e-12
        ax_pwr.set_ylim(_pw_min - _pw_mg, _pw_max + _pw_mg)

    # ── 动画更新 ──
    def update(frame):
        # 每轮开始时清屏
        # 左上：原子位置动画
        _bond_lc.set_segments(_make_bond_segs(frame))
        _scat_free.set_xdata(x_vis[frame, _free_mask])
        _scat_free.set_ydata(z_vis[frame, _free_mask])
        _scat_drv.set_xdata(x_vis[frame, _drv_mask])
        _scat_drv.set_ydata(z_vis[frame, _drv_mask])
        if n_boundary > 0:
            _q_force.set_offsets(
                np.column_stack([x_vis[frame, boundary_atom_idx],
                                 z_vis[frame, boundary_atom_idx]]))
            _q_force.set_UVC(bnd_fx_scaled[frame], bnd_fz_scaled[frame])
            _q_vel.set_offsets(
                np.column_stack([x_vis[frame, boundary_atom_idx],
                                 z_vis[frame, boundary_atom_idx]]))
            _q_vel.set_UVC(bnd_vx_scaled[frame], bnd_vz_scaled[frame])

        if frame == 0:
            all_clear = list(wave_lines) + list(energy_lines) + [flux_line]
            all_clear += [ln for ln in [ln_ek, ln_us, ln_et, ln_pw, ln_ug, ln_ugr]
                          if ln is not None]
            if has_driver:
                all_clear += [ln for ln in [ln_ek_drv, ln_ep_drv, ln_et_drv,
                                            ln_ek_free, ln_ep_free, ln_et_free]
                              if ln is not None]
            all_clear += list(ln_pwr_each.values()) + [ln_pwr_total]
            for ln in all_clear:
                ln.set_data([], [])
            _tstr0 = f"t = {t[0]:.2f} s / {_t_total_str}  |  帧 1/{n_frames}"
            for _tt in time_texts:
                _tt.set_text(_tstr0)
            return all_clear + time_texts + [_bond_lc, _scat_free, _scat_drv,
                                 _q_force, _q_vel]

        # 左中：位移波形
        for i, ln in enumerate(wave_lines):
            ln.set_data(atom_idx, wave_disps[i][frame])
        _tstr = f"t = {t[frame]:.2f} s / {_t_total_str}  |  帧 {frame+1}/{n_frames}"
        for _tt in time_texts:
            _tt.set_text(_tstr)

        # 左下：每粒子能量 + 能流密度
        for i, ln in enumerate(energy_lines):
            ln.set_data(atom_idx, energy_arrays[i][frame])
        if flux.shape[1] > 0:
            flux_line.set_data(bond_xpos, flux[frame])

        # 右上：系统能量（累计）
        cur_t = t[:frame + 1]
        ln_ek.set_data(cur_t, ek_sys[:frame + 1])
        ln_us.set_data(cur_t, us_sys[:frame + 1])
        ln_et.set_data(cur_t, e_total[:frame + 1])
        ln_pw.set_data(cur_t, power[:frame + 1])
        if ln_ug:  ln_ug.set_data(cur_t, ug_sys[:frame + 1])
        if ln_ugr: ln_ugr.set_data(cur_t, ugr_sys[:frame + 1])

        # 右中：驱动/系统粒子能量（累计）
        if has_driver:
            ln_ek_drv.set_data( cur_t, ek_drv[:frame + 1])
            ln_ep_drv.set_data( cur_t, ep_drv[:frame + 1])
            ln_et_drv.set_data( cur_t, et_drv[:frame + 1])
            ln_ek_free.set_data(cur_t, ek_free[:frame + 1])
            ln_ep_free.set_data(cur_t, ep_free[:frame + 1])
            ln_et_free.set_data(cur_t, et_free[:frame + 1])

        # 右下：驱动做功功率（累计）
        for aid, ln in ln_pwr_each.items():
            ln.set_data(cur_t, drv_powers[aid][:frame + 1])
        ln_pwr_total.set_data(cur_t, total_drv_power[:frame + 1])

        artists = (wave_lines + time_texts + energy_lines +
                   [flux_line, ln_ek, ln_us, ln_et, ln_pw])
        if ln_ug:     artists.append(ln_ug)
        if ln_ugr:    artists.append(ln_ugr)
        if has_driver:
            artists += [ln_ek_drv, ln_ep_drv, ln_et_drv,
                        ln_ek_free, ln_ep_free, ln_et_free]
        artists += list(ln_pwr_each.values()) + [ln_pwr_total]
        artists += [_bond_lc, _scat_free, _scat_drv, _q_force, _q_vel]
        return artists

    ani = FuncAnimation(fig, update, frames=n_frames, interval=50, blit=True, repeat=True)

    # ── 输出文件 ──
    gif_path = None
    if save_gif:
        gif_path = os.path.join(output_dir, "wave_animation.gif")
        ani.save(gif_path, writer="pillow", fps=min(20, max(1, n_frames // 5)))
        print(f"[plot_wave] GIF 已保存: {gif_path}")

    if save_mp4:
        try:
            import matplotlib.animation as manim
            ffmpeg_path = None
            try:
                import imageio_ffmpeg
                ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
            except Exception:
                pass
            if ffmpeg_path and os.path.exists(ffmpeg_path):
                plt.rcParams['animation.ffmpeg_path'] = ffmpeg_path
            ffps = min(20, max(1, n_frames // 5))
            writer = manim.FFMpegWriter(fps=ffps, codec="libx264",
                                        extra_args=["-pix_fmt", "yuv420p"])
            mp4_path = os.path.join(output_dir, "wave_animation.mp4")
            ani.save(mp4_path, writer=writer)
            print(f"[plot_wave] MP4 已保存: {mp4_path}")
        except FileNotFoundError:
            print("[plot_wave] 警告: 未找到 ffmpeg，跳过 MP4 输出")
        except Exception as e:
            print(f"[plot_wave] 警告: MP4 输出失败 ({e})，跳过")

    if show:
        plt.show()
    else:
        plt.close(fig)
    return gif_path


if __name__ == "__main__":
    script_dir = os.path.dirname(os.path.abspath(__file__))
    if len(sys.argv) > 1:
        output_dir = os.path.abspath(sys.argv[1])
    else:
        output_dir = compute.get_output_dir(script_dir)
    plot_wave(output_dir)
