import open3d as o3d
import numpy as np
import sqlite3
import signal
import sys
import os
import time

DB_PATH = '../inertial_data.db'

# ===== ЗАВЕРШЕНИЕ ВСЕГО ПРОЕКТА =====

def shutdown_all():
    """Останавливает все процессы проекта"""
    os.system("pkill -f imu_reader.py")
    os.system("pkill -f lidar_reader.py")
    os.system("pkill -f calculations")

# ===== ЧТЕНИЕ ИЗ SQLite =====

def read_trajectory(conn):
    cursor = conn.execute(
        "SELECT x, y, z, roll, pitch, yaw FROM trajectory ORDER BY timestamp"
    )
    rows = cursor.fetchall()
    if not rows:
        return np.zeros((0, 3)), np.zeros((0, 3))
    data = np.array(rows)
    return data[:, :3], data[:, 3:]

def read_lidar_points(conn):
    cursor = conn.execute("SELECT x, y, z FROM lidar_points")
    rows = cursor.fetchall()
    if not rows:
        return np.zeros((0, 3))
    return np.array(rows)

# ===== ПОВОРОТ =====

def get_rotation_matrix(roll, pitch, yaw):
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(roll), -np.sin(roll)],
        [0, np.sin(roll), np.cos(roll)]
    ])
    Ry = np.array([
        [np.cos(pitch), 0, np.sin(pitch)],
        [0, 1, 0],
        [-np.sin(pitch), 0, np.cos(pitch)]
    ])
    Rz = np.array([
        [np.cos(yaw), -np.sin(yaw), 0],
        [np.sin(yaw), np.cos(yaw), 0],
        [0, 0, 1]
    ])
    return Rz @ Ry @ Rx

# ===== MAIN =====

def main():
    conn = sqlite3.connect(DB_PATH)

    # Создаём окно Open3D
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="Инерциальное управление", width=1280, height=720)

    pcd = o3d.geometry.PointCloud()
    line_set = o3d.geometry.LineSet()
    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2)

    vis.add_geometry(pcd)
    vis.add_geometry(line_set)
    vis.add_geometry(frame)

    print("Визуализация запущена. Закройте окно для остановки проекта.")

    try:
        while True:
            # Проверяем закрыто ли окно
            if not vis.poll_events():
                print("Окно закрыто — останавливаем проект...")
                shutdown_all()
                break

            # Читаем новые данные из БД
            trajectory, angles = read_trajectory(conn)
            lidar_points = read_lidar_points(conn)

            # Обновляем облако точек лидара
            if len(lidar_points) > 0:
                pcd.points = o3d.utility.Vector3dVector(lidar_points)
                vis.update_geometry(pcd)

            # Обновляем траекторию
            if len(trajectory) > 1:
                lines = [[i, i+1] for i in range(len(trajectory)-1)]
                line_set.points = o3d.utility.Vector3dVector(trajectory)
                line_set.lines = o3d.utility.Vector2iVector(lines)
                vis.update_geometry(line_set)

            # Обновляем позицию объекта
            if len(trajectory) > 0:
                pos = trajectory[-1]
                roll, pitch, yaw = angles[-1]
                R = get_rotation_matrix(roll, pitch, yaw)
                frame_new = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2)
                frame_new.rotate(R, center=(0, 0, 0))
                frame_new.translate(pos)
                frame.vertices = frame_new.vertices
                frame.triangles = frame_new.triangles
                vis.update_geometry(frame)

            vis.update_renderer()

            # Обновляем каждые 200 мс
            time.sleep(0.2)

    finally:
        vis.destroy_window()
        conn.close()

main()