import open3d as o3d
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from datetime import datetime

# СОЗДАНИЕ УНИКАЛЬНОЙ ПАПКИ

timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
output_dir = f"results/run_{timestamp}"
frames_dir = os.path.join(output_dir, "frames")

os.makedirs(frames_dir, exist_ok=True)

# ЗАГРУЗКА ДАННЫХ

traj = pd.read_csv("trajectory.csv")
lidar = pd.read_csv("lidar_points.csv")

trajectory = traj[['x', 'y', 'z']].values
angles = traj[['roll', 'pitch', 'yaw']].values
lidar_points = lidar[['x', 'y', 'z']].values

# СОХРАНЕНИЕ 2D ГРАФИКОВ

plt.figure()
plt.plot(trajectory[:,0], trajectory[:,1])
plt.title("Trajectory (X-Y)")
plt.xlabel("X")
plt.ylabel("Y")
plt.grid()
plt.savefig(os.path.join(output_dir, "trajectory_plot.png"))
plt.close()

# OPEN3D

pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(lidar_points)

lines = [[i, i+1] for i in range(len(trajectory)-1)]
line_set = o3d.geometry.LineSet()
line_set.points = o3d.utility.Vector3dVector(trajectory)
line_set.lines = o3d.utility.Vector2iVector(lines)

vis = o3d.visualization.Visualizer()
vis.create_window(visible=False)  # важно для сохранения кадров

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

# ПОВОРОТ

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

# СОХРАНЕНИЕ КАДРОВ

for i in range(len(trajectory)):

    frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2)

    pos = trajectory[i]
    roll, pitch, yaw = angles[i]

    R = get_rotation_matrix(roll, pitch, yaw)

    frame.rotate(R, center=(0,0,0))
    frame.translate(pos)

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

    vis.poll_events()
    vis.update_renderer()

    # сохраняем изображение
    img_path = os.path.join(frames_dir, f"frame_{i:04d}.png")
    vis.capture_screen_image(img_path)

# СОХРАНЕНИЕ ФИНАЛЬНОГО ВИДА

vis.capture_screen_image(os.path.join(output_dir, "3d_view.png"))

vis.destroy_window()

print(f"Результаты сохранены в: {output_dir}")