"""
compute.py
----------
纯 Python 计算脚本，不依赖 VisPy。

功能：
  1. 运行 NT 步物理模拟（ kinematics / dynamics 等运动模式）
  2. 按 NSTEP 抽帧，输出 output/display.txt（新文本格式）
  3. 可选（save_trajectory=1）保留完整轨迹 output/trajectory.txt（JSON）
"""

import json
import re
import numpy as np
import os
import platform
import subprocess
import sys
import time
import datetime
from tqdm import trange

# ===========================================================================
# 物理参数（必须与 draw.py 加载的边界常量保持一致）
# ===========================================================================
# 全局参数变量，将在 load_parameters 中填充
box_a = None
alpha = None
COORD_FILE = None
X0 = None
Y0 = None
Z0 = None
VX0 = None
VY0 = None
VZ0 = None
M   = None
ATOM_IDS = None
ATOM_MASSES = None
ATOM_RADII = None
ATOM_POSITIONS = None
ATOM_VELOCITIES = None
ATOM_FIXED = None
BOND_CONNECTION_FILE = None
BOND_PARAMETER_FILE = None
BOND_PAIRS = None
BOND_NAMES = None
BOND_STIFFNESS = None
BOND_REST_LENGTHS = None
PLOT_ATOM_ID = None
PLOT_ATOM_ROW = None
G   = None
B   = None
METHOD = None
NT    = None
DT    = None
NSTEP = None
warmup_steps = None    # 预热步数（跳过不保存）
sample_start = None    # 抽帧起始索引
sample_end   = None    # 抽帧结束索引
ball_radius = None
ball_color_r = None
ball_color_g = None
ball_color_b = None
box_color_r = None
box_color_g = None
box_color_b = None
use_marker = 0
camera_keyframes_raw = ""
camera_distance = 40.0
camera_elevation = 0.0
camera_azimuth = 0.0
camera_center_x = 0.0
camera_center_y = 0.0
camera_center_z = 0.0
camera_pos_x = None
camera_pos_y = None
camera_pos_z = None
display_amp_str = ""
FIXED_MASK_X = None
FIXED_MASK_Y = None
FIXED_MASK_Z = None

# 力开关
GRAVITY_FIELD = 1         # 均匀重力场
GRAVITY_INTERACTION = 0   # 原子间万有引力
ELASTIC_FORCE = 1         # 弹簧键力
DAMPING_FORCE = 0         # 阻尼
DRIVING_FORCE = 0         # 驱动力
GRAVITY_STRENGTH = 1.0

# 驱动力数据
DRIVER_DATA = None        # 加载自 driver.txt

# 派生边界（根据 box_a 计算）
X_MIN = None
X_MAX = None
Y_MIN = None
Y_MAX = None
Z_MIN = None
Z_MAX = None


def _load_camera_motion(path):
    """读取 move_camera.txt（速度段格式），返回 JSON 字符串。
    
    格式：每行是一个运动段
      start-end  vx=f1 vy=f2 vz=f3  rx=d1 ry=d2 rz=d3
    示例：
      1-60   vx=1.0  rx=10
      30-90  vy=2.0  ry=20  rz=10
    
    返回 JSON: [{"start":N,"end":N,"v":[x,y,z],"r":[x,y,z]},...]
    """
    import re
    if not os.path.exists(path):
        print(f"[compute] 警告: 未找到 {path}，跳过运动相机")
        return ""
    segments = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            # 解析帧范围：支持 "all"（全程）或 "N-M"（区间）
            if line.lower().startswith("all") or re.match(r'^\s*all\s', line, re.IGNORECASE):
                start, end = 0, 10**9
            else:
                m = re.match(r'(\d+)\s*-\s*(\d+)', line)
                if not m:
                    continue
                start, end = int(m.group(1)), int(m.group(2))
            v = [0.0, 0.0, 0.0]
            r = [0.0, 0.0, 0.0]
            # 解析 vx=, vy=, vz=
            for i, axis in enumerate(['x', 'y', 'z']):
                m2 = re.search(r'v' + axis + r'\s*=\s*([-\d.]+)', line)
                if m2:
                    v[i] = float(m2.group(1))
            # 解析 rx=, ry=, rz=
            for i, axis in enumerate(['x', 'y', 'z']):
                m2 = re.search(r'r' + axis + r'\s*=\s*([-\d.]+)', line)
                if m2:
                    r[i] = float(m2.group(1))
            if any(v) or any(r):
                segments.append({"start": start, "end": end, "v": v, "r": r})
    if not segments:
        return ""
    import json
    return json.dumps(segments)


def _to_text_value(value):
    """Convert numpy-heavy objects into JSON-friendly plain Python values."""
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        return float(value)
    if isinstance(value, dict):
        return {key: _to_text_value(val) for key, val in value.items()}
    if isinstance(value, (list, tuple)):
        return [_to_text_value(item) for item in value]
    return value


def _from_text_value(value):
    """Convert nested numeric lists back into numpy arrays when appropriate."""
    if isinstance(value, list):
        if not value:
            return np.array([], dtype=np.float64)
        if all(not isinstance(item, (list, dict)) for item in value):
            if all(isinstance(item, (int, float, bool)) for item in value):
                return np.array(value)
            return value
        if all(isinstance(item, list) for item in value):
            return np.array(value)
        return [_from_text_value(item) for item in value]
    if isinstance(value, dict):
        return {key: _from_text_value(val) for key, val in value.items()}
    return value


def save_text_data(path, data):
    """Save structured simulation data as formatted JSON text (旧格式，保留兼容)."""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(_to_text_value(data), f, ensure_ascii=False, indent=2)
    return path


def load_text_data(path):
    """Load structured simulation data from JSON text (旧格式)."""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    return _from_text_value(data)


# ========================================================================
# 新 display.txt 格式：纯文本，按帧分块
#   第1行: number of frames: N
#   第2行: number of particles: M
#   第3行: frame:  1
#   第4行: n   x            y            z           vx         vy         vz
#   第5+行: 数据行（每个原子一行）
#   重复第3-5行直到所有帧
# ========================================================================

def save_display_txt(path, frames_x, frames_y, frames_z,
                     frames_vx, frames_vy, frames_vz,
                     atom_ids, n_total_frames, n_total_particles,
                     header_fields=None):
    """Write display.txt in new text format.
    
    Args:
        path: 输出文件路径
        frames_x/y/z/vx/vy/vz: (n_frames, n_atoms) 数组
        atom_ids: (n_atoms,) 原子编号数组
        n_total_frames: 总帧数（含未采样）
        n_total_particles: 总粒子数
        header_fields: 可选的额外元数据字典（写入文件头之后）
    """
    os.makedirs(os.path.dirname(path), exist_ok=True)
    n_frames = frames_x.shape[0]
    n_atoms = frames_x.shape[1]
    
    # 格式化辅助：固定宽度，6位小数
    def fmt(v): return f"{v:13.6f}"
    
    with open(path, "w", encoding="utf-8") as f:
        f.write(f"number of frames: {n_total_frames}\n")
        f.write(f"number of particles: {n_total_particles}\n")
        # 写入额外元数据
        if header_fields:
            for k, v in header_fields.items():
                f.write(f"{k}: {v}\n")
        
        for fr in range(n_frames):
            f.write(f"\nframe:  {fr + 1}\n")
            f.write(f"n   x            y            z           vx         vy         vz\n")
            for a in range(n_atoms):
                f.write(f"{atom_ids[a]:d}"
                        f"{fmt(frames_x[fr, a])}"
                        f"{fmt(frames_y[fr, a])}"
                        f"{fmt(frames_z[fr, a])}"
                        f"{fmt(frames_vx[fr, a])}"
                        f"{fmt(frames_vy[fr, a])}"
                        f"{fmt(frames_vz[fr, a])}\n")
    return path


def load_display_txt(path):
    """Read display.txt new text format into numpy arrays（快速版）.
    
    Returns dict with keys: frames_x/y/z/vx/vy/vz, atom_ids,
           n_total_frames, n_total_particles, header_fields
    """
    import re
    
    with open(path, "r", encoding="utf-8") as f:
        raw = f.read()
    
    # 解析 header 行
    header_fields = {}
    n_total_frames = 0
    n_total_particles = 0
    
    lines = raw.splitlines()
    data_start = 0
    for i, line in enumerate(lines):
        line_stripped = line.strip()
        if line_stripped.startswith("number of frames:"):
            n_total_frames = int(line_stripped.split(":")[1].strip())
        elif line_stripped.startswith("number of particles:"):
            n_total_particles = int(line_stripped.split(":")[1].strip())
        elif line_stripped.startswith("frame:"):
            data_start = i
            break
        else:
            if ":" in line_stripped:
                k, v = line_stripped.split(":", 1)
                header_fields[k.strip()] = v.strip()
    
    # 快速定位所有数据行：跳过 frame header 和 column header
    # 数据行格式：每行 7 个字段（n x y z vx vy vz），固定宽度列
    data_text = []
    i = data_start
    n_frames = 0
    while i < len(lines):
        line = lines[i].strip()
        if line.startswith("frame:"):
            n_frames += 1
            i += 2  # 跳过 "frame: N" 和列头行
            continue
        if line:
            data_text.append(line)
        i += 1
    
    if n_frames == 0 or not data_text:
        raise ValueError(f"{path} 中没有有效帧数据")
    
    # 用 numpy 批量解析所有数据行（远比逐行 split+float 快）
    data_array = np.genfromtxt(data_text, dtype=np.float64)
    # data_array shape: (n_frames * n_atoms, 7) — 列: n, x, y, z, vx, vy, vz
    
    n_atoms = n_total_particles
    atoms_per_frame = len(data_text) // n_frames
    
    # 提取原子ID（第一帧即可）
    atom_ids = data_array[0:n_atoms, 0].astype(np.int64)
    
    # 重塑为 (n_frames, n_atoms, 6) — 去掉第0列(原子ID)
    all_data = data_array[:, 1:].reshape(n_frames, n_atoms, 6)
    
    return {
        "frames_x":  all_data[:, :, 0],
        "frames_y":  all_data[:, :, 1],
        "frames_z":  all_data[:, :, 2],
        "frames_vx": all_data[:, :, 3],
        "frames_vy": all_data[:, :, 4],
        "frames_vz": all_data[:, :, 5],
        "atom_ids": atom_ids,
        "n_total_frames": n_total_frames,
        "n_total_particles": n_total_particles,
        "header_fields": header_fields,
    }


def save_display_npz(path, frames_x, frames_y, frames_z,
                     frames_vx, frames_vy, frames_vz,
                     atom_ids, header_fields=None):
    """Write display.npz binary format alongside display.txt.

    Stores all frame arrays in compressed NumPy format.  Metadata is
    serialised as a JSON byte-string stored in the 'meta' array so the
    file stays self-contained.

    Args:
        path: output path ending in '.npz'
        frames_x/y/z/vx/vy/vz: (n_frames, n_atoms) float64 arrays
        atom_ids: (n_atoms,) int array
        header_fields: dict of metadata key-value strings (same as for
                       save_display_txt)
    """
    meta_bytes = np.bytes_(json.dumps(header_fields or {}, ensure_ascii=False))
    stem = path[:-4] if path.endswith(".npz") else path
    np.savez_compressed(
        stem,
        frames_x=frames_x,
        frames_y=frames_y,
        frames_z=frames_z,
        frames_vx=frames_vx,
        frames_vy=frames_vy,
        frames_vz=frames_vz,
        atom_ids=atom_ids,
        meta=np.array(meta_bytes),
    )
    return path


