import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import sqlite3
import os
import numpy as np

DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'inertial_data.db')

TRAJECTORY_LIMIT = 5000
LIDAR_LIMIT = 2000


def read_trajectory():
    # Новое соединение на каждый кадр — иначе видим устаревший снимок БД.
    with sqlite3.connect(DB_PATH) as conn:
        cursor = conn.execute(
            "SELECT x, y, z FROM trajectory ORDER BY timestamp DESC LIMIT ?",
            (TRAJECTORY_LIMIT,)
        )
        rows = cursor.fetchall()
    if not rows:
        return np.zeros((0, 3))
    return np.array(rows[::-1])  # от старых к новым


def read_lidar_points():
    with sqlite3.connect(DB_PATH) as conn:
        cursor = conn.execute(
            "SELECT x, y, z FROM lidar_points ORDER BY rowid DESC LIMIT ?",
            (LIDAR_LIMIT,)
        )
        rows = cursor.fetchall()
    if not rows:
        return np.zeros((0, 3))
    return np.array(rows)


def main():
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')

    def on_close(event):
        print("Окно визуализации закрыто.")

    fig.canvas.mpl_connect('close_event', on_close)

    def update(frame):
        ax.cla()
        ax.set_title("Инерциальное управление")
        ax.set_xlabel("X, м")
        ax.set_ylabel("Y, м")
        ax.set_zlabel("Z, м")

        trajectory = read_trajectory()
        lidar_points = read_lidar_points()

        # Точка старта — для ориентира всегда показываем
        ax.scatter([0], [0], [0], c='black', s=80, marker='x',
                   label='Старт (0, 0, 0)')

        if len(lidar_points) > 0:
            ax.scatter(
                lidar_points[:, 0],
                lidar_points[:, 1],
                lidar_points[:, 2],
                c='blue', s=1, alpha=0.5, label='Лидар'
            )

        if len(trajectory) > 1:
            ax.plot(
                trajectory[:, 0],
                trajectory[:, 1],
                trajectory[:, 2],
                c='red', linewidth=2, label='Траектория'
            )
            # Текущая позиция датчика — последняя точка траектории
            ax.scatter(
                [trajectory[-1, 0]],
                [trajectory[-1, 1]],
                [trajectory[-1, 2]],
                c='green', s=100, label='Позиция (сейчас)'
            )

        ax.legend(loc='upper right')

    ani = animation.FuncAnimation(fig, update, interval=200, cache_frame_data=False)

    print("Визуализация запущена. Закройте окно для выхода.")
    print("Чёрный крестик — точка старта (0, 0, 0).")
    print("Зелёная точка — текущая позиция датчика.")
    print("Красная линия — путь от старта до текущей позиции.")
    plt.show()


main()