"""3D 显示验证 — 模拟圆周运动轨迹 (无需串口, 使用四元数)"""
import numpy as np
import time
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation
from trajectory_tracker import Tracker
from visualize_3d import TrajectoryViewer


def main():
    tracker = Tracker(zupt_frames=50)
    viewer = TrajectoryViewer(title="IMU 3D Demo — 四元数模拟")

    dt = 0.005   # 200Hz
    t = 0.0
    omega = 0.5  # 角速度 (rad/s)
    last_draw = time.time()

    print("3D 演示运行中... 按 Ctrl+C 停止")

    try:
        while plt.fignum_exists(viewer.fig.number):
            # X-Y 平面圆周运动 + Z 轴正弦起伏
            ax_linear = -0.25 * np.cos(omega * t)
            ay_linear = -0.25 * np.sin(omega * t)
            az_linear = 0.3 * np.sin(2 * t)

            # 模拟四元数 (绕 Z 轴旋转 omega*t)
            yaw = omega * t
            r = Rotation.from_euler('ZYX', [np.degrees(yaw), 0, 0])
            q = r.as_quat()  # [x, y, z, w]

            # body 加速度 = R^T @ (a_linear + gravity)
            a_world = np.array([ax_linear, ay_linear, az_linear + 9.81])
            accel_body = r.as_matrix().T @ a_world

            gyro = np.array([0.0, 0.0, omega])

            pos = tracker.update(gyro, accel_body, q[3], q[0], q[1], q[2], dt)

            now = time.time()
            if now - last_draw >= 0.033:
                viewer.update(tracker.history_array)
                plt.pause(0.001)
                last_draw = now

            t += dt

    except KeyboardInterrupt:
        print("停止。")
    finally:
        viewer.close()


if __name__ == '__main__':
    main()