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

matplotlib 3D 窗口, 30Hz 刷新, 显示:
  - 蓝色轨迹线 (最近 800 点, Tracker 端已限长)
  - 当前点红点 + 朝向指示线 (RGB = body XYZ)
  - 原点世界坐标系指示
  - 等比例坐标轴
"""

import numpy as np
import matplotlib.pyplot as plt


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

    def __init__(self, title="IMU 3D Odometry"):
        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')

        # 朝向指示线 (body XYZ, 用 plot 可原地更新)
        self.body_x_line, = self.ax.plot([], [], [], 'r-', linewidth=2.0, alpha=0.9)
        self.body_y_line, = self.ax.plot([], [], [], 'g-', linewidth=2.0, alpha=0.9)
        self.body_z_line, = self.ax.plot([], [], [], 'b-', linewidth=2.0, alpha=0.9)

        # 世界坐标系指示 (固定在原点)
        axis_len = 0.3
        self.ax.quiver(0, 0, 0, axis_len, 0, 0, color='r', arrow_length_ratio=0.15)
        self.ax.quiver(0, 0, 0, 0, axis_len, 0, color='g', arrow_length_ratio=0.15)
        self.ax.quiver(0, 0, 0, 0, 0, axis_len, color='b', arrow_length_ratio=0.15)

        self._setup_axes()
        self._frame = 0

        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.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, rotation_matrix=None):
        """更新显示

        Args:
            history_array:   Nx3 位置历史 (已由 Tracker 限长为 800)
            rotation_matrix: 3x3 body→world 旋转矩阵 (None 则不显示朝向)
        """
        n = len(history_array)
        if n < 1:
            return

        x, y, z = history_array[:, 0], history_array[:, 1], history_array[:, 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]])

        # 朝向 (原地更新, 不重建对象)
        if rotation_matrix is not None:
            self._update_orientation(last, rotation_matrix)
        else:
            self.body_x_line.set_data([], [])
            self.body_x_line.set_3d_properties([])
            self.body_y_line.set_data([], [])
            self.body_y_line.set_3d_properties([])
            self.body_z_line.set_data([], [])
            self.body_z_line.set_3d_properties([])

        # 坐标轴范围 (每 5 帧更新一次)
        self._frame += 1
        if self._frame % 8 == 0:
            self._auto_scale(x, y, z)

    def _update_orientation(self, position, R):
        """更新 body 系朝向指示线 (不重建对象)

        Args:
            position: (3,) 当前位置
            R: 3x3 body→world 旋转矩阵
        """
        px, py, pz = position
        length = 0.15

        for line, col in [(self.body_x_line, 0),
                          (self.body_y_line, 1),
                          (self.body_z_line, 2)]:
            d = R[:, col] * length
            line.set_data([px, px + d[0]], [py, py + d[1]])
            line.set_3d_properties([pz, pz + d[2]])

    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)