def load_display_npz(path):
    """Read display.npz and return the same dict as load_display_txt."""
    data = np.load(path, allow_pickle=False)
    header_fields = json.loads(data["meta"].item().decode("utf-8"))
    frames_x = data["frames_x"]
    n_total_frames = int(header_fields.get("number_of_frames",
                                           frames_x.shape[0]))
    n_total_particles = frames_x.shape[1]
    return {
        "frames_x":  frames_x,
        "frames_y":  data["frames_y"],
        "frames_z":  data["frames_z"],
        "frames_vx": data["frames_vx"],
        "frames_vy": data["frames_vy"],
        "frames_vz": data["frames_vz"],
        "atom_ids":  data["atom_ids"],
        "n_total_frames": n_total_frames,
        "n_total_particles": n_total_particles,
        "header_fields": header_fields,
    }


def load_display(output_dir):
    """Load display data, preferring binary .npz over text .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 load_display_npz(npz_path)
    if os.path.exists(txt_path):
        return load_display_txt(txt_path)
    raise FileNotFoundError(
        f"找不到 display.npz 或 display.txt in {output_dir}")


def get_output_dir(base_dir=None):
    """Return the output directory used for generated artifacts."""
    override = os.environ.get("DYNAMICS_OUTPUT_DIR")
    if override:
        output_dir = os.path.abspath(override)
    else:
        if base_dir is None:
            base_dir = os.path.dirname(os.path.abspath(__file__))
        output_dir = os.path.join(base_dir, "output")
    os.makedirs(output_dir, exist_ok=True)
    return output_dir


def get_input_dir(base_dir=None):
    """Return the input directory used for configuration and source data."""
    override = os.environ.get("DYNAMICS_INPUT_DIR")
    if override:
        input_dir = os.path.abspath(override)
    else:
        if base_dir is None:
            base_dir = os.path.dirname(os.path.abspath(__file__))
        input_dir = os.path.join(base_dir, "input")
    os.makedirs(input_dir, exist_ok=True)
    return input_dir


def load_coord_file(coord_path):
    """Load atom ids, masses, radii, coordinates, and velocities."""
    if not os.path.exists(coord_path):
        raise FileNotFoundError(f"坐标文件不存在: {coord_path}")

    rows = []
    with open(coord_path, "r", encoding="utf-8") as f:
        header = f.readline().strip().split()
        expected = ["n", "mass", "radius", "x", "y", "z", "vx", "vy", "vz"]
        legacy = expected + ["fixed_x", "fixed_y", "fixed_z"]
        legacy_new = expected + ["fix_x", "fix_y", "fix_z"]
        if header not in (expected, legacy, legacy_new):
            raise ValueError(
                f"坐标文件表头应为: {' '.join(expected)} 或加三列 fix_x fix_y fix_z，"
                f"实际为: {' '.join(header)}")

        for line_no, line in enumerate(f, start=2):
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            parts = stripped.split()
            if header == expected and len(parts) != 9:
                raise ValueError(f"{coord_path}:{line_no} 应有 9 列，实际为 {len(parts)} 列")
            if header in (legacy, legacy_new) and len(parts) != 12:
                raise ValueError(f"{coord_path}:{line_no} 应有 12 列，实际为 {len(parts)} 列")
            rows.append(parts)

    if not rows:
        raise ValueError(f"坐标文件没有原子数据: {coord_path}")

    atom_ids = np.array([int(row[0]) for row in rows], dtype=np.int64)
    masses = np.array([float(row[1]) for row in rows], dtype=np.float64)
    radii = np.array([float(row[2]) for row in rows], dtype=np.float64)
    positions = np.array([[float(row[3]), float(row[4]), float(row[5])] for row in rows],
                         dtype=np.float64)
    velocities = np.array([[float(row[6]), float(row[7]), float(row[8])] for row in rows],
                          dtype=np.float64)
    fixed = np.zeros((len(rows), 3), dtype=np.int64)
    if header in (legacy, legacy_new):
        fixed = np.array([[int(row[9]), int(row[10]), int(row[11])] for row in rows],
                         dtype=np.int64)

    if np.any(masses <= 0):
        raise ValueError("坐标文件中的质量必须为正数")
    if np.any(radii <= 0):
        raise ValueError("坐标文件中的半径必须为正数")

    print(f"[compute] 已加载坐标文件: {coord_path}，原子数={len(atom_ids)}")
    return atom_ids, masses, radii, positions, velocities, fixed


def load_bond_parameters(bond_path):
    """Load bond stiffness and optional rest length by bond name.

    File format (表头):
        bond_name k              # 2 列：从初始坐标计算键长
        bond_name k rest_length  # 3 列：显式指定键长
    """
    if not os.path.exists(bond_path):
        raise FileNotFoundError(f"键参数文件不存在: {bond_path}")

    bond_map = {}
    with open(bond_path, "r", encoding="utf-8") as f:
        header = f.readline().strip().split()

        if header == ["bond_name", "k"]:
            has_rest_length = False
        elif header == ["bond_name", "k", "rest_length"]:
            has_rest_length = True
        else:
            expected = ["bond_name", "k"] + (["rest_length"] if len(header) == 3 else [])
            raise ValueError(
                f"键参数文件表头应为: {' '.join(expected)}，实际为: {' '.join(header)}")

        ncols = 3 if has_rest_length else 2

        for line_no, line in enumerate(f, start=2):
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            parts = stripped.split()
            if len(parts) != ncols:
                raise ValueError(
                    f"{bond_path}:{line_no} 应有 {ncols} 列，实际为 {len(parts)} 列")
            bond_name = parts[0]
            stiffness = float(parts[1])
            if stiffness < 0:
                raise ValueError(f"{bond_path}:{line_no} 劲度系数必须非负")

            rest_length = None
            if has_rest_length:
                rest_length = float(parts[2])
                if rest_length <= 0:
                    raise ValueError(f"{bond_path}:{line_no} 键长必须为正数")

            bond_map[bond_name] = {"stiffness": stiffness, "rest_length": rest_length}

    if not bond_map:
        raise ValueError(f"键参数文件没有有效数据: {bond_path}")
    return bond_map


def load_bond_connections(connection_path, atom_ids, atom_positions, bond_map):
    """Load bonded atom pairs and derive rest lengths.

    键长优先级：
      1. bond_map 中显式指定的 rest_length（bond.txt 第三列）
      2. 否则从初始坐标计算
    """
    if not os.path.exists(connection_path):
        raise FileNotFoundError(f"成键连接文件不存在: {connection_path}")

    atom_index = {int(atom_id): idx for idx, atom_id in enumerate(atom_ids)}
    pairs = []
    bond_names = []
    stiffness = []
    rest_lengths = []

    with open(connection_path, "r", encoding="utf-8") as f:
        header = f.readline().strip().split()
        expected = ["n1", "n2", "bond_name"]
        if header != expected:
            raise ValueError(
                f"成键连接文件表头应为: {' '.join(expected)}，实际为: {' '.join(header)}")

        for line_no, line in enumerate(f, start=2):
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            parts = stripped.split()
            if len(parts) != 3:
                raise ValueError(f"{connection_path}:{line_no} 应有 3 列，实际为 {len(parts)} 列")

            atom_1 = int(parts[0])
            atom_2 = int(parts[1])
            bond_name = parts[2]

            if atom_1 not in atom_index or atom_2 not in atom_index:
                raise ValueError(
                    f"{connection_path}:{line_no} 中的原子序号 {atom_1}, {atom_2} 不在坐标文件里")
            if bond_name not in bond_map:
                raise ValueError(
                    f"{connection_path}:{line_no} 使用了未定义的键类型 {bond_name}")

            idx_1 = atom_index[atom_1]
            idx_2 = atom_index[atom_2]

            bond_entry = bond_map[bond_name]
            rest_length = bond_entry["rest_length"]
            if rest_length is None:
                # 未指定键长时从初始坐标计算
                delta = atom_positions[idx_2] - atom_positions[idx_1]
                rest_length = float(np.linalg.norm(delta))

            pairs.append((idx_1, idx_2))
            bond_names.append(bond_name)
            stiffness.append(bond_entry["stiffness"])
            rest_lengths.append(rest_length)

    if not pairs:
        return (
            np.zeros((0, 2), dtype=np.int64),
            [],
            np.zeros(0, dtype=np.float64),
            np.zeros(0, dtype=np.float64),
        )

    return (
        np.array(pairs, dtype=np.int64),
        bond_names,
        np.array(stiffness, dtype=np.float64),
        np.array(rest_lengths, dtype=np.float64),
    )


# ===========================================================================
# 驱动力加载 & 应用
# ===========================================================================

def load_driver_file(driver_path, atom_ids):
    """从 driver.txt 加载驱动力定义。

    格式：
        n amp_x amp_y amp_z freq_x freq_y freq_z phi_x phi_y phi_z period
       数值    0     0     0      0      0     10     0     0    90     all

    其中：
      position = A * cos(2π f t + φ),  φ 为角度制（自动转弧度）
      period = all | 数值（周期数，结束后原子静止）
    """
    if not os.path.exists(driver_path):
        print(f"[compute] 警告: 驱动力文件不存在: {driver_path}")
        return None

    atom_index = {int(aid): idx for idx, aid in enumerate(atom_ids)}
    drivers = []
    ncols = 11

    with open(driver_path, "r", encoding="utf-8") as f:
        header = f.readline().strip().split()
        expected = ["n", "amp_x", "amp_y", "amp_z",
                     "freq_x", "freq_y", "freq_z",
                     "phi_x", "phi_y", "phi_z", "period"]
        if header != expected:
            raise ValueError(
                f"driver.txt 表头应为: {' '.join(expected)}，实际为: {' '.join(header)}")

        for line_no, line in enumerate(f, start=2):
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                continue
            parts = stripped.split()
            if len(parts) != ncols:
                raise ValueError(
                    f"{driver_path}:{line_no} 应有 {ncols} 列，实际为 {len(parts)} 列")

            n = int(parts[0])
            if n not in atom_index:
                raise ValueError(f"{driver_path}:{line_no} 原子序号 {n} 不在 coord.txt 中")

            amp = np.array([float(parts[1]), float(parts[2]), float(parts[3])],
                           dtype=np.float64)
            freq = np.array([float(parts[4]), float(parts[5]), float(parts[6])],
                            dtype=np.float64)
            phi_deg = np.array([float(parts[7]), float(parts[8]), float(parts[9])],
                               dtype=np.float64)
            phi_rad = phi_deg * np.pi / 180.0

            period_str = parts[10]

            drivers.append({
                "atom_index": atom_index[n],
                "atom_id": n,
                "amp": amp,
                "freq": freq,
                "phi": phi_rad,
                "period_str": period_str,
                "period_cycles": None if period_str == "all" else float(period_str),
                # 平衡位置（原子初始坐标），驱动为相对于此位置的偏移
                "eq_pos": None,
                # 在模拟中动态记录：冻结步数索引、冻结位置
                "freeze_step": None,
                "freeze_pos": None,
            })

    if not drivers:
        print(f"[compute] 警告: driver.txt 中没有有效数据")
        return None

    print(f"[compute] 已加载驱动力: {len(drivers)} 条定义")
    if len(drivers) <= 20:
        for d in drivers:
            print(f"         原子 {d['atom_id']}: "
                  f"A=({d['amp'][0]},{d['amp'][1]},{d['amp'][2]}), "
                  f"f=({d['freq'][0]},{d['freq'][1]},{d['freq'][2]}), "
                  f"φ=({phi_deg[d['amp'].tolist().index(max(d['amp']))]}° 等), "
                  f"period={d['period_str']}")
    else:
        _first = drivers[0]
        _last  = drivers[-1]
        print(f"         原子 {_first['atom_id']}~{_last['atom_id']}: "
              f"A=({_first['amp'][0]},{_first['amp'][1]},{_first['amp'][2]}), "
              f"f=({_first['freq'][0]},{_first['freq'][1]},{_first['freq'][2]}), "
              f"period={_first['period_str']} 等 {len(drivers)} 条")
    return drivers


def apply_driving_force(x, y, z, vx, vy, vz, t, step, drivers, dt):
    """对受驱原子按驱动力函数覆盖位置/速度。

    驱动力公式：pos = A * cos(2π f t + φ)
                vel = -A * 2π f * sin(2π f t + φ)
    """
    if drivers is None:
        return x, y, z, vx, vy, vz

    for d in drivers:
        idx = d["atom_index"]

        # 确定驱动力持续到哪一步
        if d["period_cycles"] is not None:
            max_freq = np.max(np.abs(d["freq"]))
            if max_freq > 1e-12:
                period_duration = d["period_cycles"] / max_freq
                period_steps = int(period_duration / dt)
            else:
                period_steps = 0

            if step > period_steps:
                # 驱动力已结束：原子静止（保持最后位置，速度归零）
                if d["freeze_pos"] is not None:
                    x[idx], y[idx], z[idx] = d["freeze_pos"]
                vx[idx] = vy[idx] = vz[idx] = 0.0
                continue
        else:
            period_steps = None  # 全程驱动

        # 当前驱动力下的位置 / 速度（偏移叠加在初始平衡坐标上）
        eq = d["eq_pos"] if d["eq_pos"] is not None else np.zeros(3)
        pos_drive = eq + d["amp"] * np.cos(2.0 * np.pi * d["freq"] * t + d["phi"])
        vel_drive = -d["amp"] * 2.0 * np.pi * d["freq"] * np.sin(2.0 * np.pi * d["freq"] * t + d["phi"])

        x[idx] = pos_drive[0]
        y[idx] = pos_drive[1]
        z[idx] = pos_drive[2]
        vx[idx] = vel_drive[0]
        vy[idx] = vel_drive[1]
        vz[idx] = vel_drive[2]

        # 若周期有限，记录冻结位置（驱动力最后一帧的位置）
        if period_steps is not None and step == period_steps:
            d["freeze_pos"] = (float(pos_drive[0]), float(pos_drive[1]), float(pos_drive[2]))

    return x, y, z, vx, vy, vz
def parse_gravity_vector(value):
    """Parse gravity into a 3D acceleration vector.

    Backward compatibility:
    - scalar `G: 9.8` -> [0, 0, -9.8]
    - vector `G: [gx, gy, gz]` -> [gx, gy, gz]
    """
    if isinstance(value, (int, float, np.integer, np.floating)):
        return np.array([0.0, 0.0, -float(value)], dtype=np.float64)
    vector = np.asarray(value, dtype=np.float64)
    if vector.shape != (3,):
        raise ValueError(f"G 必须是长度为 3 的分量数组，实际为 {value}")
    return vector


def parse_damping_vector(value):
    """Parse damping into a 3D component vector.

    Backward compatibility:
    - scalar `B: 0.5` -> [0.5, 0.5, 0.5]
    - vector `B: [bx, by, bz]` -> [bx, by, bz]
    """
    if isinstance(value, (int, float, np.integer, np.floating)):
        scalar = float(value)
        return np.array([scalar, scalar, scalar], dtype=np.float64)
    vector = np.asarray(value, dtype=np.float64)
    if vector.shape != (3,):
        raise ValueError(f"B 必须是长度为 3 的分量数组，实际为 {value}")
    return vector


def run_from_config(config, out_dir=None):
    """直接接收配置字典，运行模拟并保存结果。

    config 字典结构：
        box_a, alpha, coord_file, connection_file, bond_file, plot_atom, G, B, method, NT, DT, NSTEP,
        warmup_steps, sample_start, sample_end（可选）,
        ball_radius, ball_color_r/g/b, box_color_r/g/b（可选）

    out_dir: 输出目录，默认同脚本所在目录
    返回：(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz)
    """
    global box_a, alpha, COORD_FILE, X0, Y0, Z0, VX0, VY0, VZ0, M, G, B, METHOD, NT, DT, NSTEP
    global PLOT_ATOM_ID, PLOT_ATOM_ROW
    global ATOM_IDS, ATOM_MASSES, ATOM_RADII, ATOM_POSITIONS, ATOM_VELOCITIES, ATOM_FIXED
    global BOND_CONNECTION_FILE, BOND_PARAMETER_FILE, BOND_PAIRS, BOND_NAMES, BOND_STIFFNESS, BOND_REST_LENGTHS
    global X_MIN, X_MAX, Y_MIN, Y_MAX, Z_MIN, Z_MAX
    global ball_radius, ball_color_r, ball_color_g, ball_color_b
    global box_color_r, box_color_g, box_color_b
    global use_marker, camera_keyframes_raw, camera_distance, camera_elevation, camera_azimuth
    global camera_center_x, camera_center_y, camera_center_z
    global camera_pos_x, camera_pos_y, camera_pos_z, display_amp_str
    global FIXED_MASK_X, FIXED_MASK_Y, FIXED_MASK_Z
    global warmup_steps, sample_start, sample_end
    global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH
    global DRIVING_FORCE, DRIVER_DATA

    box_a = float(config.get("box_a", 10.0))
    raw_alpha = config.get("alpha", 0.2)
    if isinstance(raw_alpha, (list, tuple)):
        alpha = [float(a) for a in raw_alpha]
    else:
        alpha = float(raw_alpha)
    COORD_FILE = str(config.get("coord_file", os.path.join("input", "coord.txt")))
    coord_path = COORD_FILE
    if out_dir is not None and not os.path.isabs(coord_path):
        coord_path = os.path.join(out_dir, coord_path)
    (ATOM_IDS, ATOM_MASSES, ATOM_RADII, ATOM_POSITIONS,
     ATOM_VELOCITIES, ATOM_FIXED) = load_coord_file(coord_path)
    FIXED_MASK_X = ATOM_FIXED[:, 0] != 0
    FIXED_MASK_Y = ATOM_FIXED[:, 1] != 0
    FIXED_MASK_Z = ATOM_FIXED[:, 2] != 0
    BOND_CONNECTION_FILE = str(config.get("connection_file", os.path.join("input", "connection.txt")))
    BOND_PARAMETER_FILE = str(config.get("bond_file", os.path.join("input", "bond.txt")))
    connection_path = BOND_CONNECTION_FILE
    bond_path = BOND_PARAMETER_FILE
    if out_dir is not None and not os.path.isabs(connection_path):
        connection_path = os.path.join(out_dir, connection_path)
    if out_dir is not None and not os.path.isabs(bond_path):
        bond_path = os.path.join(out_dir, bond_path)
    bond_map = load_bond_parameters(bond_path)
    (BOND_PAIRS, BOND_NAMES, BOND_STIFFNESS,
     BOND_REST_LENGTHS) = load_bond_connections(
        connection_path, ATOM_IDS, ATOM_POSITIONS, bond_map)
    PLOT_ATOM_ID = int(config.get("plot_atom", int(ATOM_IDS[0])))
    matches = np.where(ATOM_IDS == PLOT_ATOM_ID)[0]
    if len(matches) == 0:
        raise ValueError(f"plot_atom={PLOT_ATOM_ID} 不在坐标文件原子序号中")
    PLOT_ATOM_ROW = int(matches[0])

    M = float(ATOM_MASSES[PLOT_ATOM_ROW])
    X0, Y0, Z0 = [float(v) for v in ATOM_POSITIONS[PLOT_ATOM_ROW]]
    VX0, VY0, VZ0 = [float(v) for v in ATOM_VELOCITIES[PLOT_ATOM_ROW]]
    G = parse_gravity_vector(config["G"])
    B = parse_damping_vector(config["B"])
    METHOD = normalize_method_name(config.get("method", "explicit_euler"))
    NT = int(config["NT"])
    DT = float(config["DT"])
    NSTEP = int(config["NSTEP"])

    # 步骤控制参数（可选，有默认值）
    warmup_steps = int(config.get("warmup_steps", 0))
    sample_start = config.get("sample_start")  # None 表示从头开始
    sample_end   = config.get("sample_end")    # None 表示到末尾

    X_MIN = -box_a; X_MAX = box_a
    Y_MIN = -box_a; Y_MAX = box_a
    Z_MIN = -box_a; Z_MAX = box_a

    ball_radius = float(config.get("ball_radius", ATOM_RADII[PLOT_ATOM_ROW]))
    ball_color_r = float(config.get("ball_color_r", 0.9))
    ball_color_g = float(config.get("ball_color_g", 0.2))
    ball_color_b = float(config.get("ball_color_b", 0.2))
    box_color_r = float(config.get("box_color_r", 0.8))
    box_color_g = float(config.get("box_color_g", 0.8))
    box_color_b = float(config.get("box_color_b", 0.85))
    use_marker = int(config.get("use_marker", 0))
    camera_distance = float(config.get("camera_distance", 40.0))
    camera_elevation = float(config.get("camera_elevation", 0.0))
    camera_azimuth = float(config.get("camera_azimuth", 0.0))
    camera_center_x = float(config.get("camera_center_x", 0.0))
    camera_center_y = float(config.get("camera_center_y", 0.0))
    camera_center_z = float(config.get("camera_center_z", 0.0))
    camera_pos_x = config.get("camera_pos_x", None)
    camera_pos_y = config.get("camera_pos_y", None)
    camera_pos_z = config.get("camera_pos_z", None)
    _damp_cfg = config.get("display_amp", None)
    if _damp_cfg is not None:
        import ast as _ast_cfg
        _dv = np.array(_ast_cfg.literal_eval(str(_damp_cfg).strip()), dtype=np.float64)
        display_amp_str = f"[{_dv[0]},{_dv[1]},{_dv[2]}]"
    else:
        display_amp_str = ""

    # 力开关
    global GRAVITY_FIELD, GRAVITY_INTERACTION, ELASTIC_FORCE, DAMPING_FORCE, GRAVITY_STRENGTH
    global DRIVING_FORCE, DRIVER_DATA
    GRAVITY_FIELD = int(config.get("gravity_field", 1))
    GRAVITY_INTERACTION = int(config.get("gravity_interaction", 0))
    ELASTIC_FORCE = int(config.get("elastic_force", 1))
    DAMPING_FORCE = int(config.get("damping_force", 0))
    GRAVITY_STRENGTH = float(config.get("gravity_strength", 1.0))
    DRIVING_FORCE = int(config.get("driving_force", 0))

    # 加载驱动力定义
    DRIVER_DATA = None
    if DRIVING_FORCE:
        driver_rel = str(config.get("driver_file", os.path.join("input", "driver.txt")))
        driver_path = driver_rel
        if out_dir is not None and not os.path.isabs(driver_rel):
            driver_path = os.path.join(out_dir, driver_rel)
        DRIVER_DATA = load_driver_file(driver_path, ATOM_IDS)
        # 将每个驱动原子的初始坐标作为平衡位置
        if DRIVER_DATA:
            for _d in DRIVER_DATA:
                _i = _d["atom_index"]
                _d["eq_pos"] = np.array([
                    ATOM_POSITIONS[_i, 0],
                    ATOM_POSITIONS[_i, 1],
                    ATOM_POSITIONS[_i, 2],
                ], dtype=np.float64)

    # 加载运动相机关键帧
    camera_keyframes_raw = ""
    move_camera = int(config.get("move_camera", 0))
    camera_keyframes_raw = ""
    if move_camera:
        cam_rel = str(config.get("move_camera_file", os.path.join("input", "move_camera.txt")))
        cam_path = cam_rel
        if out_dir is not None and not os.path.isabs(cam_rel):
            cam_path = os.path.join(out_dir, cam_rel)
        camera_keyframes_raw = _load_camera_motion(cam_path)

    print(f"[compute] 使用算法: {METHOD}")
    print(f"[compute] 已加载成键信息: {len(BOND_PAIRS)} 条键")
    if config.get("_skip_run", False):
        return None, None, None, None, None, None
    t_start = time.time()
    t_start_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation(
        save_trajectory=int(config.get("save_trajectory", 0)))
    t_end = time.time()
    t_end_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    elapsed = t_end - t_start

    # 写入 dynamics.log
    log_dir = get_output_dir(out_dir)
    log_path = os.path.join(log_dir, "dynamics.log")
    with open(log_path, "w", encoding="utf-8") as f:
        f.write("=" * 50 + "\n")
        f.write("Dynamics 计算日志\n")
        f.write("=" * 50 + "\n")
        f.write(f"计算开始时间: {t_start_str}\n")
        f.write(f"计算结束时间: {t_end_str}\n")
        f.write(f"计算耗时:     {elapsed:.3f} s\n")
        f.write("-" * 50 + "\n")
        f.write("计算参数:\n")
        f.write(f"  box_a:      {box_a}\n")
        _alpha_str = alpha if isinstance(alpha, str) else (str(alpha) if isinstance(alpha, list) else str(alpha))
        f.write(f"  alpha:      {_alpha_str}\n")
        f.write(f"  coord_file: {COORD_FILE}\n")
        f.write(f"  method:     {METHOD}\n")
        f.write(f"  NT:         {NT}\n")
        f.write(f"  DT:         {DT}\n")
        f.write(f"  NSTEP:      {NSTEP}\n")
        f.write(f"  G:          ({G[0]:.2f}, {G[1]:.2f}, {G[2]:.2f})\n")
        f.write(f"  B:          ({B[0]:.2f}, {B[1]:.2f}, {B[2]:.2f})\n")
        f.write(f"  warmup:     {warmup_steps}\n")
        f.write(f"  原子数:     {len(ATOM_IDS)}\n")
        f.write(f"  键数:       {len(BOND_PAIRS)}\n")
        f.write("-" * 50 + "\n")
        f.write("初始状态 (plot_atom):\n")
        f.write(f"  position:   ({X0:.6f}, {Y0:.6f}, {Z0:.6f})\n")
        f.write(f"  velocity:   ({VX0:.6f}, {VY0:.6f}, {VZ0:.6f})\n")
        f.write(f"  mass:       {M}\n")
        f.write("=" * 50 + "\n")
    print(f"[compute] 日志已保存至: {log_path}")
    record = len(traj_x)
    print(f"[compute] 计算完成: {record} 步, {elapsed:.3f} s")

    return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz


def run_engine_dll(engine, output_dir, config):
    """通过 DLL（ctypes）调用计算引擎，不经过文件 I/O，直接返回轨迹数组。

    Args:
        engine:     引擎名称 "c", "cpp", 或 "fortran"
        output_dir: 输出目录（用于保存 display.npz）
        config:     YAML 配置字典
    Returns:
        None（结果直接写入 output_dir/display.npz）
    Raises:
        FileNotFoundError: DLL 尚未编译
        RuntimeError: DLL 运算出错
    """
    import sys as _sys
    import datetime as _datetime
    _eng_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "engines")
    if _eng_dir not in _sys.path:
        _sys.path.insert(0, _eng_dir)
    from engine_dll import load_dll, run_dynamics_dll, is_dll_available

    if not is_dll_available(engine):
        raise FileNotFoundError(
            f"DLL 未找到（引擎 {engine}）。"
            f"请先编译：cd engines/{engine} && make dll")

    lib = load_dll(engine)

    # ── 构造原子数据 ──────────────────────────────────────────
    if ATOM_POSITIONS is None:
        raise RuntimeError("run_engine_dll: 请先调用 load_parameters() 加载配置")

    pos  = np.asarray(ATOM_POSITIONS, dtype=np.float64)   # (n, 3)
    vel  = np.asarray(ATOM_VELOCITIES, dtype=np.float64)
    mass = np.asarray(ATOM_MASSES, dtype=np.float64)
    fixed= np.asarray(ATOM_FIXED, dtype=np.int32)         # (n, 3)

    # ── 键数据 ────────────────────────────────────────────────
    n_bonds = len(BOND_PAIRS) if BOND_PAIRS is not None else 0
    bp  = np.asarray(BOND_PAIRS,       dtype=np.int32)   if n_bonds else np.zeros((0,2), dtype=np.int32)
    bk  = np.asarray(BOND_STIFFNESS,   dtype=np.float64) if n_bonds else np.zeros(0)
    br0 = np.asarray(BOND_REST_LENGTHS,dtype=np.float64) if n_bonds else np.zeros(0)

    # ── 驱动数据 ──────────────────────────────────────────────
    drv_list = []
    if int(config.get("driving_force", 0)) and DRIVER_DATA:
        atom_id_to_local = {int(aid): i for i, aid in enumerate(ATOM_IDS)}
        for d in DRIVER_DATA:
            aid = int(d.get("atom_id", -1))
            if aid not in atom_id_to_local:
                continue
            local_idx = atom_id_to_local[aid]
            # d["amp"], d["freq"], d["phi"] are numpy arrays; d["phi"] is already in radians
            amp  = [float(v) for v in d["amp"]]
            freq = [float(v) for v in d["freq"]]
            phi  = [float(v) for v in d["phi"]]   # radians
            # eq_pos is set by run_from_config; fall back to initial position
            eq_pos = d.get("eq_pos")
            eq = ([float(v) for v in eq_pos] if eq_pos is not None
                  else [float(pos[local_idx, 0]), float(pos[local_idx, 1]), float(pos[local_idx, 2])])
            pc = d.get("period_cycles")            # None → unlimited, float → finite
            nc = float(pc) if pc is not None else 0.0
            hp = 1 if nc > 0 else 0
            drv_list.append({"local_idx": local_idx, "amp": amp, "freq": freq,
                              "phi": phi, "eq_pos": eq, "n_cycles": nc, "has_period": hp})

    # ── 进度回调 ──────────────────────────────────────────────
    total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0))
    try:
        from tqdm import tqdm as _tqdm
        _pbar = _tqdm(total=total_steps, desc=f"[compute] DLL {engine}",
                      unit="步", bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]')
        def _cb(step, total):
            _pbar.n = step
            _pbar.refresh()
    except ImportError:
        _pbar = None
        _cb = None

    _t0 = time.time()
    try:
        result = run_dynamics_dll(
            lib, config,
            pos, vel, mass, fixed,
            bp, bk, br0,
            drv_list, np.asarray(ATOM_IDS),
            progress_cb=_cb,
        )
    finally:
        if _pbar is not None:
            _pbar.n = total_steps
            _pbar.close()
    elapsed = time.time() - _t0

    n_frames, n_atoms = result["x"].shape
    print(f"[compute] DLL 完成: {n_frames} 帧  {n_atoms} 原子  {elapsed:.3f} s")

    # ── 构建 header 并保存 display.npz ────────────────────────
    # 与 run_simulation 写入的 header 保持字段完全一致，
    # 确保 draw.py / plot_wave.py 读到所有必要参数。
    G_vec = parse_gravity_vector(config.get("G", [0, 0, 0]))
    B_vec = parse_damping_vector(config.get("B", [0, 0, 0]))
    record_steps_hdr = int(config["NT"]) - int(config.get("warmup_steps", 0))
    header = {
        "DT":                  str(config["DT"]),
        "NSTEP":               str(config.get("NSTEP", 1)),
        "method":              str(config.get("method", "leapfrog")),
        "NT":                  str(config["NT"]),
        "warmup_steps":        str(config.get("warmup_steps", 0)),
        "dynamic_steps":       str(record_steps_hdr),
        "T_total":             str(int(config["NT"]) * float(config["DT"])),
        "box_a":               str(config.get("box_a", 300.0)),
        "gravity_field":       str(config.get("gravity_field", 0)),
        "gravity_interaction": str(config.get("gravity_interaction", 0)),
        "elastic_force":       str(config.get("elastic_force", 1)),
        "damping_force":       str(config.get("damping_force", 0)),
        "driving_force":       str(config.get("driving_force", 0)),
        "gravity_strength":    str(config.get("gravity_strength", 1.0)),
        "G":                   json.dumps(G_vec.tolist()),
        "B":                   json.dumps(B_vec.tolist()),
        "number_of_frames":    str(n_frames),
        "number_of_particles": str(n_atoms),
        # draw.py 需要的渲染参数
        "use_marker":          str(use_marker),
        "display_color":       json.dumps(config.get("display_color",  
            {"x":[0,[1.0,0.0,0.0]],"y":[0,[0.0,1.0,0.0]],"z":[0,[0.0,0.0,1.0]],
             "xy":[0,[1.0,1.0,0.0]],"yz":[0,[0.0,1.0,1.0]],"zx":[0,[1.0,0.0,1.0]],
             "xyz":[1,[1.0,1.0,1.0]]})),
        "ball_radius":         str(config.get("ball_radius", float(ATOM_RADII[0]) if ATOM_RADII is not None else 0.5)),
        "ball_color_r":        str(config.get("ball_color_r", 0.9)),
        "ball_color_g":        str(config.get("ball_color_g", 0.2)),
        "ball_color_b":        str(config.get("ball_color_b", 0.2)),
        "box_color_r":         str(config.get("box_color_r", 0.8)),
        "box_color_g":         str(config.get("box_color_g", 0.8)),
        "box_color_b":         str(config.get("box_color_b", 0.85)),
        "alpha":               ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])),
        # draw.py / plot_wave.py 需要的原子、键数据
        "atom_radii":          ",".join(str(r) for r in ATOM_RADII),
        "atom_masses":         json.dumps([float(v) for v in ATOM_MASSES]),
        "atom_positions":      json.dumps(ATOM_POSITIONS.tolist()),
        "bond_pairs":          json.dumps(BOND_PAIRS.tolist() if BOND_PAIRS is not None else []),
        "bond_stiffness":      json.dumps(BOND_STIFFNESS.tolist() if BOND_STIFFNESS is not None else []),
        "bond_rest_lengths":   json.dumps(BOND_REST_LENGTHS.tolist() if BOND_REST_LENGTHS is not None else []),
        # 边界（draw.py 用于场景缩放）
        "X_MIN": str(-float(config.get("box_a", 300.0))),
        "X_MAX": str( float(config.get("box_a", 300.0))),
        "Y_MIN": str(-float(config.get("box_a", 300.0))),
        "Y_MAX": str( float(config.get("box_a", 300.0))),
        "Z_MIN": str(-float(config.get("box_a", 300.0))),
        "Z_MAX": str( float(config.get("box_a", 300.0))),
        # 相机参数
        "camera_distance":     str(camera_distance),
        "camera_elevation":    str(camera_elevation),
        "camera_azimuth":      str(camera_azimuth),
        "camera_center_x":     str(camera_center_x),
        "camera_center_y":     str(camera_center_y),
        "camera_center_z":     str(camera_center_z),
        "camera_keyframes":    str(camera_keyframes_raw),
    }
    if display_amp_str:
        header["display_amp"] = display_amp_str
    if camera_pos_x is not None:
        header["camera_pos_x"] = str(camera_pos_x)
        header["camera_pos_y"] = str(camera_pos_y)
        header["camera_pos_z"] = str(camera_pos_z)

    os.makedirs(output_dir, exist_ok=True)
    npz_path = os.path.join(output_dir, "display.npz")
    save_display_npz(
        npz_path,
        result["x"], result["y"], result["z"],
        result["vx"], result["vy"], result["vz"],
        np.asarray(ATOM_IDS),
        header_fields=header,
    )
    print(f"[compute] display.npz 已生成: {npz_path}")


def run_engine(engine, input_dir, output_dir, config):
    """调用外部计算引擎（C/C++/Fortran），生成 trajectory.txt。

    Args:
        engine:    引擎名称 ("c", "cpp", "fortran")
        input_dir: 输入目录（含 coord.txt, connection.txt, bond.txt）
        output_dir: 输出目录（轨迹文件写入位置）
        config:    YAML 配置字典
    Returns:
        (traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz)
    """
    script_dir = os.path.dirname(os.path.abspath(__file__))
    system = platform.system().lower()
    engine_map = {
        "c":       "engines/release/dynamics_c",
        "cpp":     "engines/release/dynamics_cpp",
        "c++":     "engines/release/dynamics_cpp",
        "fortran": "engines/release/dynamics_f90",
        "f90":     "engines/release/dynamics_f90",
        "python":  None,  # 特殊处理：用 sys.executable 调用 main.py
    }
    if engine not in engine_map:
        raise ValueError(f"不支持的引擎: {engine}，可选: c, cpp, fortran, python")

    if engine == "python":
        # Python 引擎：用当前解释器运行 engines/python/main.py
        py_main = os.path.join(script_dir, "engines", "python", "main.py")
        if not os.path.exists(py_main):
            raise FileNotFoundError(f"Python 引擎脚本不存在: {py_main}")
        found = py_main
        engine_path = sys.executable
    else:
        engine_rel = engine_map[engine]
        engine_path = os.path.join(script_dir, engine_rel)

        # 自动检测可执行文件后缀和平台专用版本
        candidates = [
            engine_path,
            engine_path + ".exe",
            engine_path + f"_{system}.exe",
        ]
        found = None
        for p in candidates:
            if os.path.exists(p):
                found = p
                break
        if found is None:
            raise FileNotFoundError(
                f"引擎可执行文件不存在: 尝试了 {candidates}\n"
                f"请先编译: cd engines/{engine} && make\n"
                f"或安装交叉编译器后: cd engines/{engine} && make {system}")

    # 构造 param.json（数值参数）
    G = parse_gravity_vector(config.get("G", [0, 0, -9.8]))
    B = parse_damping_vector(config.get("B", [0, 0, 0]))
    param_json = {
        "box_a": float(config.get("box_a", 10.0)),
        "NT": int(config["NT"]),
        "DT": float(config["DT"]),
        "NSTEP": int(config.get("NSTEP", 1)),
        "warmup_steps": int(config.get("warmup_steps", 0)),
        "method": normalize_method_name(config.get("method", "leapfrog")),
        "G": [float(v) for v in G],
        "B": [float(v) for v in B],
        "gravity_field": int(config.get("gravity_field", 1)),
        "gravity_interaction": int(config.get("gravity_interaction", 0)),
        "elastic_force": int(config.get("elastic_force", 1)),
        "damping_force": int(config.get("damping_force", 0)),
        "gravity_strength": float(config.get("gravity_strength", 1.0)),
        "driving_force": int(config.get("driving_force", 0)),
        "save_trajectory": int(config.get("save_trajectory", 0)),
        # 渲染参数（用于 display.txt header）
        "alpha": config.get("alpha", 0.2),
        "ball_radius": float(config.get("ball_radius", 0.5)),
        "ball_color": [float(config.get("ball_color_r", 0.9)),
                       float(config.get("ball_color_g", 0.2)),
                       float(config.get("ball_color_b", 0.2))],
        "box_color": [float(config.get("box_color_r", 0.8)),
                      float(config.get("box_color_g", 0.8)),
                      float(config.get("box_color_b", 0.85))],
        "use_marker": int(config.get("use_marker", 0)),
        "camera_distance": float(config.get("camera_distance", 40.0)),
        "camera_elevation": float(config.get("camera_elevation", 0)),
        "camera_azimuth": float(config.get("camera_azimuth", 0)),
        "camera_center_x": float(config.get("camera_center_x", 0.0)),
        "camera_center_y": float(config.get("camera_center_y", 0.0)),
        "camera_center_z": float(config.get("camera_center_z", 0.0)),
    }
    param_path = os.path.join(script_dir, "engines", "release", f"{engine}.json")
    os.makedirs(os.path.dirname(param_path), exist_ok=True)
    with open(param_path, "w", encoding="utf-8") as f:
        json.dump(param_json, f, indent=2)

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    print(f"[compute] 引擎: {engine} → {os.path.basename(engine_path)}")
    print(f"[compute]   input:  {input_dir}")
    print(f"[compute]   output: {output_dir}")

    # ── 预校准：跑少量步测速（结果缓存，避免每次重跑）────────────
    total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0))
    _calib_nt = min(1000, max(100, total_steps // 10))
    n_atoms_calib = len(ATOM_IDS) if ATOM_IDS is not None else 0

    # 尝试读取缓存；当 n_atoms 相同且 NT 在 50% 范围内时视为有效
    _cache_path = os.path.join(script_dir, "engines", "release", f"_calib_{engine}.json")
    _step_time = None
    try:
        with open(_cache_path, encoding="utf-8") as _cf:
            _cache = json.load(_cf)
        if (_cache.get("n_atoms") == n_atoms_calib and
                abs(_cache.get("nt", 0) - total_steps) / max(total_steps, 1) < 0.5):
            _step_time = float(_cache["step_time"])
            print(f"[compute] 使用校准缓存: {_step_time*1e6:.2f} μs/步")
    except Exception:
        pass

    if _step_time is None:
        _calib_param = dict(param_json)
        _calib_param["NT"] = _calib_nt
        _calib_path = os.path.join(script_dir, "engines", "release", f"_calib_{engine}.json")
        with open(_calib_path, "w", encoding="utf-8") as _cf:
            json.dump(_calib_param, _cf, indent=2)
        _calib_outdir = os.path.join(script_dir, "engines", "release", f"_calib_{engine}_out")
        os.makedirs(_calib_outdir, exist_ok=True)
        _ct0 = time.time()
        subprocess.run(
            [engine_path, os.path.abspath(input_dir), _calib_outdir, _calib_path],
            capture_output=True, timeout=60)
        # 清理校准临时文件
        for _f in os.listdir(_calib_outdir):
            try: os.remove(os.path.join(_calib_outdir, _f))
            except OSError: pass
        try: os.rmdir(_calib_outdir)
        except OSError: pass
        os.remove(_calib_path)
        _calib_elapsed = max(time.time() - _ct0, 0.001)
        _overhead = _calib_elapsed * 0.15
        _step_time = max(_calib_elapsed - _overhead, 0.0001) / _calib_nt
        # 保存缓存
        try:
            with open(_cache_path, "w", encoding="utf-8") as _cf:
                json.dump({"n_atoms": n_atoms_calib, "nt": total_steps,
                           "step_time": _step_time}, _cf)
        except OSError:
            pass

    _est_total = _step_time * total_steps

    t_start = time.time()
    t_start_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # Python 引擎：[python, main.py, args]；其他引擎：[exe, args]
    _cmd = ([engine_path, found] if engine == "python" else [engine_path]) + \
           [os.path.abspath(input_dir), os.path.abspath(output_dir), param_path]
    _p = subprocess.Popen(
        _cmd,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, encoding='utf-8', errors='replace')
    _engine_lines = []

    try:
        from tqdm import tqdm as _tqdm
        _pbar = _tqdm(total=total_steps, desc=f"[compute] 引擎 {engine}",
                      unit="步", bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]')
    except ImportError:
        _pbar = None

    try:
        while True:
            _line = _p.stdout.readline() if _p.stdout else ''
            if _line:
                _line = _line.strip()
                if _line:
                    _engine_lines.append(_line)
                    # 读取外部引擎真实进度：格式 "[xxx-engine] progress: N/total"
                    _prog_match = re.search(r'progress:\s*(\d+)/(\d+)', _line)
                    if _pbar is not None and _prog_match:
                        _prog_done = int(_prog_match.group(1))
                        _prog_total = int(_prog_match.group(2))
                        if _prog_total > 0:
                            _pbar.n = min(_prog_done, total_steps)
                            _pbar.refresh()
            if _p.poll() is not None:
                if _p.stdout:
                    for _r in _p.stdout:
                        _r = _r.strip()
                        if _r:
                            _engine_lines.append(_r)
                            # 读取残留在管道中的进度消息，避免 20%→100% 跳变
                            _prog_match = re.search(r'progress:\s*(\d+)/(\d+)', _r)
                            if _pbar is not None and _prog_match:
                                _p_done = min(int(_prog_match.group(1)), total_steps)
                                _pbar.n = max(_pbar.n, _p_done)
                    if _pbar is not None: _pbar.refresh()
                break
            time.sleep(0.05)
    finally:
        if _pbar is not None:
            _pbar.n = total_steps
            _pbar.refresh()
            _pbar.close()
        if _p.poll() is None:
            _p.kill()
            _p.wait(timeout=5)

    _rc = _p.returncode
    t_end = time.time()
    t_end_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    elapsed = t_end - t_start

    if _rc != 0:
        _err = _p.stderr.read() if _p.stderr else ''
        if _err:
            print(f"[compute] 引擎错误:\n{_err}")
        raise RuntimeError(f"引擎 {engine} 返回错误码 {_rc}")

    if _engine_lines:
        _log_path = os.path.join(output_dir, "dynamics.log")
        try:
            with open(_log_path, "a", encoding="utf-8") as _lf:
                _lf.write("\n--- 引擎输出 ---\n")
                for _ln in _engine_lines:
                    _lf.write(_ln + "\n")
        except OSError:
            pass

    # 加载输出的 trajectory.txt / display.txt
    save_traj = int(config.get("save_trajectory", 0))
    if not save_traj:
        # save_trajectory=0：引擎只写 display.txt
        disp_path = os.path.join(os.path.abspath(output_dir), "display.txt")
        if not os.path.exists(disp_path):
            raise FileNotFoundError(f"引擎未生成 display.txt: {disp_path}")
        print(f"[compute] 引擎已生成 {disp_path}")
        # 将 display.txt 转换为 display.npz（加快后续 draw.py / plot_wave.py 加载）
        _npz_path = disp_path.replace("display.txt", "display.npz")
        try:
            _d = load_display_txt(disp_path)
            # 外部引擎不写摄像机参数到 display.txt，在此从全局变量补入
            _cam_extra = {
                "camera_distance":  str(camera_distance),
                "camera_elevation": str(camera_elevation),
                "camera_azimuth":   str(camera_azimuth),
                "camera_center_x":  str(camera_center_x),
                "camera_center_y":  str(camera_center_y),
                "camera_center_z":  str(camera_center_z),
                "camera_keyframes": str(camera_keyframes_raw),
            }
            if camera_pos_x is not None:
                _cam_extra["camera_pos_x"] = str(camera_pos_x)
                _cam_extra["camera_pos_y"] = str(camera_pos_y)
                _cam_extra["camera_pos_z"] = str(camera_pos_z)
            if display_amp_str:
                _cam_extra["display_amp"] = display_amp_str
            save_display_npz(
                _npz_path,
                _d["frames_x"], _d["frames_y"], _d["frames_z"],
                _d["frames_vx"], _d["frames_vy"], _d["frames_vz"],
                _d["atom_ids"],
                header_fields={**_d["header_fields"],
                                **_cam_extra,
                                "number_of_frames": str(_d["n_total_frames"]),
                                "number_of_particles": str(_d["n_total_particles"])})
            print(f"[compute] display.npz 已生成: {_npz_path}")
        except Exception as _e:
            print(f"[compute] 警告: display.npz 生成失败({_e})，将使用 display.txt")
        return None, None, None, None, None, None
    
    # save_trajectory=1：加载完整 trajectory.txt
    traj_path = os.path.join(os.path.abspath(output_dir), "trajectory.txt")
    if not os.path.exists(traj_path):
        raise FileNotFoundError(f"引擎未生成 trajectory.txt: {traj_path}")
    data = load_text_data(traj_path)
    traj_x  = data["traj_x"]
    traj_y  = data["traj_y"]
    traj_z  = data["traj_z"]
    traj_vx = data["traj_vx"]
    traj_vy = data["traj_vy"]
    traj_vz = data["traj_vz"]

    # 写入日志文件
    log_path = os.path.join(output_dir, "dynamics.log")
    with open(log_path, "w", encoding="utf-8") as f:
        f.write("=" * 50 + "\n")
        f.write("Dynamics 计算日志\n")
        f.write("=" * 50 + "\n")
        f.write(f"引擎:       {engine}\n")
        f.write(f"计算开始时间: {t_start_str}\n")
        f.write(f"计算结束时间: {t_end_str}\n")
        f.write(f"计算耗时:     {elapsed:.3f} s\n")

    print(f"[compute] 引擎完成: {len(traj_x)} 步, {traj_x.shape[1]} 原子, {elapsed:.3f} s")
    return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz


def save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, out_dir=None):
    """将轨迹数据保存到结构化 txt 文件（含所有参数元数据）。"""
    if out_dir is None:
        out_dir = os.path.dirname(os.path.abspath(__file__))
    out_path = os.path.join(get_output_dir(out_dir), "trajectory.txt")

    # 记录的实际步数（扣除预热后）
    record_steps = len(traj_x)
    sample_start_to_save = -1 if sample_start is None else int(sample_start)
    sample_end_to_save = -1 if sample_end is None else int(sample_end)

    payload = {
        "traj_x": traj_x, "traj_y": traj_y, "traj_z": traj_z,
        "traj_vx": traj_vx, "traj_vy": traj_vy, "traj_vz": traj_vz,
        "NT": record_steps, "DT": DT, "NSTEP": NSTEP,
        "method": METHOD,
        "coord_file": COORD_FILE,
        "connection_file": BOND_CONNECTION_FILE,
        "bond_file": BOND_PARAMETER_FILE,
        "atom_ids": ATOM_IDS,
        "atom_masses": ATOM_MASSES,
        "atom_radii": ATOM_RADII,
        "atom_positions": ATOM_POSITIONS,
        "atom_velocities": ATOM_VELOCITIES,
        "atom_fixed": ATOM_FIXED,
        "bond_pairs": BOND_PAIRS,
        "bond_names": BOND_NAMES,
        "bond_stiffness": BOND_STIFFNESS,
        "bond_rest_lengths": BOND_REST_LENGTHS,
        "G": G.tolist() if G is not None else None,
        "B": B.tolist() if B is not None else None,
        "plot_atom_id": PLOT_ATOM_ID,
        "plot_atom_row": PLOT_ATOM_ROW,
        "warmup_steps": warmup_steps or 0,
        "sample_start": sample_start_to_save,
        "sample_end": sample_end_to_save,
        "X_MIN": X_MIN, "X_MAX": X_MAX,
        "Y_MIN": Y_MIN, "Y_MAX": Y_MAX,
        "Z_MIN": Z_MIN, "Z_MAX": Z_MAX,
        "X0": X0, "Y0": Y0, "Z0": Z0,
        "VX0": VX0, "VY0": VY0, "VZ0": VZ0,
        "M": M,
        "alpha": alpha,
        "ball_radius": ball_radius,
        "ball_color_r": ball_color_r,
        "ball_color_g": ball_color_g,
        "ball_color_b": ball_color_b,
        "box_color_r": box_color_r,
        "box_color_g": box_color_g,
        "box_color_b": box_color_b,
        "gravity_field": GRAVITY_FIELD,
        "gravity_interaction": GRAVITY_INTERACTION,
        "elastic_force": ELASTIC_FORCE,
        "damping_force": DAMPING_FORCE,
        "gravity_strength": GRAVITY_STRENGTH,
    }
    save_text_data(out_path, payload)
    print(f"[compute] 轨迹数据已保存至: {out_path}")
    return out_path


def load_parameters(param_file):
    """从文本文件加载参数，文件格式：参数 = 值"""
    global box_a, alpha, X0, Y0, Z0, VX0, VY0, VZ0, M, G, B, METHOD, NT, DT, NSTEP
    global X_MIN, X_MAX, Y_MIN, Y_MAX, Z_MIN, Z_MAX
    global ball_radius, ball_color_r, ball_color_g, ball_color_b
    global box_color_r, box_color_g, box_color_b

    print(f"[compute] 正在加载参数文件: {param_file}")
    with open(param_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    for line in lines:
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        if '=' not in line:
            continue
        # 分割键值，并去除行内注释
        key, value = line.split('=', 1)
        key = key.strip()
        value = value.strip()
        # 去除值中可能存在的注释（# 之后的部分）
        if '#' in value:
            value = value.split('#', 1)[0].strip()
        # 尝试转换为 int 或 float
        try:
            if '.' in value or 'e' in value.lower():
                value = float(value)
            else:
                value = int(value)
        except ValueError:
            pass  # 保持字符串
        
        # 设置全局变量
        if key in globals():
            globals()[key] = value
        else:
            print(f"警告: 未知参数 '{key}'，忽略")
    
    # 计算派生边界
    if box_a is not None:
        X_MIN = -box_a
        X_MAX =  box_a
        Y_MIN = -box_a
        Y_MAX =  box_a
        Z_MIN = -box_a
        Z_MAX =  box_a
    else:
        print("错误: 未找到 box_a 参数")
        sys.exit(1)
    
    # 检查必需参数
    required = ['alpha', 'X0', 'Y0', 'Z0', 'VX0', 'VY0', 'VZ0', 'M', 'G', 'B', 'NT', 'DT', 'NSTEP']
    for param in required:
        if globals()[param] is None:
            print(f"错误: 未找到必需参数 '{param}'")
            sys.exit(1)

    METHOD = normalize_method_name(METHOD or "explicit_euler")
    
    # 设置绘图参数的默认值（如果未提供）
    if ball_radius is None:
        ball_radius = 0.28
    if ball_color_r is None:
        ball_color_r = 0.9
    if ball_color_g is None:
        ball_color_g = 0.2
    if ball_color_b is None:
        ball_color_b = 0.2
    if box_color_r is None:
        box_color_r = 0.8
    if box_color_g is None:
        box_color_g = 0.8
    if box_color_b is None:
        box_color_b = 0.85
    
    print(f"[compute] 参数加载完成: box_a={box_a}, alpha={alpha}, NT={NT}, DT={DT}, method={METHOD}")

def normalize_method_name(method):
    """Normalize method aliases from YAML/text config."""
    aliases = {
        "explicit_euler": "explicit_euler",
        "explicit": "explicit_euler",
        "euler_explicit": "explicit_euler",
        "显式欧拉": "explicit_euler",
        "显式欧拉法": "explicit_euler",
        "implicit_euler": "implicit_euler",
        "implicit": "implicit_euler",
        "euler_implicit": "implicit_euler",
        "隐式欧拉": "implicit_euler",
        "隐式欧拉法": "implicit_euler",
        "midpoint": "midpoint",
        "mid_point": "midpoint",
        "midpoint_method": "midpoint",
        "中点法": "midpoint",
        "中点差分法": "midpoint",
        "leapfrog": "leapfrog",
        "leap_frog": "leapfrog",
        "蛙跳法": "leapfrog",
    }
    key = str(method).strip().lower()
    if key not in aliases:
        valid = ", ".join(["explicit_euler", "implicit_euler", "midpoint", "leapfrog"])
        raise ValueError(f"未知算法 '{method}'。可选值: {valid}")
    return aliases[key]


def compute_force(x, y, z, vx, vy, vz, m, g, b):
    """Compute total force from the current state.

    Each force type is independently controlled by a global switch:
    - GRAVITY_FIELD: uniform gravity F = m * g_vec
    - DAMPING_FORCE: linear drag F_drag = -B_vec * v
    - ELASTIC_FORCE: spring bonds based on bonded pair distance
    - GRAVITY_INTERACTION: Newtonian gravity between atom pairs
    """
    fx = np.zeros_like(x)
    fy = np.zeros_like(y)
    fz = np.zeros_like(z)

    if GRAVITY_FIELD:
        fx += m * g[0]
        fy += m * g[1]
        fz += m * g[2]

    if DAMPING_FORCE:
        fx -= b[0] * vx
        fy -= b[1] * vy
        fz -= b[2] * vz

    if ELASTIC_FORCE and BOND_PAIRS is not None and len(BOND_PAIRS) > 0:
        idx_1 = BOND_PAIRS[:, 0]
        idx_2 = BOND_PAIRS[:, 1]
        dx = x[idx_2] - x[idx_1]
        dy = y[idx_2] - y[idx_1]
        dz = z[idx_2] - z[idx_1]
        dist = np.sqrt(dx * dx + dy * dy + dz * dz)
        valid = dist > 1e-12
        if np.any(valid):
            force_scale = np.zeros_like(dist)
            stretch = dist[valid] - BOND_REST_LENGTHS[valid]
            force_scale[valid] = BOND_STIFFNESS[valid] * stretch / dist[valid]

            fx_bond = force_scale * dx
            fy_bond = force_scale * dy
            fz_bond = force_scale * dz

            np.add.at(fx, idx_1, fx_bond)
            np.add.at(fx, idx_2, -fx_bond)
            np.add.at(fy, idx_1, fy_bond)
            np.add.at(fy, idx_2, -fy_bond)
            np.add.at(fz, idx_1, fz_bond)
            np.add.at(fz, idx_2, -fz_bond)

    if GRAVITY_INTERACTION:
        n = len(m)
        for i in range(n):
            for j in range(i + 1, n):
                dx = x[j] - x[i]
                dy = y[j] - y[i]
                dz = z[j] - z[i]
                r2 = dx * dx + dy * dy + dz * dz
                if r2 <= 1e-12:
                    continue
                mi = m[i]; mj = m[j]
                f_mag = GRAVITY_STRENGTH * mi * mj / r2
                r = np.sqrt(r2)
                fx_i = f_mag * dx / r
                fy_i = f_mag * dy / r
                fz_i = f_mag * dz / r
                fx[i] += fx_i; fy[i] += fy_i; fz[i] += fz_i
                fx[j] -= fx_i; fy[j] -= fy_i; fz[j] -= fz_i

    return fx, fy, fz


def compute_acceleration(x, y, z, vx, vy, vz, m, g, b):
    """Compute acceleration from the shared force model."""
    fx, fy, fz = compute_force(x, y, z, vx, vy, vz, m, g, b)
    return fx / m, fy / m, fz / m


def compute_accel_conservative(x, y, z, m, g):
    """保守力加速度（弹簧 + 重力 + 原子间引力），不含阻尼。
    供真蛙跳法使用：阻尼项由 leapfrog_staggered_step 半隐式处理。
    """
    _v0 = np.zeros_like(x)
    _b0 = np.zeros(3)
    fx, fy, fz = compute_force(x, y, z, _v0, _v0, _v0, m, g, _b0)
    return fx / m, fy / m, fz / m


def leapfrog_staggered_step(x, y, z, vx_h, vy_h, vz_h, dt, m, g, b):
    """真蛙跳一步：x(t), v(t-dt/2) → x(t+dt), v(t+dt/2)

    - 无阻尼 (DAMPING_FORCE=False 或 b≈0)：
        纯保守蛙跳，每步仅 1 次力计算，辛积分器。
        v(t+dt/2) = v(t-dt/2) + a_c(t)·dt

    - 有阻尼 (DAMPING_FORCE=True 且 b≠0)：
        半隐式处理阻尼，仍只 1 次力计算，对任意阻尼无条件稳定。
        利用 v(t) ≈ [v(t-dt/2) + v(t+dt/2)] / 2 解析求解：
        v(t+dt/2) = [v(t-dt/2)·(1 - α) + a_c(t)·dt] / (1 + α)
        其中 α = b·dt / (2·m)，每个方向独立。
    """
    ax_c, ay_c, az_c = compute_accel_conservative(x, y, z, m, g)

    has_damping = DAMPING_FORCE and np.any(b != 0)
    if has_damping:
        alpha_x = b[0] * dt / (2.0 * m)
        alpha_y = b[1] * dt / (2.0 * m)
        alpha_z = b[2] * dt / (2.0 * m)
        vx_h_new = (vx_h * (1.0 - alpha_x) + ax_c * dt) / (1.0 + alpha_x)
        vy_h_new = (vy_h * (1.0 - alpha_y) + ay_c * dt) / (1.0 + alpha_y)
        vz_h_new = (vz_h * (1.0 - alpha_z) + az_c * dt) / (1.0 + alpha_z)
    else:
        vx_h_new = vx_h + ax_c * dt
        vy_h_new = vy_h + ay_c * dt
        vz_h_new = vz_h + az_c * dt

    x_new = x + vx_h_new * dt
    y_new = y + vy_h_new * dt
    z_new = z + vz_h_new * dt
    return x_new, y_new, z_new, vx_h_new, vy_h_new, vz_h_new


def Explicit_Euler_Method(x, y, z, vx, vy, vz, dt, m, g, b):
    ax, ay, az = compute_acceleration(x, y, z, vx, vy, vz, m, g, b)
    x = x + vx * dt
    y = y + vy * dt
    z = z + vz * dt
    vx = vx + ax * dt
    vy = vy + ay * dt
    vz = vz + az * dt
    return x, y, z, vx, vy, vz


def Implicit_Euler_Method(x, y, z, vx, vy, vz, dt, m, g, b):
    gamma_x = b[0] / m
    gamma_y = b[1] / m
    gamma_z = b[2] / m
    vx_next = (vx + g[0] * dt) / (1.0 + gamma_x * dt)
    vy_next = (vy + g[1] * dt) / (1.0 + gamma_y * dt)
    vz_next = (vz + g[2] * dt) / (1.0 + gamma_z * dt)
    ax_next, ay_next, az_next = compute_acceleration(
        x, y, z, vx_next, vy_next, vz_next, m, g, b)

    vx = vx + ax_next * dt
    vy = vy + ay_next * dt
    vz = vz + az_next * dt
    x = x + vx * dt
    y = y + vy * dt
    z = z + vz * dt
    return x, y, z, vx, vy, vz


def Midpoint_Method(x, y, z, vx, vy, vz, dt, m, g, b):
    ax, ay, az = compute_acceleration(x, y, z, vx, vy, vz, m, g, b)
    x_mid = x + 0.5 * vx * dt
    y_mid = y + 0.5 * vy * dt
    z_mid = z + 0.5 * vz * dt
    vx_mid = vx + 0.5 * ax * dt
    vy_mid = vy + 0.5 * ay * dt
    vz_mid = vz + 0.5 * az * dt

    x = x + vx_mid * dt
    y = y + vy_mid * dt
    z = z + vz_mid * dt

    ax_mid, ay_mid, az_mid = compute_acceleration(
        x_mid, y_mid, z_mid, vx_mid, vy_mid, vz_mid, m, g, b)
    vx = vx + ax_mid * dt
    vy = vy + ay_mid * dt
    vz = vz + az_mid * dt
    return x, y, z, vx, vy, vz


def Leapfrog_Method(x, y, z, vx, vy, vz, dt, m, g, b):
    ax, ay, az = compute_acceleration(x, y, z, vx, vy, vz, m, g, b)
    vx_half = vx + 0.5 * ax * dt
    vy_half = vy + 0.5 * ay * dt
    vz_half = vz + 0.5 * az * dt

    x = x + vx_half * dt
    y = y + vy_half * dt
    z = z + vz_half * dt

    # 显式预测器：用原始加速度外推半步（标准 Velocity-Verlet 预测步）
    # v_pred = v_half + 0.5 * a(t)*dt，包含重力+阻尼+弹簧的所有贡献
    vx_pred = vx_half + 0.5 * ax * dt
    vy_pred = vy_half + 0.5 * ay * dt
    vz_pred = vz_half + 0.5 * az * dt
    ax_next, ay_next, az_next = compute_acceleration(
        x, y, z, vx_pred, vy_pred, vz_pred, m, g, b)
    vx = vx_half + 0.5 * ax_next * dt
    vy = vy_half + 0.5 * ay_next * dt
    vz = vz_half + 0.5 * az_next * dt
    return x, y, z, vx, vy, vz

def Limit_in_box(a, amin, amax, va):
    """限制物体在边界内，发生碰撞时反弹。"""
    over = a > amax
    under = a < amin
    a = np.where(over, amax, a)
    a = np.where(under, amin, a)
    va = np.where(over | under, -va, va)
    return a, va

def apply_motion_update(x, y, z, vx, vy, vz, dt, m, g, b):
    """按配置选择的位置更新算法推进一步。"""
    if METHOD == "explicit_euler":
        x, y, z, vx, vy, vz = Explicit_Euler_Method(x, y, z, vx, vy, vz, dt, m, g, b)
    elif METHOD == "implicit_euler":
        x, y, z, vx, vy, vz = Implicit_Euler_Method(x, y, z, vx, vy, vz, dt, m, g, b)
    elif METHOD == "midpoint":
        x, y, z, vx, vy, vz = Midpoint_Method(x, y, z, vx, vy, vz, dt, m, g, b)
    elif METHOD == "leapfrog":
        x, y, z, vx, vy, vz = Leapfrog_Method(x, y, z, vx, vy, vz, dt, m, g, b)
    else:
        raise ValueError(f"未知算法: {METHOD}")
    return x, y, z, vx, vy, vz


def apply_fixed_constraints(x, y, z, vx, vy, vz):
    """Keep fixed degrees of freedom at their initial coordinate with zero speed."""
    if FIXED_MASK_X is not None and np.any(FIXED_MASK_X):
        x[FIXED_MASK_X] = ATOM_POSITIONS[FIXED_MASK_X, 0]
        vx[FIXED_MASK_X] = 0.0
    if FIXED_MASK_Y is not None and np.any(FIXED_MASK_Y):
        y[FIXED_MASK_Y] = ATOM_POSITIONS[FIXED_MASK_Y, 1]
        vy[FIXED_MASK_Y] = 0.0
    if FIXED_MASK_Z is not None and np.any(FIXED_MASK_Z):
        z[FIXED_MASK_Z] = ATOM_POSITIONS[FIXED_MASK_Z, 2]
        vz[FIXED_MASK_Z] = 0.0
    return x, y, z, vx, vy, vz

def wrap_position(x, y, z):
    """边界回绕（ dynamics 模式）。"""
    x = np.where(x > X_MAX, X_MIN, x)
    x = np.where(x < X_MIN, X_MAX, x)
    y = np.where(y > Y_MAX, Y_MIN, y)
    y = np.where(y < Y_MIN, Y_MAX, y)
    z = np.where(z > Z_MAX, Z_MIN, z)
    z = np.where(z < Z_MIN, Z_MAX, z)
    return x, y, z


# ===========================================================================
# 主计算流程
# ===========================================================================

def run_simulation(save_trajectory=0):
    """计算 NT 步轨迹，直接抽帧并保存 display.txt。

    步骤控制：
    - warmup_steps: 预热步数，跳过不记录（用于稳定初始状态）
    - 按 NSTEP 抽帧保存到 display.txt（新格式）
    - save_trajectory=1 时额外保存完整 trajectory.txt（JSON 旧格式）
    """
    # 预热阶段
    x = ATOM_POSITIONS[:, 0].copy()
    y = ATOM_POSITIONS[:, 1].copy()
    z = ATOM_POSITIONS[:, 2].copy()
    vx = ATOM_VELOCITIES[:, 0].copy()
    vy = ATOM_VELOCITIES[:, 1].copy()
    vz = ATOM_VELOCITIES[:, 2].copy()
    x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)
    x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, 0.0, 0, DRIVER_DATA, DT)

    # 真蛙跳初始化：从 v(0) 反推 v(-dt/2)，使后续每步只需 1 次力计算。
    # 其他方法（euler/midpoint）vx/vy/vz 始终存整步速度，不受影响。
    if METHOD == "leapfrog":
        _ax0, _ay0, _az0 = compute_accel_conservative(x, y, z, ATOM_MASSES, G)
        vx = vx - 0.5 * _ax0 * DT
        vy = vy - 0.5 * _ay0 * DT
        vz = vz - 0.5 * _az0 * DT
        # 从此 vx/vy/vz 存 v(t-dt/2)（蛙跳半步速度）

    if warmup_steps is not None and warmup_steps > 0:
        print(f"[compute] 预热阶段: 前 {warmup_steps} 步不记录")
        for step in trange(warmup_steps, desc="[compute] 预热"):
            t = (step + 1) * DT
            x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, t, step, DRIVER_DATA, DT)
            if METHOD == "leapfrog":
                x, y, z, vx, vy, vz = leapfrog_staggered_step(
                    x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
            else:
                x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
            x, vx = Limit_in_box(x, X_MIN, X_MAX, vx)
            y, vy = Limit_in_box(y, Y_MIN, Y_MAX, vy)
            z, vz = Limit_in_box(z, Z_MIN, Z_MAX, vz)
            x, y, z = wrap_position(x, y, z)
            x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)
        print(
            f"[compute] 预热完成，展示原子位置: "
            f"({x[PLOT_ATOM_ROW]:.4f}, {y[PLOT_ATOM_ROW]:.4f}, {z[PLOT_ATOM_ROW]:.4f})"
        )

    # 记录阶段 - 按 NSTEP 抽帧保存
    record_steps = NT - (warmup_steps or 0)
    n_atoms = len(ATOM_IDS)
    n_frames = (record_steps + NSTEP - 1) // NSTEP
    # 按 NSTEP 抽帧的临时缓冲区（远小于全量轨迹）
    sampled_x  = np.zeros((n_frames, n_atoms), dtype=np.float64)
    sampled_y  = np.zeros((n_frames, n_atoms), dtype=np.float64)
    sampled_z  = np.zeros((n_frames, n_atoms), dtype=np.float64)
    sampled_vx = np.zeros((n_frames, n_atoms), dtype=np.float64)
    sampled_vy = np.zeros((n_frames, n_atoms), dtype=np.float64)
    sampled_vz = np.zeros((n_frames, n_atoms), dtype=np.float64)

    # 如果 save_trajectory，准备完整轨迹缓冲区
    if save_trajectory:
        traj_x  = np.zeros((record_steps, n_atoms), dtype=np.float64)
        traj_y  = np.zeros((record_steps, n_atoms), dtype=np.float64)
        traj_z  = np.zeros((record_steps, n_atoms), dtype=np.float64)
        traj_vx = np.zeros((record_steps, n_atoms), dtype=np.float64)
        traj_vy = np.zeros((record_steps, n_atoms), dtype=np.float64)
        traj_vz = np.zeros((record_steps, n_atoms), dtype=np.float64)

    for step in trange(record_steps, desc="[compute] 计算中"):
        t = (step + (warmup_steps or 0)) * DT
        x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, t, step, DRIVER_DATA, DT)

        if save_trajectory:
            traj_x[step]  = x
            traj_y[step]  = y
            traj_z[step]  = z
            traj_vx[step] = vx
            traj_vy[step] = vy
            traj_vz[step] = vz

        # 抽帧：NSTEP 间隔保存
        if step % NSTEP == 0:
            fi = step // NSTEP
            sampled_x[fi]  = x
            sampled_y[fi]  = y
            sampled_z[fi]  = z
            sampled_vx[fi] = vx
            sampled_vy[fi] = vy
            sampled_vz[fi] = vz
        if METHOD == "leapfrog":
            x, y, z, vx, vy, vz = leapfrog_staggered_step(
                x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
        else:
            x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
        x, vx = Limit_in_box(x, X_MIN, X_MAX, vx)
        y, vy = Limit_in_box(y, Y_MIN, Y_MAX, vy)
        z, vz = Limit_in_box(z, Z_MIN, Z_MAX, vz)
        x, y, z = wrap_position(x, y, z)
        x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)

    # 写入 display.txt（新格式）
    output_dir = get_output_dir()
    disp_path = os.path.join(output_dir, "display.txt")
    n_frames_actual = (record_steps + NSTEP - 1) // NSTEP
    save_display_txt(
        disp_path,
        sampled_x[:n_frames_actual], sampled_y[:n_frames_actual], sampled_z[:n_frames_actual],
        sampled_vx[:n_frames_actual], sampled_vy[:n_frames_actual], sampled_vz[:n_frames_actual],
        np.array(ATOM_IDS), record_steps, n_atoms,
        header_fields={"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD),
                       "warmup_steps": str(warmup_steps or 0),
                       "dynamic_steps": str(record_steps),
                       "T_total": str(NT * DT),
                       "X_MAX": str(X_MAX), "X_MIN": str(X_MIN),
                       "Y_MAX": str(Y_MAX), "Y_MIN": str(Y_MIN),
                       "Z_MAX": str(Z_MAX), "Z_MIN": str(Z_MIN),
                       "ball_radius": str(ball_radius),
                       "ball_color_r": str(ball_color_r),
                       "ball_color_g": str(ball_color_g),
                       "ball_color_b": str(ball_color_b),
                       "box_color_r": str(box_color_r),
                       "box_color_g": str(box_color_g),
                       "box_color_b": str(box_color_b),
                       "gravity_field": str(GRAVITY_FIELD),
                       "gravity_interaction": str(GRAVITY_INTERACTION),
                       "elastic_force": str(ELASTIC_FORCE),
                       "damping_force": str(DAMPING_FORCE),
                       "gravity_strength": str(GRAVITY_STRENGTH),
                       "driving_force": str(DRIVING_FORCE),
                       "use_marker": str(use_marker),
                       "alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])),
                       "atom_masses": json.dumps([float(v) for v in ATOM_MASSES]),
                       "atom_positions": json.dumps(ATOM_POSITIONS.tolist()),
                       "bond_pairs": json.dumps(BOND_PAIRS.tolist() if BOND_PAIRS is not None else []),
                       "bond_stiffness": json.dumps(BOND_STIFFNESS.tolist() if BOND_STIFFNESS is not None else []),
                       "bond_rest_lengths": json.dumps(BOND_REST_LENGTHS.tolist() if BOND_REST_LENGTHS is not None else []),
                       "G": json.dumps(G.tolist() if G is not None else [0.0, 0.0, 0.0]),
                       "atom_radii": ",".join(str(r) for r in ATOM_RADII),
                       "camera_distance": str(camera_distance),
                       "camera_elevation": str(camera_elevation),
                       "camera_azimuth": str(camera_azimuth),
                       "camera_center_x": str(camera_center_x),
                       "camera_center_y": str(camera_center_y),
                       "camera_center_z": str(camera_center_z),
                       **({
                           "camera_pos_x": str(camera_pos_x),
                           "camera_pos_y": str(camera_pos_y),
                           "camera_pos_z": str(camera_pos_z),
                       } if camera_pos_x is not None else {}),
                       "camera_keyframes": str(camera_keyframes_raw),
                       **({"display_amp": display_amp_str} if display_amp_str else {})}
    )
    print(f"[compute] display.txt 已保存至: {disp_path} ({n_frames_actual} 帧)")

    # 同时保存二进制 display.npz（I/O 速度提升 5-10x）
    _npz_path = os.path.join(output_dir, "display.npz")
    _hdr = {"DT": str(DT), "NSTEP": str(NSTEP), "method": str(METHOD),
            "warmup_steps": str(warmup_steps or 0),
            "dynamic_steps": str(record_steps),
            "T_total": str(NT * DT),
            "X_MAX": str(X_MAX), "X_MIN": str(X_MIN),
            "Y_MAX": str(Y_MAX), "Y_MIN": str(Y_MIN),
            "Z_MAX": str(Z_MAX), "Z_MIN": str(Z_MIN),
            "ball_radius": str(ball_radius),
            "ball_color_r": str(ball_color_r),
            "ball_color_g": str(ball_color_g),
            "ball_color_b": str(ball_color_b),
            "box_color_r": str(box_color_r),
            "box_color_g": str(box_color_g),
            "box_color_b": str(box_color_b),
            "gravity_field": str(GRAVITY_FIELD),
            "gravity_interaction": str(GRAVITY_INTERACTION),
            "elastic_force": str(ELASTIC_FORCE),
            "damping_force": str(DAMPING_FORCE),
            "gravity_strength": str(GRAVITY_STRENGTH),
            "driving_force": str(DRIVING_FORCE),
            "use_marker": str(use_marker),
            "alpha": ",".join(str(a) for a in (alpha if isinstance(alpha, list) else [alpha])),
            "atom_masses": json.dumps([float(v) for v in ATOM_MASSES]),
            "atom_positions": json.dumps(ATOM_POSITIONS.tolist()),
            "bond_pairs": json.dumps(BOND_PAIRS.tolist() if BOND_PAIRS is not None else []),
            "bond_stiffness": json.dumps(BOND_STIFFNESS.tolist() if BOND_STIFFNESS is not None else []),
            "bond_rest_lengths": json.dumps(BOND_REST_LENGTHS.tolist() if BOND_REST_LENGTHS is not None else []),
            "G": json.dumps(G.tolist() if G is not None else [0.0, 0.0, 0.0]),
            "atom_radii": ",".join(str(r) for r in ATOM_RADII),
            "camera_distance": str(camera_distance),
            "camera_elevation": str(camera_elevation),
            "camera_azimuth": str(camera_azimuth),
            "camera_center_x": str(camera_center_x),
            "camera_center_y": str(camera_center_y),
            "camera_center_z": str(camera_center_z),
            "camera_keyframes": str(camera_keyframes_raw),
            "number_of_frames": str(record_steps),
            "number_of_particles": str(n_atoms),
            **({"display_amp": display_amp_str} if display_amp_str else {})}
    save_display_npz(_npz_path,
                     sampled_x[:n_frames_actual], sampled_y[:n_frames_actual],
                     sampled_z[:n_frames_actual],
                     sampled_vx[:n_frames_actual], sampled_vy[:n_frames_actual],
                     sampled_vz[:n_frames_actual],
                     np.array(ATOM_IDS), header_fields=_hdr)
    print(f"[compute] display.npz 已保存至: {_npz_path}")

    # 可选：保存完整 trajectory.txt
    if save_trajectory:
        save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz)
        print(f"[compute] trajectory.txt 已保存（完整轨迹）")

    return sampled_x[:n_frames_actual], sampled_y[:n_frames_actual], sampled_z[:n_frames_actual], \
           sampled_vx[:n_frames_actual], sampled_vy[:n_frames_actual], sampled_vz[:n_frames_actual]


def save_trajectory_table_txt(txt_path, traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT):
    """将轨迹数据保存为逐行表格文本，便于教学查看。"""
    if np.ndim(traj_x) > 1:
        row = PLOT_ATOM_ROW if PLOT_ATOM_ROW is not None else 0
        traj_x = traj_x[:, row]
        traj_y = traj_y[:, row]
        traj_z = traj_z[:, row]
        traj_vx = traj_vx[:, row]
        traj_vy = traj_vy[:, row]
        traj_vz = traj_vz[:, row]

    with open(txt_path, 'w', encoding='utf-8') as f:
        # 写入表头
        f.write("# 步数, 时间(s), x, y, z, vx, vy, vz\n")
        for step in range(NT):
            t = step * DT
            x = traj_x[step]
            y = traj_y[step]
            z = traj_z[step]
            vx = traj_vx[step]
            vy = traj_vy[step]
            vz = traj_vz[step]
            f.write(f"{step+1:6d}, {t:.6f}, {x:.6f}, {y:.6f}, {z:.6f}, {vx:.6f}, {vy:.6f}, {vz:.6f}\n")
    print(f"[compute] 轨迹文本文件已保存至: {txt_path}")


def main():
    # 默认参数文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    param_file = os.path.join(get_input_dir(script_dir), "input.txt")
    if len(sys.argv) > 1:
        param_file = sys.argv[1]
    
    # 加载参数
    load_parameters(param_file)
    
    script_dir = os.path.dirname(os.path.abspath(__file__))
    output_dir = get_output_dir(script_dir)

    print(f"[compute] 开始计算  NT={NT}  DT={DT} ")
    traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz = run_simulation(save_trajectory=0)
    print(f"[compute] 计算完成，共 {NT} 步")
    print(f"[compute] display.txt 已在 run_simulation 中保存")

    # 如果需要完整轨迹，以上传 save_trajectory=1 重新运行
    # 以下旧函数保留兼容但不再自动调用
    # save_trajectory_txt(...)

    # 同时保存为逐行表格，便于直接查看
    txt_path = os.path.join(output_dir, "trajectory_table.txt")
    save_trajectory_table_txt(txt_path, traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT)


if __name__ == "__main__":
    main()
