"""
IMU 3D 轨迹实时可视化

matplotlib 3D 窗口, 30Hz 刷新, 显示:
  - 蓝色轨迹线 (自动降采样, 最多 800 点)
  - 当前点红点
  - 原点坐标系指示
  - 等比例坐标轴
"""

import numpy as np
import matplotlib.pyplot as plt


class TrajectoryViewer:
    """3D 轨迹实时显示窗口"""

    def __init__(self, title="IMU 3D Odometry", max_trail_points=800):
        """
        Args:
            title:            窗口标题
            max_trail_points: 轨迹线最多显示点数 (降采样, 避免渲染卡顿)
        """
        self.max_trail_points = max_trail_points
        self._full_count = 0

        self.fig = plt.figure(figsize=(8, 7))
        self.fig.canvas.manager.set_window_title(title)
        self.ax = self.fig.add_subplot(111, projection='3d')

        # 轨迹线 (蓝色)
        self.traj_line, = self.ax.plot([], [], [], 'b-', linewidth=1.0, label='Trajectory')

        # 当前点 (红色)
        self.current_point, = self.ax.plot([], [], [], 'ro', markersize=6, label='Current')

        # 坐标系指示 (原点处)
        axis_len = 0.3
        self.origin_axes = [
            self.ax.quiver(0, 0, 0, axis_len, 0, 0, color='r', arrow_length_ratio=0.15, label='X'),
            self.ax.quiver(0, 0, 0, 0, axis_len, 0, color='g', arrow_length_ratio=0.15, label='Y'),
            self.ax.quiver(0, 0, 0, 0, 0, axis_len, color='b', arrow_length_ratio=0.15, label='Z'),
        ]

        self._setup_axes()

        # 显示非阻塞窗口
        plt.show(block=False)
        plt.pause(0.1)

    def _setup_axes(self):
        """初始化坐标轴"""
        self.ax.set_xlabel('X (front)')
        self.ax.set_ylabel('Y (left)')
        self.ax.set_zlabel('Z (up)')
        self.ax.set_title("IMU 3D Trajectory (Z-up)")
        self.ax.legend(loc='upper left')

        self.ax.set_xlim([-1, 1])
        self.ax.set_ylim([-1, 1])
        self.ax.set_zlim([-1, 1])

        try:
            self.ax.set_box_aspect([1, 1, 1])
        except NotImplementedError:
            pass

        self.ax.grid(True)

    def update(self, history_array):
        """更新显示的轨迹数据 (自动降采样)

        Args:
            history_array: Nx3 numpy array, 位置历史
        """
        n = len(history_array)
        if n < 1:
            return

        # 降采样: 超过 max_trail_points 时取等间隔子集
        if n > self.max_trail_points:
            step = n // self.max_trail_points
            indices = np.arange(0, n, step)
            sampled = history_array[indices]
        else:
            sampled = history_array

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

        self.traj_line.set_data(x, y)
        self.traj_line.set_3d_properties(z)

        # 当前点
        last = history_array[-1]
        self.current_point.set_data([last[0]], [last[1]])
        self.current_point.set_3d_properties([last[2]])

        # 自适应坐标轴 (每 20 次完整更新做一次, 减少开销)
        self._full_count += 1
        if self._full_count % 5 == 0:
            self._auto_scale(x, y, z)

    def _auto_scale(self, x, y, z):
        """自适应等比例坐标轴"""
        all_coords = np.concatenate([x, y, z])
        margin = max(np.ptp(all_coords) * 0.2, 0.5)
        mid = (all_coords.min() + all_coords.max()) / 2
        half = (all_coords.max() - all_coords.min()) / 2 + margin

        self.ax.set_xlim([mid - half, mid + half])
        self.ax.set_ylim([mid - half, mid + half])
        self.ax.set_zlim([mid - half, mid + half])

    def close(self):
        """关闭窗口"""
        plt.close(self.fig)