from Neko.RealSense import RealSenseController
import cv2
from ultralytics import YOLO
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from datetime import datetime

controller = RealSenseController()
model = YOLO('./models/yolov9e_object_classification.pt')

def inline_detection(controller):
    depth_matrix = controller.acquisition.get_depth_image()
    color_matrix = controller.acquisition.get_color_image()

    if depth_matrix is None or color_matrix is None:
        print("Не удалось получить изображения. Проверьте подключение камеры.")
        return []

    results = model(color_matrix)

    annotated_image = results[0].plot()

    detections_info = []

    for detection in results[0].boxes:
        # Извлекаем координаты ограничивающей рамки
        x1, y1, x2, y2 = map(int, detection.xyxy[0])

        # Извлекаем класс объекта
        class_id = int(detection.cls[0])
        class_name = model.names[class_id]

        # Извлекаем соответствующую область из карты глубины
        depth_values = depth_matrix[y1:y2, x1:x2]

        # Вычисляем среднее значение глубины
        mean_depth = np.mean(depth_values)

        # Сохраняем информацию в массив
        detections_info.append({
            'class_id': class_id,
            'class_name': class_name,
            'bbox': [x1, y1, x2, y2],
            'mean_depth': float(mean_depth)  # Преобразуем в стандартный float
        })

        # Наносим информацию на изображение
        cv2.putText(annotated_image, f'{class_name} {mean_depth:.2f}m', (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # Сохраняем аннотированное изображение
    output_path = './annotated_image.jpg'
    cv2.imwrite(output_path, annotated_image)
    print(f"Annotated image saved to {output_path}")

    # Опционально: выводим массив с информацией
    print("Detections info:")
    for info in detections_info:
        print(info)

    return detections_info

def plot_3d_bounding_boxes(detections_info):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    for detection in detections_info:
        # Извлекаем координаты бокса и среднюю глубину
        x1, y1, x2, y2 = detection['bbox']
        mean_depth = detection['mean_depth']
        class_name = detection['class_name']

        # Определяем размеры бокса
        width = x2 - x1
        height = y2 - y1
        depth = 0.05  # Задаём фиксированную глубину для визуализации

        # Определяем 8 точек для 3D-бокса
        box_points = [
            [x1, y1, mean_depth],
            [x2, y1, mean_depth],
            [x2, y2, mean_depth],
            [x1, y2, mean_depth],
            [x1, y1, mean_depth + depth],
            [x2, y1, mean_depth + depth],
            [x2, y2, mean_depth + depth],
            [x1, y2, mean_depth + depth]
        ]

        # Определяем грани бокса (шесть сторон)
        faces = [
            [box_points[0], box_points[1], box_points[5], box_points[4]],  # Верхняя грань
            [box_points[3], box_points[2], box_points[6], box_points[7]],  # Нижняя грань
            [box_points[0], box_points[3], box_points[7], box_points[4]],  # Левая грань
            [box_points[1], box_points[2], box_points[6], box_points[5]],  # Правая грань
            [box_points[0], box_points[1], box_points[2], box_points[3]],  # Передняя грань
            [box_points[4], box_points[5], box_points[6], box_points[7]]   # Задняя грань
        ]

        # Добавляем бокс в виде параллелепипеда
        box = Poly3DCollection(faces, facecolors='cyan', linewidths=1, edgecolors='r', alpha=0.25)
        ax.add_collection3d(box)

        # Добавляем метку с названием класса и глубиной
        ax.text((x1 + x2) / 2, (y1 + y2) / 2, mean_depth + depth, f'{class_name} {mean_depth:.2f}m',
                color='blue', fontsize=8)

    # Настраиваем оси для лучшей визуализации
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Depth (m)")
    ax.set_title("3D Bounding Boxes with Depth Information")

    # Опционально: установить ограничения осей для лучшего отображения
    all_x = [d['bbox'][0] for d in detections_info] + [d['bbox'][2] for d in detections_info]
    all_y = [d['bbox'][1] for d in detections_info] + [d['bbox'][3] for d in detections_info]
    all_z = [d['mean_depth'] for d in detections_info] + [d['mean_depth'] + 0.05 for d in detections_info]

    ax.set_xlim(min(all_x) - 50, max(all_x) + 50)
    ax.set_ylim(min(all_y) - 50, max(all_y) + 50)
    ax.set_zlim(min(all_z) - 50, max(all_z) + 50)

    plt.show()

def test_time_delta():
    start_time = datetime.now()
    detections_info = inline_detection(controller)
    end_time = datetime.now()
    delta_time = end_time - start_time
    print(f"Время выполнения: {delta_time}")

if __name__ == "__main__":
    detections_info = inline_detection(controller)
    if detections_info:
        plot_3d_bounding_boxes(detections_info)