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

class CameraProcessor:
    def __init__(self, model_path='./models/yolov9e_object_classification.pt'):
        self.controller = RealSenseController()
        self.model = YOLO(model_path)

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

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

        results = self.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 = self.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)
            })

            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)

        return detections_info, annotated_image

    def plot_3d_bounding_boxes(self, 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

            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(self):
        start_time = datetime.now()
        detections_info, _ = self.inline_detection()
        end_time = datetime.now()
        delta_time = end_time - start_time
        print(f"Время выполнения: {delta_time}")