"""
图 2（v2）：三维轨迹可视化 — 100 条轨迹版。

相对 fig2_trajectory_3d.py 的改动：
  - 轨迹数量：10 → 100
  - 目标器标记：大方块 → 小圆点 + 十字准星线
  - 视线锥：提升可见性（半透明锥面 + 更粗的边缘线）

用法:
    python -m Plots.fig2_trajectory_3d_v2
"""

import os
import sys
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from Plots.plot_config import apply_style, save_fig, COLORS, DOUBLE_COL

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.cm as cm


def draw_safety_sphere(ax, center, radius, color='red', alpha=0.08):
    """绘制防碰撞球面。"""
    u = np.linspace(0, 2 * np.pi, 30)
    v = np.linspace(0, np.pi, 20)
    xs = center[0] + radius * np.outer(np.cos(u), np.sin(v))
    ys = center[1] + radius * np.outer(np.sin(u), np.sin(v))
    zs = center[2] + radius * np.outer(np.ones_like(u), np.cos(v))
    ax.plot_surface(xs, ys, zs, alpha=alpha, color=color, linewidth=0)


def draw_los_cone(ax, apex, theta_deg, length, n_lines=36, alpha=0.10):
    """
    绘制视线锥。锥尖在 apex，沿 -y 方向展开。
    v2: 增大 alpha、线数和线宽以提升可见性；增加半透明锥面。
    """
    theta = np.radians(theta_deg)
    r = length * np.tan(theta)
    t = np.linspace(0, 2 * np.pi, n_lines, endpoint=False)

    # 锥面线（更粗、更高 alpha）
    for ti in t:
        x_end = apex[0] + r * np.cos(ti)
        z_end = apex[2] + r * np.sin(ti)
        y_end = apex[1] - length
        ax.plot([apex[0], x_end], [apex[1], y_end], [apex[2], z_end],
                color=COLORS['yellow'], alpha=0.25, linewidth=0.4)

    # 锥底圆（加粗）
    circle_t = np.linspace(0, 2 * np.pi, 100)
    cx = apex[0] + r * np.cos(circle_t)
    cz = apex[2] + r * np.sin(circle_t)
    cy = np.full_like(cx, apex[1] - length)
    ax.plot(cx, cy, cz, color=COLORS['yellow'], alpha=0.45, linewidth=0.8)

    # 半透明锥面（用三角面片填充，增强锥体感）
    n_seg = 48
    t_seg = np.linspace(0, 2 * np.pi, n_seg + 1)
    for j in range(n_seg):
        tri = [
            [apex[0], apex[1], apex[2]],
            [apex[0] + r * np.cos(t_seg[j]),   apex[1] - length, apex[2] + r * np.sin(t_seg[j])],
            [apex[0] + r * np.cos(t_seg[j+1]), apex[1] - length, apex[2] + r * np.sin(t_seg[j+1])],
        ]
        face = Poly3DCollection([tri], alpha=0.04, facecolor=COLORS['yellow'],
                                 edgecolor='none')
        ax.add_collection3d(face)


def draw_target_crosshair(ax, pos, size=12, color='k'):
    """
    用小圆点 + 三轴短线（十字准星）标记目标器位置，
    比大方块/菱形更精确、更美观。
    """
    ax.scatter(*pos, marker='o', s=30, c=color, zorder=12)
    half = size
    for dx, dy, dz in [(half,0,0), (0,half,0), (0,0,half)]:
        ax.plot([pos[0]-dx, pos[0]+dx],
                [pos[1]-dy, pos[1]+dy],
                [pos[2]-dz, pos[2]+dz],
                color=color, linewidth=0.7, alpha=0.6, zorder=11)


