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

功能：
  1. 运行 NT 步物理模拟（ kinematics / dynamics 等运动模式）
  2. 将每一步的 (x, y, z, vx, vy, vz) 保存到 output/trajectory.txt
  3. 同时保存所有模拟参数元数据
"""

import json
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

# 力开关
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 _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)


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),
                # 在模拟中动态记录：冻结步数索引、冻结位置
                "freeze_step": None,
                "freeze_pos": None,
            })

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

    print(f"[compute] 已加载驱动力: {len(drivers)} 条定义")
    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']}")
    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  # 全程驱动

        # 当前驱动力下的位置 / 速度
        t_vec = np.array([t, t, t], dtype=np.float64)
        pos_drive = d["amp"] * np.cos(2.0 * np.pi * d["freq"] * t_vec + d["phi"])
        vel_drive = -d["amp"] * 2.0 * np.pi * d["freq"] * np.sin(2.0 * np.pi * d["freq"] * t_vec + 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 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)
    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))

    # 力开关
    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)

    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()
    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(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/c/build/dynamics_c",
        "cpp": "engines/cpp/build/dynamics_cpp",
        "fortran": "engines/fortran/build/dynamics_f90",
    }
    if engine not in engine_map:
        raise ValueError(f"不支持的引擎: {engine}，可选: {list(engine_map.keys())}")

    engine_rel = engine_map[engine]
    engine_path = os.path.join(script_dir, engine_rel)

    # 自动检测可执行文件后缀和平台专用版本
    candidates = [
        engine_path,                          # 无后缀
        engine_path + ".exe",                 # Windows .exe
        engine_path + f"_{system}.exe",       # 平台专用 (c_linux.exe, c_darwin.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)),
    }
    param_path = os.path.join(script_dir, "engines", engine, "param.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}")

    # ── 预校准：跑 NT=1000 步测速 ────────────────────────────────
    total_steps = int(config["NT"]) - int(config.get("warmup_steps", 0))
    _calib_nt = min(1000, max(100, total_steps // 10))
    _calib_param = dict(param_json)
    _calib_param["NT"] = _calib_nt
    _calib_path = os.path.join(script_dir, "engines", engine, "_calib.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", engine, "_calib_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
    _est_total = max(_calib_elapsed, _overhead + _step_time * total_steps)

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

    _p = subprocess.Popen(
        [engine_path, os.path.abspath(input_dir), os.path.abspath(output_dir), param_path],
        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)
            if _p.poll() is not None:
                if _p.stdout:
                    for _r in _p.stdout:
                        _r = _r.strip()
                        if _r:
                            _engine_lines.append(_r)
                break
            if _pbar is not None and _est_total > 0:
                _pbar.n = int(min((time.time() - t_start) / _est_total, 0.99) * total_steps)
                _pbar.refresh()
            time.sleep(0.2)
    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
    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:
        for bond_idx, (idx_1, idx_2) in enumerate(BOND_PAIRS):
            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)
            if dist <= 1e-12:
                continue

            stretch = dist - BOND_REST_LENGTHS[bond_idx]
            force_mag = BOND_STIFFNESS[bond_idx] * stretch
            ux = dx / dist
            uy = dy / dist
            uz = dz / dist

            fx_bond = force_mag * ux
            fy_bond = force_mag * uy
            fz_bond = force_mag * uz

            fx[idx_1] += fx_bond
            fy[idx_1] += fy_bond
            fz[idx_1] += fz_bond
            fx[idx_2] -= fx_bond
            fy[idx_2] -= fy_bond
            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 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}")

    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)
    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."""
    fixed = ATOM_FIXED != 0
    positions = np.column_stack((x, y, z))
    velocities = np.column_stack((vx, vy, vz))
    positions = np.where(fixed, ATOM_POSITIONS, positions)
    velocities = np.where(fixed, 0.0, velocities)
    return (
        positions[:, 0], positions[:, 1], positions[:, 2],
        velocities[:, 0], velocities[:, 1], velocities[:, 2],
    )

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():
    """计算 NT 步轨迹，返回位置/速度数组。

    步骤控制：
    - warmup_steps: 预热步数，跳过不记录（用于稳定初始状态）
    - 实际记录步数 = NT - warmup_steps
    """
    # 预热阶段
    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)
    # 初始时刻驱动力（t=0 时原子 1 的位置由驱动力决定而非 coord.txt）
    x, y, z, vx, vy, vz = apply_driving_force(x, y, z, vx, vy, vz, 0.0, 0, DRIVER_DATA, DT)

    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)
            x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
            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})"
        )

    # 记录阶段
    record_steps = NT - (warmup_steps or 0)
    n_atoms = len(ATOM_IDS)
    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)

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

        x, y, z, vx, vy, vz = apply_motion_update(x, y, z, vx, vy, vz, DT, ATOM_MASSES, G, B)
        x, y, z = wrap_position(x, y, z)
        x, y, z, vx, vy, vz = apply_fixed_constraints(x, y, z, vx, vy, vz)

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


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()
    print(f"[compute] 计算完成，共 {NT} 步")

    save_trajectory_txt(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, script_dir)

    # 同时保存为逐行表格，便于直接查看
    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()
