"""VisPy 演示：加载预计算轨迹数据，驱动小球运动动画。

计算与显示完全分离：
  1. 运行 run_dynamics.py  → 生成 output/display.txt（新格式，直接抽帧）
  2. 本文件加载 output/display.txt，按帧播放动画

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

import numpy as np
import json
import os
import sys
from vispy import app, scene
from vispy.visuals.transforms import STTransform

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

# ===========================================================================
# 加载预计算轨迹
# ===========================================================================
script_dir = os.path.dirname(os.path.abspath(__file__))
if len(sys.argv) > 1:
    # 用户指定了输出目录
    output_dir = os.path.abspath(sys.argv[1])
else:
    output_dir = compute.get_output_dir(script_dir)
os.environ["DYNAMICS_OUTPUT_DIR"] = output_dir
disp_path = os.path.join(output_dir, "display.txt")
npz_path  = os.path.join(output_dir, "display.npz")

if not os.path.exists(npz_path) and not os.path.exists(disp_path):
    raise FileNotFoundError(
        f"找不到 display.npz 或 display.txt！\n"
        f"期望路径: {output_dir}\n"
        f"请先运行 compute.py 计算轨迹，再运行 sample.py 生成显示数组。\n"
        f"用法: python draw.py [案例输出目录]"
    )

# 优先读二进制 npz（加载速度约快 5-10x）
if os.path.exists(npz_path):
    disp_data = compute.load_display_npz(npz_path)
else:
    disp_data = compute.load_display_txt(disp_path)

# ── 从 input.txt 读取参数（替代 display.npz 中的 meta）──
try:
    import yaml
    _have_yaml = True
except ImportError:
    _have_yaml = False

input_dir = os.path.join(os.path.dirname(output_dir), "input")
input_path = os.path.join(input_dir, "input.txt")

if _have_yaml and os.path.exists(input_path):
    try:
        with open(input_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
    except Exception:
        config = {}
else:
    config = {}
# 兼容旧版：若 input.txt 不存在或解析失败，降级到 display.npz 的 meta
if not config:
    config = disp_data.get("header_fields", {})

# ── 从 coord.txt 读取平衡位置 ─────────────────
# 注意: config 中的 *_file 可能带 "input/" 前缀，但 input_dir 已指向 input/ 目录
_coord_file_raw = config.get("coord_file", "coord.txt")
coord_path = os.path.join(input_dir, os.path.basename(_coord_file_raw))
if os.path.exists(coord_path):
    try:
        _ids, _masses, _radii, _pos, _vel, _fixed = compute.load_coord_file(coord_path)
        EQ_POS = _pos                             # (n_atoms, 3)
        ATOM_FIXED = _fixed                        # (n_atoms, 3)
    except Exception:
        EQ_POS = None
        ATOM_FIXED = None
else:
    EQ_POS = None
    ATOM_FIXED = None

# ── 从 connection.txt 读取成键信息（若 meta 中没有）──
BOND_PAIRS = disp_data.get("bond_pairs", [])
if not BOND_PAIRS and 'bond_pairs' not in disp_data:
    conn_path = os.path.join(input_dir, os.path.basename(config.get("connection_file", "connection.txt")))
    if os.path.exists(conn_path):
        try:
            _pairs = []
            with open(conn_path, "r", encoding="utf-8") as _f:
                _f.readline()  # skip header
                for _line in _f:
                    _line = _line.strip()
                    if not _line or _line.startswith("#"):
                        continue
                    _parts = _line.split()
                    if len(_parts) >= 2:
                        _pairs.append([int(_parts[0]) - 1, int(_parts[1]) - 1])
            BOND_PAIRS = np.array(_pairs, dtype=np.int32) if _pairs else []
        except Exception:
            BOND_PAIRS = []
    BOND_PAIRS = BOND_PAIRS.tolist() if hasattr(BOND_PAIRS, 'tolist') else BOND_PAIRS

# 全原子帧数据
DISP_ALL_X  = disp_data["frames_x"]       # (n_frames, n_atoms)
DISP_ALL_Y  = disp_data["frames_y"]
DISP_ALL_Z  = disp_data["frames_z"]
DISP_ALL_VX = disp_data["frames_vx"]
DISP_ALL_VY = disp_data["frames_vy"]
DISP_ALL_VZ = disp_data["frames_vz"]

# 第一个原子的轨迹（用于信息显示）
DISP_X    = DISP_ALL_X[:, 0]
DISP_Y    = DISP_ALL_Y[:, 0]
DISP_Z    = DISP_ALL_Z[:, 0]
DISP_VX   = DISP_ALL_VX[:, 0]
DISP_VY   = DISP_ALL_VY[:, 0]
DISP_VZ   = DISP_ALL_VZ[:, 0]

N_FRAMES  = DISP_ALL_X.shape[0]
NT        = int(disp_data["n_total_frames"])
N_ATOMS   = int(disp_data["n_total_particles"])
DT        = float(config.get("DT", 0.001))

# 视觉位移放大：display_amp: [ax, ay, az]，对偏离平衡位置的位移乘以倍数
_damp_raw = config.get("display_amp", "")
if isinstance(_damp_raw, (list, tuple)):
    _damp = np.array(_damp_raw, dtype=np.float64)
elif isinstance(_damp_raw, str) and _damp_raw.strip():
    import ast as _ast
    _damp = np.array(_ast.literal_eval(_damp_raw.strip()), dtype=np.float64)
else:
    _damp = None
if _damp is not None and _damp.shape == (3,) and not np.allclose(_damp, 1.0):
        if EQ_POS is not None:
            _eq_x = EQ_POS[None, :, 0]   # coord.txt 平衡位置
            _eq_y = EQ_POS[None, :, 1]
            _eq_z = EQ_POS[None, :, 2]
        else:
            _eq_x = DISP_ALL_X[0:1, :]   # 第0帧作为平衡位置参考
            _eq_y = DISP_ALL_Y[0:1, :]
            _eq_z = DISP_ALL_Z[0:1, :]
        DISP_ALL_X = _eq_x + (DISP_ALL_X - _eq_x) * _damp[0]
        DISP_ALL_Y = _eq_y + (DISP_ALL_Y - _eq_y) * _damp[1]
        DISP_ALL_Z = _eq_z + (DISP_ALL_Z - _eq_z) * _damp[2]

# ── 位移颜色映射 ──────────────────────────────
# display_color: { mode: [enabled, [R,G,B]], ... }
#   mode: x, y, z, xy, yz, zx, xyz
#   enabled: 0=off, 1=on
#   [R,G,B]: 最大位移时的颜色（渐变起点为白色）
#   例：xyz: [1, [1.0,1.0,1.0]] → 三方向合成位移 → 白色渐变
FRAME_COLORS = None
_DC_MODE = None
_DC_COLOR = None

_raw_dc = config.get("display_color", "")
if _raw_dc:
    try:
        dc = json.loads(_raw_dc) if isinstance(_raw_dc, str) else _raw_dc
        if isinstance(dc, dict):
            for mode, (enabled, color) in dc.items():
                if int(enabled) and mode in ("x","y","z","xy","yz","zx","xyz"):
                    _DC_MODE = mode
                    _DC_COLOR = np.array(color, dtype=np.float32)
                    break
    except Exception:
        pass

if _DC_MODE is not None and _DC_COLOR is not None:
    if EQ_POS is not None:
        eq_x = EQ_POS[None, :, 0]  # coord.txt 平衡位置
        eq_y = EQ_POS[None, :, 1]
        eq_z = EQ_POS[None, :, 2]
    else:
        eq_x = DISP_ALL_X[0:1, :]  # 第0帧
        eq_y = DISP_ALL_Y[0:1, :]
        eq_z = DISP_ALL_Z[0:1, :]
    dx = DISP_ALL_X - eq_x
    dy = DISP_ALL_Y - eq_y
    dz = DISP_ALL_Z - eq_z

    if _DC_MODE == "x":
        disp_mag = np.abs(dx)
    elif _DC_MODE == "y":
        disp_mag = np.abs(dy)
    elif _DC_MODE == "z":
        disp_mag = np.abs(dz)
    elif _DC_MODE == "xy":
        disp_mag = np.sqrt(dx**2 + dy**2)
    elif _DC_MODE == "yz":
        disp_mag = np.sqrt(dy**2 + dz**2)
    elif _DC_MODE == "zx":
        disp_mag = np.sqrt(dz**2 + dx**2)
    else:  # xyz
        disp_mag = np.sqrt(dx**2 + dy**2 + dz**2)

    # ── color_xrange: 限定归一化基准的原子范围 ───────────
    # 格式: [['min'|'mid'|'max'|数值, 'min'|'mid'|'max'|数值], ...] 对应 x,y,z
    # 在此范围内的原子计算最大位移 d_max，所有原子以此基准归一化着色
    _cxr = config.get("color_xrange", None)
    _d_max_source = disp_mag  # 默认：所有原子
    _range_label = "全部原子"
    if _cxr and isinstance(_cxr, (list, tuple)) and len(_cxr) == 3:
        try:
            # 获取坐标极值
            if EQ_POS is not None:
                _eq_all = EQ_POS
            else:
                _eq_all = np.column_stack([
                    DISP_ALL_X[0], DISP_ALL_Y[0], DISP_ALL_Z[0]])
            _cmin = _eq_all.min(axis=0)
            _cmax = _eq_all.max(axis=0)
            _cmid = (_cmin + _cmax) / 2

            _key_map = {"min": _cmin, "mid": _cmid, "max": _cmax}
            _range_lo = np.zeros(3, dtype=np.float64)
            _range_hi = np.zeros(3, dtype=np.float64)

            for _i in range(3):
                _lo = _cxr[_i][0]
                _hi = _cxr[_i][1]
                _range_lo[_i] = _key_map[_lo][_i] if _lo in _key_map else float(_lo)
                _range_hi[_i] = _key_map[_hi][_i] if _hi in _key_map else float(_hi)

            _in_x = (_eq_all[:, 0] >= _range_lo[0]) & (_eq_all[:, 0] <= _range_hi[0])
            _in_y = (_eq_all[:, 1] >= _range_lo[1]) & (_eq_all[:, 1] <= _range_hi[1])
            _in_z = (_eq_all[:, 2] >= _range_lo[2]) & (_eq_all[:, 2] <= _range_hi[2])
            _color_mask = _in_x & _in_y & _in_z

            _n_in_range = _color_mask.sum()
            if _n_in_range > 0:
                _d_max_source = disp_mag[:, _color_mask]  # 仅在范围内找最大位移
                _range_label = (f"x[{_range_lo[0]:.1f},{_range_hi[0]:.1f}] "
                                f"y[{_range_lo[1]:.1f},{_range_hi[1]:.1f}] "
                                f"z[{_range_lo[2]:.1f},{_range_hi[2]:.1f}]")
        except Exception as _e:
            print(f"[draw] color_xrange 解析失败: {_e}")

    # 用（范围限定的）最大位移归一化，所有原子统一着色
    d_max = _d_max_source.max()
    if d_max > 1e-12:
        disp_norm = disp_mag / d_max
    else:
        disp_norm = disp_mag

    FRAME_COLORS = np.ones((N_FRAMES, N_ATOMS, 4), dtype=np.float32)
    t = disp_norm[..., None]
    FRAME_COLORS[..., :3] = 1.0 + (_DC_COLOR - 1.0) * t
    print(f"[draw] 位移颜色映射: mode={_DC_MODE}, color={_DC_COLOR.tolist()}, "
          f"d_max={d_max:.4f}, 范围: {_range_label}")

    # ── color_driver: 驱动原子 → 统一颜色 ──
    _cd_raw = config.get("color_driver", None)
    if _cd_raw is not None:
        try:
            _cd = np.array(_cd_raw, dtype=np.float32)
            # 读取 driver.txt 获取驱动原子 ID
            _driver_file = config.get("driver_file", "driver.txt")
            _driver_path = os.path.join(input_dir, os.path.basename(_driver_file))
            if os.path.exists(_driver_path):
                _driver_ids = []
                with open(_driver_path, "r", encoding="utf-8") as _df:
                    _df.readline()  # skip header
                    for _line in _df:
                        _line = _line.strip()
                        if not _line or _line.startswith("#"):
                            continue
                        _parts = _line.split()
                        if _parts:
                            _driver_ids.append(int(_parts[0]) - 1)  # 1-based → 0-based
                if _driver_ids:
                    FRAME_COLORS[:, _driver_ids, :3] = _cd
                    print(f"[draw] color_driver: {len(_driver_ids)} 驱动原子 → RGB{_cd_raw}")
        except Exception as _e:
            print(f"[draw] color_driver 解析失败: {_e}")

    # ── color_fix: 全固定原子 (fix_x=fix_y=fix_z=1) → 统一颜色 ──
    _cf_raw = config.get("color_fix", None)
    if _cf_raw is not None and ATOM_FIXED is not None:
        try:
            _cf = np.array(_cf_raw, dtype=np.float32)
            _full_fixed = (ATOM_FIXED[:, 0] == 1) & (ATOM_FIXED[:, 1] == 1) & (ATOM_FIXED[:, 2] == 1)
            _n_fix = _full_fixed.sum()
            if _n_fix > 0:
                FRAME_COLORS[:, _full_fixed, :3] = _cf
                print(f"[draw] color_fix: {_n_fix} 全固定原子 → RGB{_cf_raw}")
        except Exception as _e:
            print(f"[draw] color_fix 解析失败: {_e}")
NSTEP     = int(config.get("NSTEP", 1))
DISP_STEP = np.arange(N_FRAMES) * NSTEP
DISP_T    = DISP_STEP * DT

# 原子信息
ATOM_IDS  = disp_data["atom_ids"]
# 优先使用 per-atom 半径，否则用统一的 ball_radius
_raw_radii = config.get("atom_radii", "")
if _raw_radii.strip():
    ATOM_RADII = np.array([float(x) for x in _raw_radii.split(",")])
else:
    ATOM_RADII = np.full(N_ATOMS, float(config.get("ball_radius", 0.5)))
PLOT_ATOM_ROW = 0
PLOT_ATOM_ID  = int(ATOM_IDS[0])
# 成键信息已在上面从 connection.txt 加载

# 渲染方式：0=Sphere(网格球体), 1=Marker(GPU点精灵)
USE_MARKER = int(config.get("use_marker", 0))

if N_FRAMES <= 0:
    raise ValueError(
        "output/display.txt 中没有可播放的帧，请检查 sample_start/sample_end/NSTEP 配置。")

# 模拟边界（从 input.txt 的 box_a 计算）
_box_a = float(config.get("box_a", 10.0))
X_MIN = -_box_a; X_MAX = _box_a
Y_MIN = -_box_a; Y_MAX = _box_a
Z_MIN = -_box_a; Z_MAX = _box_a
raw_alpha = config.get("alpha", "0.2")
if isinstance(raw_alpha, (list, tuple)):
    alpha_list = [float(x) for x in raw_alpha]
else:
    try:
        alpha_list = [float(x) for x in raw_alpha.split(",")]
    except (ValueError, AttributeError):
        alpha_list = [float(raw_alpha)] * 6
if len(alpha_list) != 6:
    alpha_list = (alpha_list * 6)[:6]

# 绘图参数
ball_radius = float(config.get("ball_radius", 0.5))
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))


# ===========================================================================
# 图形界面无关的几何参数（不参与物理计算，仅用于场景外观）
# ===========================================================================

info_margin  = 8
axis_length  = 10.0

import math as _math_cam

_cx = float(config.get("camera_center_x", 0.0))
_cy = float(config.get("camera_center_y", 0.0))
_cz = float(config.get("camera_center_z", 0.0))

# 若 input.txt 指定了摄像机自身坐标，则由坐标反推 distance/elevation/azimuth
if config.get("camera_pos_x") is not None:
    _px = float(config.get("camera_pos_x"))
    _py = float(config.get("camera_pos_y"))
    _pz = float(config.get("camera_pos_z"))
    _dx, _dy, _dz = _px - _cx, _py - _cy, _pz - _cz
    _dist = _math_cam.sqrt(_dx*_dx + _dy*_dy + _dz*_dz) or 1.0
    _elev = _math_cam.degrees(_math_cam.asin(max(-1.0, min(1.0, _dy / _dist))))
    _azim = _math_cam.degrees(_math_cam.atan2(_dx, _dz))
else:
    _dist = float(config.get("camera_distance", 40.0))
    _elev = float(config.get("camera_elevation", 0))
    _azim = float(config.get("camera_azimuth", 0))

initial_camera = {
    "distance":  _dist,
    "elevation": _elev,
    "azimuth":   _azim,
    "center":    (_cx, _cy, _cz),
}


# ===========================================================================
# 创建画布与相机
# ===========================================================================
canvas = scene.SceneCanvas(
    keys="interactive",
    size=(1000, 700),
    bgcolor=(0.08, 0.08, 0.10, 1.0),
    show=True,
)

view = canvas.central_widget.add_view()
view.camera = "turntable"
view.camera.distance  = initial_camera["distance"]
view.camera.elevation = initial_camera["elevation"]
view.camera.azimuth   = initial_camera["azimuth"]
view.camera.center    = initial_camera["center"]


# ===========================================================================
# 场景对象
# ===========================================================================
axis_width = 3
arrow_size = 14
axes_visible = True
axes_group = []

axes_group.append(scene.visuals.Arrow(
    pos=np.array([[0, 0, 0], [axis_length, 0, 0]], dtype=np.float32),
    color=(1.0, 0.2, 0.2, 1.0),
    width=axis_width,
    arrows=np.array([[0, 0, 0, axis_length, 0, 0]], dtype=np.float32),
    arrow_size=arrow_size,
    parent=view.scene,
))
axes_group.append(scene.visuals.Arrow(
    pos=np.array([[0, 0, 0], [0, axis_length, 0]], dtype=np.float32),
    color=(0.2, 1.0, 0.2, 1.0),
    width=axis_width,
    arrows=np.array([[0, 0, 0, 0, axis_length, 0]], dtype=np.float32),
    arrow_size=arrow_size,
    parent=view.scene,
))
axes_group.append(scene.visuals.Arrow(
    pos=np.array([[0, 0, 0], [0, 0, axis_length]], dtype=np.float32),
    color=(0.3, 0.6, 1.0, 1.0),
    width=axis_width,
    arrows=np.array([[0, 0, 0, 0, 0, axis_length]], dtype=np.float32),
    arrow_size=arrow_size,
    parent=view.scene,
))

axes_group.append(scene.visuals.Text(text="x", color=(1.0, 0.2, 0.2, 1.0), font_size=14,
    pos=(axis_length + 0.2, 0, 0), anchor_x="left", anchor_y="center", parent=view.scene))
axes_group.append(scene.visuals.Text(text="y", color=(0.2, 1.0, 0.2, 1.0), font_size=14,
    pos=(0, axis_length + 0.2, 0), anchor_x="left", anchor_y="bottom", parent=view.scene))
axes_group.append(scene.visuals.Text(text="z", color=(0.3, 0.6, 1.0, 1.0), font_size=14,
    pos=(0, 0, axis_length + 0.2), anchor_x="left", anchor_y="bottom", parent=view.scene))

# ── 原子渲染 ──────────────────────────────────
# 两种渲染模式：
#   Sphere = 独立网格球体（精细，适合少量原子）
#   Marker = GPU 实例化点精灵（高效，适合大量原子）
TAB10_RGB = np.array([
    [0.1216, 0.4667, 0.7059],   # 蓝
    [0.8902, 0.4667, 0.1137],   # 橙
    [0.1725, 0.6275, 0.1725],   # 绿
    [0.8392, 0.1529, 0.1569],   # 红
    [0.5804, 0.4039, 0.7412],   # 紫
    [0.5490, 0.3373, 0.2941],   # 棕
    [0.8902, 0.4667, 0.7608],   # 粉
    [0.4980, 0.4980, 0.4980],   # 灰
    [0.7373, 0.7412, 0.1333],   # 黄绿
    [0.0902, 0.7451, 0.8118],   # 青
])
# 每个原子的颜色（循环使用 tab10 色板，或按位移着色）
atom_colors = np.zeros((N_ATOMS, 4), dtype=np.float32)
if FRAME_COLORS is not None:
    atom_colors[:] = FRAME_COLORS[0]   # 初始帧颜色
else:
    for i in range(N_ATOMS):
        r, g, b = TAB10_RGB[i % len(TAB10_RGB)]
        atom_colors[i] = [r, g, b, 1.0]

if USE_MARKER:
    # ── Marker 模式：GPU 实例化，一次 draw call ──
    marker_pos = np.zeros((N_ATOMS, 3), dtype=np.float32)
    marker_sizes = np.array([float(ATOM_RADII[i]) * 40 for i in range(N_ATOMS)], dtype=np.float32)
    balls = scene.visuals.Markers(parent=view.scene)
    balls.set_data(
        pos=marker_pos, face_color=atom_colors,
        size=marker_sizes, symbol='disc', edge_width=0,
    )
else:
    # ── Sphere 模式：每个原子一个独立网格球体 ──
    balls = []
    for i in range(N_ATOMS):
        r, g, b, _ = atom_colors[i]
        s = scene.visuals.Sphere(
            radius=float(ATOM_RADII[i]), method="latitude",
            color=(r, g, b, 1.0), edge_color=None, parent=view.scene)
        s.mesh.shading = "smooth"
        balls.append(s)

# 成键线（原子之间的连接）
if len(BOND_PAIRS) > 0:
    n_bonds = len(BOND_PAIRS)
    bond_pos_buffer = np.zeros((n_bonds * 2, 3), dtype=np.float32)  # 预分配，后续原地修改
    bond_lines = scene.visuals.Line(
        pos=bond_pos_buffer, color=(0.7, 0.7, 0.7, 0.8), width=3,
        connect="segments", parent=view.scene)
else:
    bond_pos_buffer = None
    bond_lines = None

# 六个面形成立方体边界（每个面独立透明度，alpha<=0 时隐藏该面）
box_size = X_MAX - X_MIN
faces = [
    ((X_MAX, 0, 0), "-x"),
    ((X_MIN, 0, 0), "+x"),
    ((0, Y_MAX, 0), "-y"),
    ((0, Y_MIN, 0), "+y"),
    ((0, 0, Z_MAX), "-z"),
    ((0, 0, Z_MIN), "+z"),
]
for f_idx, (pos, direction) in enumerate(faces):
    a = alpha_list[f_idx]
    if a <= 0:
        continue
    face_color = (box_color_r, box_color_g, box_color_b, a)
    plane = scene.visuals.Plane(
        width=box_size, height=box_size, width_segments=1, height_segments=1,
        direction=direction, color=face_color, parent=view.scene)
    plane.set_gl_state(depth_test=False, blend=True)
    plane.transform = STTransform(translate=pos)

# 右上角：相机信息
camera_info = scene.visuals.Text(
    text="", color="white", font_size=12,
    pos=(0, 0), anchor_x="right", anchor_y="top", parent=canvas.scene)

# 左上角：小球信息
ball_info = scene.visuals.Text(
    text="", color=(0.2, 1.0, 0.2, 1.0), font_size=14,
    pos=(0, 0), anchor_x="left", anchor_y="top",
    face="黑体", bold=True, parent=canvas.scene)

# 左上角：reset 按钮（在 info 上方）
reset_btn_size = (60, 30)
reset_button = scene.visuals.Rectangle(
    center=(reset_btn_size[0] / 2 + 8, reset_btn_size[1] / 2 + 8),
    width=reset_btn_size[0], height=reset_btn_size[1],
    radius=6, color=(0.18, 0.35, 0.65, 0.85),
    border_color="white", parent=canvas.scene)
reset_button_label = scene.visuals.Text(
    text="reset", color="white", font_size=13,
    pos=(reset_btn_size[0] / 2 + 8, reset_btn_size[1] / 2 + 8),
    anchor_x="center", anchor_y="center",
    bold=True, parent=canvas.scene)

# 信息显示/隐藏 切换按钮（左上角小方块，在 reset 下方）
info_btn_size = (60, 30)
info_toggle_visible = True
info_button = scene.visuals.Rectangle(
    center=(info_btn_size[0] / 2 + 8, info_btn_size[1] / 2 + 8),
    width=info_btn_size[0], height=info_btn_size[1],
    radius=6, color=(0.9, 0.3, 0.3, 0.9),
    border_color="white", parent=canvas.scene)
info_button_label = scene.visuals.Text(
    text="info", color="white", font_size=13,
    pos=(info_btn_size[0] / 2 + 8, info_btn_size[1] / 2 + 8),
    anchor_x="center", anchor_y="center",
    bold=True, parent=canvas.scene)

# 强制在初始化时定位到左上角（canvas.scene坐标系：左下角为原点）
_cw, _ch = canvas.size
reset_button.center = (reset_btn_size[0] / 2 + 8, _ch - reset_btn_size[1] / 2 - 8 - info_btn_size[1] - 4)
reset_button_label.pos = reset_button.center
info_button.center = (info_btn_size[0] / 2 + 8, _ch - info_btn_size[1] / 2 - 8)
info_button_label.pos = info_button.center

# axes 显示/隐藏 按钮（在 info 下方）
axes_btn_size = (60, 30)
axes_button = scene.visuals.Rectangle(
    center=(axes_btn_size[0] / 2 + 8, axes_btn_size[1] / 2 + 8),
    width=axes_btn_size[0], height=axes_btn_size[1],
    radius=6, color=(0.3, 0.7, 0.3, 0.9),
    border_color="white", parent=canvas.scene)
axes_button_label = scene.visuals.Text(
    text="axes", color="white", font_size=13,
    pos=(axes_btn_size[0] / 2 + 8, axes_btn_size[1] / 2 + 8),
    anchor_x="center", anchor_y="center",
    bold=True, parent=canvas.scene)
axes_button.center = (axes_btn_size[0] / 2 + 8, _ch - axes_btn_size[1] / 2 - 8 - info_btn_size[1] - 4 - axes_btn_size[1] - 4)
axes_button_label.pos = axes_button.center


# ===========================================================================
# 回调函数
# ===========================================================================

def update_camera_info(event=None):
    c = view.camera
    camera_info.text = (
        "Camera\n"
        f"center = ({c.center[0]:.1f}, {c.center[1]:.1f}, {c.center[2]:.1f})\n"
        f"distance = {c.distance:.1f}\n"
        f"elevation = {c.elevation:.1f}\n"
        f"azimuth = {c.azimuth:.1f}\n"
        f"step = {PAN_SPEED:.1f}  |  {'透视' if _PERSPECTIVE else '正交'}"
    )


def update_ball_info(frame_idx, x, y, z, vx, vy, vz):
    step = int(DISP_STEP[frame_idx])
    t    = float(DISP_T[frame_idx])
    ball_info.text = (
        f"原子 {PLOT_ATOM_ID}  (共 {N_ATOMS} 个原子)\n"
        f"frame {frame_idx+1}/{N_FRAMES}  |  saved step {step}/{NT-1}\n"
        f"t = {t:.2f} s  |  dt = {DT:.3f} s  |  nstep = {NSTEP}\n"
        f"Position: ({x:.2f}, {y:.2f}, {z:.2f})\n"
        f"Velocity: ({vx:.2f}, {vy:.2f}, {vz:.2f})\n"
        f"W/S 沿Z轴  |  A/D 左右  |  Q/E 升降  |  C/X 步长  |  V 透视/正交"
    )


def reposition_camera_info(event=None):
    width, height = canvas.size
    camera_info.pos  = (width - 20, height - 20)
    # reset → info → axes 三按钮纵向排列
    gap = 4
    info_y = info_btn_size[1] / 2 + 8
    reset_y = info_y + info_btn_size[1] + gap
    axes_y = info_y + info_btn_size[1] + gap + axes_btn_size[1] + gap
    reset_button.center = (reset_btn_size[0] / 2 + 8, height - reset_y)
    reset_button_label.pos = reset_button.center
    info_button.center = (info_btn_size[0] / 2 + 8, height - info_y)
    info_button_label.pos = info_button.center
    axes_button.center = (axes_btn_size[0] / 2 + 8, height - axes_y)
    axes_button_label.pos = axes_button.center
    # info 文字放在所有按钮下方
    buttons_bottom = height - (axes_y + axes_btn_size[1] / 2)
    ball_info.pos = (info_margin, buttons_bottom - 10)
    update_ball_info(frame_idx, DISP_X[frame_idx], DISP_Y[frame_idx], DISP_Z[frame_idx],
                     DISP_VX[frame_idx], DISP_VY[frame_idx], DISP_VZ[frame_idx])
    update_camera_info()


# ── 平移速度 & 投影模式（全局变量）──
PAN_SPEED = 1.0
_PERSPECTIVE = True          # True=透视, False=正交


def handle_view_interaction(event):
    update_camera_info()


def handle_key_press(event):
    global PAN_SPEED, _PERSPECTIVE
    key_name = ""
    if getattr(event, "text", None):
        key_name = event.text.lower()
    elif getattr(event, "key", None) is not None:
        key_name = str(event.key).lower()

    c = view.camera

    # ── 投影切换 ──
    if key_name == "v":
        _PERSPECTIVE = not _PERSPECTIVE
        try:
            if _PERSPECTIVE:
                c.fov = 60.0
            else:
                c.fov = 0.0
            print(f"[draw] 投影模式: {'透视' if _PERSPECTIVE else '正交'}")
        except Exception:
            print("[draw] 当前相机不支持 fov 切换")
        update_camera_info()
        return

    # ── 步长控制 ──
    if key_name == "c":
        PAN_SPEED = min(PAN_SPEED * 1.5, 50.0)
        print(f"[draw] 步长: {PAN_SPEED:.1f}")
        update_camera_info()
        return
    elif key_name == "x":
        PAN_SPEED = max(PAN_SPEED / 1.5, 0.05)
        print(f"[draw] 步长: {PAN_SPEED:.1f}")
        update_camera_info()
        return

    # ── 计算相机方向向量 ──
    import math as _math
    azim_rad = _math.radians(c.azimuth)
    elev_rad = _math.radians(c.elevation)

    # 视线方向（从 center 指向相机）
    vd = np.array([
        _math.cos(elev_rad) * _math.sin(azim_rad),
        _math.sin(elev_rad),
        _math.cos(elev_rad) * _math.cos(azim_rad),
    ])
    vd /= np.linalg.norm(vd)

    # 屏幕右方向
    world_up = np.array([0.0, 1.0, 0.0])
    right = np.cross(vd, world_up)
    rn = np.linalg.norm(right)
    if rn > 1e-10:
        right /= rn
    else:
        right = np.array([1.0, 0.0, 0.0])

    # 屏幕上方向
    up = np.cross(right, vd)
    un = np.linalg.norm(up)
    if un > 1e-10:
        up /= un
    else:
        up = np.array([0.0, 1.0, 0.0])

    pan = PAN_SPEED * 0.3

    if key_name == "a":           # 右移（屏幕右方向）
        c.center = tuple(np.array(c.center) + right * pan)
    elif key_name == "d":         # 左移（屏幕右方向负向）
        c.center = tuple(np.array(c.center) - right * pan)
    elif key_name == "e":         # 上升（屏幕上方向，原 W 的功能）
        c.center = tuple(np.array(c.center) + up * pan)
    elif key_name == "q":         # 下降（屏幕上方向负向，原 S 的功能）
        c.center = tuple(np.array(c.center) - up * pan)
    elif key_name == "w":         # 相机沿 Z 轴上移（靠近场景）
        c.center = (c.center[0], c.center[1], c.center[2] + pan)
    elif key_name == "s":         # 相机沿 Z 轴下移（远离场景）
        c.center = (c.center[0], c.center[1], c.center[2] - pan)
    else:
        return

    update_camera_info()


def reset_camera_view():
    global frame_idx
    frame_idx = 0
    # 立即复位所有原子到第 0 帧
    _update_atom_positions(frame_idx)
    if bond_lines is not None and len(BOND_PAIRS) > 0:
        _update_bond_positions(frame_idx)
    # 复位相机
    view.camera.distance  = initial_camera["distance"]
    view.camera.elevation  = initial_camera["elevation"]
    view.camera.azimuth    = initial_camera["azimuth"]
    view.camera.center     = initial_camera["center"]
    if hasattr(view.camera, "roll"):
        view.camera.roll = 0
    update_camera_info()


def handle_mouse_press(event):
    global info_toggle_visible, axes_visible
    if event.pos is None:
        return
    ex, ey = event.pos

    # reset 按钮
    bx, by = reset_button.center
    lw = reset_btn_size[0] / 2;  lh = reset_btn_size[1] / 2
    if (bx - lw) <= ex <= (bx + lw) and (by - lh) <= ey <= (by + lh):
        reset_camera_view()
        return

    # info 按钮
    bx, by = info_button.center
    if (bx - info_btn_size[0]/2) <= ex <= (bx + info_btn_size[0]/2) and (by - info_btn_size[1]/2) <= ey <= (by + info_btn_size[1]/2):
        info_toggle_visible = not info_toggle_visible
        ball_info.visible = info_toggle_visible
        camera_info.visible = info_toggle_visible
        return

    # axes 按钮
    bx, by = axes_button.center
    if (bx - axes_btn_size[0]/2) <= ex <= (bx + axes_btn_size[0]/2) and (by - axes_btn_size[1]/2) <= ey <= (by + axes_btn_size[1]/2):
        axes_visible = not axes_visible
        for axis in axes_group:
            axis.visible = axes_visible


# ===========================================================================
# 位置/键线更新辅助函数（两种渲染模式共用）
# ===========================================================================

def _update_atom_positions(f_idx):
    """更新所有原子到第 f_idx 帧的位置，必要时更新颜色。"""
    if USE_MARKER:
        marker_pos[:, 0] = DISP_ALL_X[f_idx]
        marker_pos[:, 1] = DISP_ALL_Y[f_idx]
        marker_pos[:, 2] = DISP_ALL_Z[f_idx]
        if FRAME_COLORS is not None:
            balls.set_data(pos=marker_pos, face_color=FRAME_COLORS[f_idx])
        else:
            balls.set_data(pos=marker_pos)
    else:
        for i in range(N_ATOMS):
            balls[i].transform = STTransform(translate=(
                float(DISP_ALL_X[f_idx, i]),
                float(DISP_ALL_Y[f_idx, i]),
                float(DISP_ALL_Z[f_idx, i]),
            ))


def _update_bond_positions(f_idx):
    """原地更新键线顶点位置，避免重新分配内存。"""
    for b_idx, (i, j) in enumerate(BOND_PAIRS):
        b2 = b_idx * 2
        bond_pos_buffer[b2]     = [DISP_ALL_X[f_idx, i], DISP_ALL_Y[f_idx, i], DISP_ALL_Z[f_idx, i]]
        bond_pos_buffer[b2 + 1] = [DISP_ALL_X[f_idx, j], DISP_ALL_Y[f_idx, j], DISP_ALL_Z[f_idx, j]]
    bond_lines.set_data(pos=bond_pos_buffer)


# ===========================================================================
# 动画初始化
# ===========================================================================
frame_idx = 0

# 初始帧：摆放所有原子并刷新 UI
_update_atom_positions(frame_idx)
if bond_lines is not None and len(BOND_PAIRS) > 0:
    _update_bond_positions(frame_idx)
reposition_camera_info()
update_ball_info(frame_idx,
    float(DISP_X[frame_idx]), float(DISP_Y[frame_idx]), float(DISP_Z[frame_idx]),
    float(DISP_VX[frame_idx]), float(DISP_VY[frame_idx]), float(DISP_VZ[frame_idx]))

mode_str = "Marker (GPU 实例化)" if USE_MARKER else "Sphere (独立网格)"
print(f"[draw] 加载 output/display.txt: {N_FRAMES} 帧, {N_ATOMS} 个原子, NT={NT}, DT={DT}, NSTEP={NSTEP}")
print(f"[draw] 渲染方式: {mode_str}")
print(f"[draw] 绘图参数: ball_radius={ball_radius}, box_color=({box_color_r:.2f},{box_color_g:.2f},{box_color_b:.2f}), alpha={alpha_list}")


# 运动相机（速度段驱动）：优先读取 move_camera.txt，其次用 display.txt header 缓存
import re as _re

def _load_move_camera_txt():
    """直接读取 input/move_camera.txt（与 output 同级的 input 目录）。"""
    input_dir = os.path.join(os.path.dirname(output_dir), "input")
    cam_path = os.path.join(input_dir, "move_camera.txt")
    if not os.path.exists(cam_path):
        return None
    segs = []
    with open(cam_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, r = [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]
            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))
                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):
                segs.append({"start": start, "end": end, "v": v, "r": r})
    return segs if segs else None

# 先试 move_camera.txt 直读，没有则用 display.txt 缓存
# header 中 camera_keyframes 为空字符串表示 move_camera=0（开关关闭），跳过文件加载
_camera_motion_enabled = bool(config.get("camera_keyframes", ""))
_CAM_MOTION = _load_move_camera_txt() if _camera_motion_enabled else None
if not _CAM_MOTION:
    _CAM_MOTION = json.loads(config.get("camera_keyframes", "null")) if config.get("camera_keyframes") else None
if _CAM_MOTION:
    _cam_center = [0.0, 0.0, 0.0]
    _cam_elev = initial_camera["elevation"]
    _cam_azim = initial_camera["azimuth"]
    _cam_dist = initial_camera["distance"]
    src = "move_camera.txt" if _load_move_camera_txt() else "display.txt header"
    print(f"[draw] 运动相机已启用（数据来源: {src}，{len(_CAM_MOTION)} 段）")


def _update_motion_camera(f_idx):
    """速度段驱动：每帧累加平移/旋转。

    时间交叠时所有段同时生效，按文件中出现的顺序依次作用。
    只有当前帧存在活动段时才覆写相机，否则保留用户键盘/鼠标操作的结果。
    """
    if not _CAM_MOTION:
        return
    global _cam_center, _cam_elev, _cam_azim, _cam_dist
    active = False
    for seg in _CAM_MOTION:
        if seg["start"] <= f_idx < seg["end"]:
            _cam_center[0] += seg["v"][0]
            _cam_center[1] += seg["v"][1]
            _cam_center[2] += seg["v"][2]
            _cam_elev += seg["r"][0]
            _cam_azim += seg["r"][1]
            active = True

    if active:
        view.camera.center    = tuple(_cam_center)
        view.camera.distance  = _cam_dist
        view.camera.elevation = _cam_elev
        view.camera.azimuth   = _cam_azim


def update(event):
    global frame_idx
    frame_idx = (frame_idx + 1) % N_FRAMES     # 循环播放

    # 更新所有原子位置（两种模式共用逻辑）
    _update_atom_positions(frame_idx)

    # 更新成键线（原地修改，无内存分配）
    if bond_lines is not None and len(BOND_PAIRS) > 0:
        _update_bond_positions(frame_idx)

    # 运动相机：速度段驱动
    _update_motion_camera(frame_idx)

    # 信息面板显示 plot_atom 的数据
    x = float(DISP_X[frame_idx])
    y = float(DISP_Y[frame_idx])
    z = float(DISP_Z[frame_idx])
    vx = float(DISP_VX[frame_idx])
    vy = float(DISP_VY[frame_idx])
    vz = float(DISP_VZ[frame_idx])
    update_ball_info(frame_idx, x, y, z, vx, vy, vz)
    update_camera_info()


# ===========================================================================
# 事件绑定与启动
# ===========================================================================
timer = app.Timer(interval=0.02, connect=update, start=True)
canvas.events.mouse_move.connect(handle_view_interaction)
canvas.events.mouse_press.connect(handle_mouse_press)
canvas.events.mouse_wheel.connect(handle_view_interaction)
canvas.events.resize.connect(reposition_camera_info)
canvas.events.key_press.connect(handle_key_press)

if hasattr(canvas, "native") and hasattr(canvas.native, "setFocus"):
    canvas.native.setFocus()


if __name__ == "__main__":
    app.run()
