"""
plot_trajectory.py
------------------
绘制轨迹数据的x-t, y-t, z-t, vx-t, vy-t, vz-t函数图像，支持多原子同时绘制。

用法：
    python plot_trajectory.py                              # 绘制所有原子
    python plot_trajectory.py output/trajectory.txt        # 指定文件
    python plot_trajectory.py output/trajectory.txt --atom-id 1  # 仅绘制原子1

参数：
    trajectory_file: 轨迹数据文件，支持结构化 .txt 格式（默认：output/trajectory.txt）
    --atom-id: 可选，指定只绘制单个原子的轨迹
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

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

def select_atom_series(data, atom_id=None):
    """Select one atom from trajectory arrays when the file stores many atoms."""
    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']

    if traj_x.ndim == 1:
        selected_id = int(data["plot_atom_id"]) if "plot_atom_id" in data else 1
        return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, selected_id

    atom_ids = data["atom_ids"] if "atom_ids" in data else np.arange(traj_x.shape[1]) + 1
    if atom_id is None:
        row = int(data["plot_atom_row"]) if "plot_atom_row" in data else 0
    else:
        matches = np.where(atom_ids == int(atom_id))[0]
        if len(matches) == 0:
            raise ValueError(f"原子序号 {atom_id} 不在轨迹文件中")
        row = int(matches[0])
    selected_id = int(atom_ids[row])
    return (
        traj_x[:, row], traj_y[:, row], traj_z[:, row],
        traj_vx[:, row], traj_vy[:, row], traj_vz[:, row],
        selected_id,
    )


def load_trajectory_txt(file_path, atom_id=None):
    """加载结构化 .txt 格式的轨迹数据"""
    data = compute.load_text_data(file_path)
    traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, selected_id = select_atom_series(data, atom_id)
    NT = int(data['NT'])
    DT = float(data['DT'])
    return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, selected_id


def load_full_trajectory(file_path):
    """加载完整轨迹数据（所有原子），返回 2D 数组及原始数据字典。"""
    data = compute.load_text_data(file_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']
    NT = int(data['NT'])
    DT = float(data['DT'])
    n_atoms = traj_x.shape[1] if traj_x.ndim > 1 else 1
    atom_ids = data.get("atom_ids", np.arange(n_atoms) + 1)
    return traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, atom_ids, data

def create_plots(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, output_dir='.', atom_ids=None, extra_data=None):
    """创建六个子图的图表，支持多原子绘制与能量曲线。"""
    # 配置中文字体支持
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False

    # 生成时间数组
    time = np.arange(NT) * DT

    # 判断单原子（1D）或多原子（2D）
    is_multi = traj_x.ndim == 2
    n_atoms = traj_x.shape[1] if is_multi else 1

    # 是否绘制能量子图（需要 multi-atom + extra_data 中有 G）
    has_energy = is_multi and extra_data is not None and "G" in extra_data
    n_rows = 3 if has_energy else 2
    n_cols = 3
    figsize = (15, 13) if has_energy else (15, 9)

    # 创建图形和子图
    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    title = '轨迹与能量分析' if has_energy else '轨迹分析'
    if atom_ids is not None and not is_multi:
        title += f' - 原子 {int(atom_ids)}'
    fig.suptitle(title, fontsize=16)

    plot_configs = [
        (axes[0, 0], traj_x, 'x'),
        (axes[0, 1], traj_y, 'y'),
        (axes[0, 2], traj_z, 'z'),
        (axes[1, 0], traj_vx, 'vx'),
        (axes[1, 1], traj_vy, 'vy'),
        (axes[1, 2], traj_vz, 'vz'),
    ]

    if is_multi:
        colors = plt.cm.tab10(np.linspace(0, 1, n_atoms))
        if atom_ids is None:
            atom_ids = np.arange(n_atoms) + 1

    for ax, data_arr, label in plot_configs:
        if is_multi:
            for i in range(n_atoms):
                aid = int(atom_ids[i])
                ax.plot(time, data_arr[:, i], color=colors[i], linewidth=1.5, label=f"原子 {aid}")
            ax.legend()
        else:
            ax.plot(time, data_arr, 'b-', linewidth=1.5)
        ax.set_title(f'{label} - 时间')
        ax.set_xlabel('时间 (s)')
        ax.set_ylabel(label)
        ax.grid(True, alpha=0.3)

    # ── 能量子图 ─────────────────────────────────────
    if has_energy:
        masses = np.array(extra_data["atom_masses"])                    # (n_atoms,)
        G_vec = np.array(extra_data["G"])                               # [gx, gy, gz]

        # 动能 Ek = ½ m v²
        ek = 0.5 * masses[np.newaxis, :] * (traj_vx**2 + traj_vy**2 + traj_vz**2)
        # 重力势能 Ug = -m G·r
        ug = -masses[np.newaxis, :] * (G_vec[0] * traj_x + G_vec[1] * traj_y + G_vec[2] * traj_z)
        # 弹性势能 Us = ½ k (d - d₀)²
        us = np.zeros_like(ek)
        bond_pairs = extra_data.get("bond_pairs")
        bond_stiffness = extra_data.get("bond_stiffness")
        bond_rest_lengths = extra_data.get("bond_rest_lengths")
        if bond_pairs is not None and len(bond_pairs) > 0:
            for b_idx in range(len(bond_pairs)):
                i, j = bond_pairs[b_idx]
                dx = traj_x[:, j] - traj_x[:, i]
                dy = traj_y[:, j] - traj_y[:, i]
                dz = traj_z[:, j] - traj_z[:, i]
                dist = np.sqrt(dx**2 + dy**2 + dz**2)
                stretch = dist - bond_rest_lengths[b_idx]
                us[:, i] += 0.5 * bond_stiffness[b_idx] * stretch**2
                us[:, j] += 0.5 * bond_stiffness[b_idx] * stretch**2

        e_total = ek + ug + us                                        # (NT, n_atoms)
        ek_sys = np.sum(ek, axis=1)
        ug_sys = np.sum(ug, axis=1)
        us_sys = np.sum(us, axis=1)
        e_sys = ek_sys + ug_sys + us_sys

        # 第 3 行左：各原子总能量
        ax_e = axes[2, 0]
        for i in range(n_atoms):
            aid = int(atom_ids[i])
            ax_e.plot(time, e_total[:, i], color=colors[i], linewidth=1.5, label=f"原子 {aid}")
        ax_e.set_title("各原子总能量")
        ax_e.set_xlabel("时间 (s)")
        ax_e.set_ylabel("能量")
        ax_e.grid(True, alpha=0.3)
        ax_e.legend(loc="upper right")

        # 第 3 行右：系统总能量
        ax_sys = axes[2, 1]
        ax_sys.plot(time, ek_sys, 'b-', linewidth=1.5, label="系统动能")
        ax_sys.plot(time, ug_sys, 'g-', linewidth=1.5, label="系统重力势能")
        if bond_pairs is not None and len(bond_pairs) > 0:
            ax_sys.plot(time, us_sys, color='orange', linewidth=1.5, label="系统弹性势能")
        ax_sys.plot(time, e_sys, 'r--', linewidth=1.5, label="系统总能量")
        ax_sys.set_title("系统总能量")
        ax_sys.set_xlabel("时间 (s)")
        ax_sys.set_ylabel("能量")
        ax_sys.grid(True, alpha=0.3)
        ax_sys.legend(loc="upper right")

        # 隐藏第 3 行第 3 列空子图
        axes[2, 2].set_visible(False)

    # 调整布局
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])

    # 保存图像
    output_path = os.path.join(output_dir, 'trajectory_plots.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"图表已保存至: {output_path}")

    # 显示图像（如果环境支持）
    plt.show()

    return output_path

def main():
    import argparse

    parser = argparse.ArgumentParser(description="绘制轨迹数据图表")
    parser.add_argument("trajectory_file", nargs="?",
                        help="轨迹数据文件路径（默认: output/trajectory.txt）")
    parser.add_argument("--atom-id", type=int, default=None,
                        help="指定只绘制某个原子的轨迹（默认绘制所有原子）")
    args = parser.parse_args()

    # 默认文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    default_file = os.path.join(compute.get_output_dir(script_dir), "trajectory.txt")
    input_file = args.trajectory_file or default_file

    # 检查文件是否存在
    if not os.path.exists(input_file):
        print(f"错误: 文件 '{input_file}' 不存在")
        print(f"请先运行 compute.py 生成轨迹数据")
        sys.exit(1)

    file_ext = os.path.splitext(input_file)[1].lower()

    try:
        if args.atom_id is not None:
            # 单原子模式（旧版兼容）
            traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, selected_id = \
                load_trajectory_txt(input_file, args.atom_id)
            ids_for_plot = selected_id
            extra_data = None
            n_atoms_str = f"  绘图原子序号: {selected_id}"
        else:
            # 全原子模式
            traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, NT, DT, atom_ids, raw_data = \
                load_full_trajectory(input_file)
            ids_for_plot = atom_ids
            extra_data = raw_data
            n_atoms = traj_x.shape[1] if traj_x.ndim > 1 else 1
            n_atoms_str = f"  绘图原子数: {n_atoms}"
    except Exception as e:
        print(f"加载文件时出错: {e}")
        sys.exit(1)

    print(f"加载轨迹数据: NT={NT}, DT={DT}")
    if args.atom_id is not None:
        print(n_atoms_str)
    else:
        print(n_atoms_str)
    print(f"  位置范围: x [{traj_x.min():.3f}, {traj_x.max():.3f}], "
          f"y [{traj_y.min():.3f}, {traj_y.max():.3f}], "
          f"z [{traj_z.min():.3f}, {traj_z.max():.3f}]")
    print(f"  速度范围: vx [{traj_vx.min():.3f}, {traj_vx.max():.3f}], "
          f"vy [{traj_vy.min():.3f}, {traj_vy.max():.3f}], "
          f"vz [{traj_vz.min():.3f}, {traj_vz.max():.3f}]")

    # 创建图表
    output_path = create_plots(
        traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz,
        NT, DT, compute.get_output_dir(script_dir), ids_for_plot, extra_data,
    )

    print("绘图完成！")

if __name__ == "__main__":
    main()
