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

class TargetFollower:
    def __init__(self, detections, annotated_image, camera_position=np.array([0, 0, 0])):
        self.detections = detections
        self.camera_position = camera_position
        self.target = None

        # Сохраняем аннотированное изображение
        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:
            print(info)

    def select_target(self, class_name, nearest=False):
        """
        Выбирает объект по указанному имени класса.
        Если nearest=True, выбирается ближайший объект по значению mean_depth.

        :param class_name: Имя класса, который нужно выбрать.
        :param nearest: Если True, выбирает ближайший объект с данным классом.
        :return: Найденный объект или None, если объект не найден.
        """
        if nearest:
            closest_detection = None
            min_depth = float('inf')

            for detection in self.detections:
                if detection['class_name'] == class_name:
                    if detection['mean_depth'] < min_depth:
                        min_depth = detection['mean_depth']
                        closest_detection = detection

            self.target = closest_detection
            if self.target is not None:
                print(f"Closest target selected: {self.target}")
            else:
                print(f"Target with class name '{class_name}' not found.")
            return self.target
        else:
            for detection in self.detections:
                if detection['class_name'] == class_name:
                    self.target = detection
                    print(f"Target selected: {self.target}")
                    return self.target

            print(f"Target with class name '{class_name}' not found.")
            return None

    def calculate_follow_vector(self):
        if self.target is None:
            print("Target not selected.")
            return None

        bbox = self.target['bbox']
        target_center_x = (bbox[0] + bbox[2]) / 2
        target_center_y = (bbox[1] + bbox[3]) / 2
        target_depth = self.target['mean_depth']

        target_position = np.array([target_center_x, target_center_y, target_depth])

        follow_vector = target_position - self.camera_position
        return follow_vector

    def get_target_position(self):
        if self.target is None:
            print("Target not selected.")
            return None

        bbox = self.target['bbox']
        target_center_x = (bbox[0] + bbox[2]) / 2
        target_center_y = (bbox[1] + bbox[3]) / 2
        target_depth = self.target['mean_depth']

        target_position = np.array([target_center_x, target_center_y, target_depth])
        return target_position