def plot_trajectory_3d_v2(
    data_path='Plots/data/eval_100full_trajectories.npz',
    out_dir='Plots',
):
    """绘制 3D 轨迹图（100 条版本）。"""
    apply_style()

    # 环境参数
    x_h = np.array([0.0, -60.0, 0.0])
    rho_safe = 15.0
    theta_los_deg = 60.0

    if not os.path.exists(data_path):
        print(f"数据文件不存在: {data_path}")
        print("请先运行带 100 条完整轨迹的评估脚本。")
        return

    data = np.load(data_path, allow_pickle=True)
    n_traj = int(data.get('n_full_saved', 100))

    fig = plt.figure(figsize=(DOUBLE_COL, DOUBLE_COL * 0.85))
    ax = fig.add_subplot(111, projection='3d')

    # ── 绘制安全约束 ──
    draw_safety_sphere(ax, [0, 0, 0], rho_safe, color='red', alpha=0.06)
    draw_los_cone(ax, [0, 0, 0], theta_los_deg, length=900, alpha=0.10)

    # ── 标记保持点 ──
    ax.scatter(*x_h, marker='*', s=100, c=COLORS['red'], zorder=10,
              label=r'Hold point $\mathbf{x}_h$')

    # ── 标记目标器（十字准星） ──
    draw_target_crosshair(ax, [0, 0, 0], size=18, color=COLORS['black'])

    # ── 绘制轨迹 ──
    n_success = 0
    n_fail = 0
    n_timeout = 0

    for i in range(n_traj):
        key = f'traj{i}_states'
        if key not in data:
            continue
        states = data[key]
        reason = str(data.get(f'traj{i}_reason', 'none'))

        x, y, z = states[:, 0], states[:, 1], states[:, 2]

        if reason == 'success':
            color = COLORS['blue']
            alpha = 0.55
            lw = 0.5
            n_success += 1
        elif reason == 'collision':
            color = COLORS['red']
            alpha = 0.7
            lw = 0.6
            n_fail += 1
        else:
            color = COLORS['grey']
            alpha = 0.25
            lw = 0.3
            n_timeout += 1

        # 轨迹线
        ax.plot(x, y, z, color=color, alpha=alpha, linewidth=lw)

        # 起始点（小圆点）
        ax.scatter(x[0], y[0], z[0], marker='o', s=6, c=color, alpha=0.5,
                  edgecolors='none')

    # ── 后向接近半空间指示线 ──
    lim = 250
    ax.plot([-lim, lim], [0, 0], [0, 0], 'k--', alpha=0.15, linewidth=0.3)

    # ── 坐标轴标签 ──
    ax.set_xlabel(r'$x$ (radial) [m]', labelpad=6)
    ax.set_ylabel(r'$y$ (along-track) [m]', labelpad=6)
    ax.set_zlabel(r'$z$ (normal) [m]', labelpad=6)

    # ── 视角 ──
    ax.view_init(elev=25, azim=-55)

    # ── 图例 ──
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color=COLORS['blue'], lw=1,
               label=f'Success ({n_success})'),
        Line2D([0], [0], color=COLORS['grey'], lw=1,
               label=f'Timeout ({n_timeout})'),
        Line2D([0], [0], marker='*', color='w', markerfacecolor=COLORS['red'],
               markersize=8, label=r'Hold point $\mathbf{x}_h$'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor=COLORS['black'],
               markersize=5, label='Target'),
    ]
    # 只在有碰撞轨迹时才加红色图例
    if n_fail > 0:
        legend_elements.insert(1,
            Line2D([0], [0], color=COLORS['red'], lw=1,
                   label=f'Collision ({n_fail})'))

    ax.legend(handles=legend_elements, loc='upper left', fontsize=6,
              framealpha=0.9)

    # ── 刻度 ──
    ax.tick_params(axis='both', which='major', labelsize=6, pad=2)

    save_fig(fig, 'fig2_trajectory_3d_v2', out_dir)
    plt.close(fig)
    print(f"✓ 图 2 (v2) 完成 — {n_success} success / {n_fail} collision / "
          f"{n_timeout} timeout (共 {n_traj} 轨迹)")


if __name__ == '__main__':
    plot_trajectory_3d_v2()