"""
sample.py
---------
纯 Python 抽帧脚本，不依赖 VisPy。

功能：
  1. 从 output/trajectory.txt 读取完整轨迹
  2. 每隔 NSTEP 帧抽取一帧，生成 output/display.txt
  3. output/display.txt 可直接被 draw.py 加载驱动动画
"""

import numpy as np
import os
import sys

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


def build_sample_indices(total_steps, sample_step, sample_start, sample_end):
    """Validate sampling settings and build frame indices."""
    if sample_step <= 0:
        raise ValueError(f"NSTEP 必须为正整数，实际为 {sample_step}")
    if sample_start < 0:
        raise ValueError(f"sample_start 不能小于 0，实际为 {sample_start}")
    if sample_end > total_steps:
        raise ValueError(
            f"sample_end 不能大于记录步数 {total_steps}，实际为 {sample_end}")
    if sample_start >= sample_end:
        raise ValueError(
            f"sample_start 必须小于 sample_end，实际为 [{sample_start}, {sample_end})")

    n_frames = (sample_end - sample_start) // sample_step
    if n_frames <= 0:
        raise ValueError(
            f"抽帧范围 [{sample_start}, {sample_end}) 过短，按 NSTEP={sample_step} 无法抽出任何帧")

    return np.arange(n_frames, dtype=np.int64) * sample_step + sample_start


def read_optional_index(data, key, default_value):
    """Read an optional integer index from txt metadata."""
    if key not in data:
        return default_value
    value = data[key]
    if value is None or int(value) < 0:
        return default_value
    return int(value)


def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    output_dir = compute.get_output_dir(script_dir)
    traj_path  = os.path.join(output_dir, "trajectory.txt")
    disp_path  = os.path.join(output_dir, "display.txt")

    # -------------------------------------------------------------------------
    # 1. 加载完整轨迹
    # -------------------------------------------------------------------------
    data = compute.load_text_data(traj_path)
    NT       = int(data["NT"])
    DT       = float(data["DT"])
    NSTEP    = int(data["NSTEP"])
    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"]
    plot_atom_row = int(data["plot_atom_row"]) if "plot_atom_row" in data else 0
    plot_atom_id = int(data["plot_atom_id"]) if "plot_atom_id" in data else int(data["atom_ids"][plot_atom_row])

    print(f"[sample] 加载轨迹数据: NT={NT}, DT={DT}, NSTEP={NSTEP}")

    # -------------------------------------------------------------------------
    # 2. 抽帧：支持配置文件中保存的 [sample_start, sample_end) 区间
    # -------------------------------------------------------------------------
    sample_start = read_optional_index(data, "sample_start", 0)
    sample_end = read_optional_index(data, "sample_end", NT)
    indices = build_sample_indices(NT, NSTEP, sample_start, sample_end)
    n_frames = len(indices)

    if traj_x.ndim == 1:
        selected_x = traj_x
        selected_y = traj_y
        selected_z = traj_z
        selected_vx = traj_vx
        selected_vy = traj_vy
        selected_vz = traj_vz
        all_x = traj_x[:, None]
        all_y = traj_y[:, None]
        all_z = traj_z[:, None]
        all_vx = traj_vx[:, None]
        all_vy = traj_vy[:, None]
        all_vz = traj_vz[:, None]
    else:
        selected_x = traj_x[:, plot_atom_row]
        selected_y = traj_y[:, plot_atom_row]
        selected_z = traj_z[:, plot_atom_row]
        selected_vx = traj_vx[:, plot_atom_row]
        selected_vy = traj_vy[:, plot_atom_row]
        selected_vz = traj_vz[:, plot_atom_row]
        all_x = traj_x
        all_y = traj_y
        all_z = traj_z
        all_vx = traj_vx
        all_vy = traj_vy
        all_vz = traj_vz

    disp_x   = selected_x[indices]
    disp_y   = selected_y[indices]
    disp_z   = selected_z[indices]
    disp_vx  = selected_vx[indices]
    disp_vy  = selected_vy[indices]
    disp_vz  = selected_vz[indices]
    disp_t   = indices * DT   # 物理时间 = step * DT
    disp_step = indices       # 对应保存轨迹的步编号（0-based）

    print(f"[sample] 抽帧完成: [{sample_start}, {sample_end}) -> {n_frames} 帧")
    print(f"[sample]   每帧时间跨度: {NSTEP*DT:.3f} s  (即每隔 {NSTEP} 步取一帧)")

    # -------------------------------------------------------------------------
    # 3. 保存显示数组
    # -------------------------------------------------------------------------
    payload = {
        "disp_x": disp_x,
        "disp_y": disp_y,
        "disp_z": disp_z,
        "disp_vx": disp_vx,
        "disp_vy": disp_vy,
        "disp_vz": disp_vz,
        "disp_all_x": all_x[indices],
        "disp_all_y": all_y[indices],
        "disp_all_z": all_z[indices],
        "disp_all_vx": all_vx[indices],
        "disp_all_vy": all_vy[indices],
        "disp_all_vz": all_vz[indices],
        "disp_t": disp_t,
        "disp_step": disp_step,
        "n_frames": n_frames,
        "NT": NT,
        "DT": DT,
        "NSTEP": NSTEP,
        "plot_atom_id": plot_atom_id,
        "plot_atom_row": plot_atom_row,
        "method": str(data["method"]) if "method" in data else "explicit_euler",
        "coord_file": str(data["coord_file"]) if "coord_file" in data else os.path.join("input", "coord.txt"),
        "atom_ids": data["atom_ids"] if "atom_ids" in data else np.array([1]),
        "atom_masses": data["atom_masses"] if "atom_masses" in data else np.array([float(data["M"])]),
        "atom_radii": data["atom_radii"] if "atom_radii" in data else np.array([float(data["ball_radius"])]),
        "atom_positions": data["atom_positions"] if "atom_positions" in data else np.array([[float(data["X0"]), float(data["Y0"]), float(data["Z0"])]]),
        "atom_velocities": data["atom_velocities"] if "atom_velocities" in data else np.array([[float(data["VX0"]), float(data["VY0"]), float(data["VZ0"])]]),
        "atom_fixed": data["atom_fixed"] if "atom_fixed" in data else np.array([[0, 0, 0]]),
        "warmup_steps": int(data["warmup_steps"]) if "warmup_steps" in data else 0,
        "sample_start": sample_start,
        "sample_end": sample_end,
        "X_MIN": float(data["X_MIN"]),
        "X_MAX": float(data["X_MAX"]),
        "Y_MIN": float(data["Y_MIN"]),
        "Y_MAX": float(data["Y_MAX"]),
        "Z_MIN": float(data["Z_MIN"]),
        "Z_MAX": float(data["Z_MAX"]),
        "X0": float(data["X0"]),
        "Y0": float(data["Y0"]),
        "Z0": float(data["Z0"]),
        "VX0": float(data["VX0"]),
        "VY0": float(data["VY0"]),
        "VZ0": float(data["VZ0"]),
        "M": float(data["M"]) if "M" in data else 1.0,
        "alpha": float(data["alpha"]),
        "ball_radius": float(data["ball_radius"]),
        "ball_color_r": float(data["ball_color_r"]),
        "ball_color_g": float(data["ball_color_g"]),
        "ball_color_b": float(data["ball_color_b"]),
        "box_color_r": float(data["box_color_r"]),
        "box_color_g": float(data["box_color_g"]),
        "box_color_b": float(data["box_color_b"]),
    }
    compute.save_text_data(disp_path, payload)
    print(f"[sample] 显示数组已保存至: {disp_path}")


if __name__ == "__main__":
    main()
