# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import math
from typing import Any

import cv2

from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
from ultralytics.utils.plotting import colors


class DistanceCalculation(BaseSolution):
    """A class to calculate distance between two objects in a real-time video stream based on their tracks.

    This class extends BaseSolution to provide functionality for selecting objects and calculating the distance between
    them in a video stream using YOLO object detection and tracking.

    Attributes:
        selected_boxes (dict[int, Any]): Dictionary to store selected bounding boxes keyed by track ID.
        centroids (list[list[int]]): List to store centroids of selected bounding boxes.

    Methods:
        mouse_event_for_distance: Handle mouse events for selecting objects in the video stream.
        process: Process video frames and calculate the distance between selected objects.

    Examples:
        >>> distance_calc = DistanceCalculation()
        >>> frame = cv2.imread("frame.jpg")
        >>> results = distance_calc.process(frame)
        >>> cv2.imshow("Distance Calculation", results.plot_im)
        >>> cv2.waitKey(0)
    """

    def __init__(self, **kwargs: Any) -> None:
        """Initialize the DistanceCalculation class for measuring object distances in video streams."""
        super().__init__(**kwargs)

        self.selected_boxes: dict[int, list[float]] = {}
        self.centroids: list[list[int]] = []  # Store centroids of selected objects

    def mouse_event_for_distance(self, event: int, x: int, y: int, flags: int, param: Any) -> None:
        """Handle mouse events to select objects in a real-time video stream for distance calculation.

        A left click selects the tracked object under the cursor (up to two objects); a right click clears the
        selection.

        Args:
            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN).
            x (int): X-coordinate of the mouse pointer.
            y (int): Y-coordinate of the mouse pointer.
            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY).
            param (Any): Additional parameters passed to the function.

        Examples:
            >>> # Assuming 'dc' is an instance of DistanceCalculation
            >>> cv2.setMouseCallback("window_name", dc.mouse_event_for_distance)
        """
        if event == cv2.EVENT_LBUTTONDOWN:
            if len(self.selected_boxes) < 2:
                for box, track_id in zip(self.boxes, self.track_ids):
                    x0, y0, x1, y1 = self.get_enclosing_box(box)
                    if x0 < x < x1 and y0 < y < y1 and track_id not in self.selected_boxes:
                        self.selected_boxes[track_id] = box
                        break  # one object per click, even where boxes overlap

        elif event == cv2.EVENT_RBUTTONDOWN:
            self.selected_boxes = {}

    def forget_tracks(self, track_ids):
        """Drop retired IDs from `selected_boxes` so a vanished object can't freeze the distance."""
        super().forget_tracks(track_ids)
        for track_id in track_ids:
            self.selected_boxes.pop(track_id, None)

    def process(self, im0) -> SolutionResults:
        """Process a video frame and calculate the distance between two selected bounding boxes.

        This method extracts tracks from the input frame, annotates bounding boxes, and calculates the distance between
        two user-selected objects if they have been chosen.

        Args:
            im0 (np.ndarray): The input image frame to process.

        Returns:
            (SolutionResults): Contains processed image `plot_im`, `total_tracks` (int) representing the total number of
                tracked objects, and `pixels_distance` (float) representing the distance between selected objects
                in pixels.

        Examples:
            >>> import numpy as np
            >>> from ultralytics.solutions import DistanceCalculation
            >>> dc = DistanceCalculation()
            >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
            >>> results = dc.process(frame)
            >>> print(f"Distance: {results.pixels_distance:.2f} pixels")
        """
        self.extract_tracks(im0)  # Extract tracks
        annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator

        pixels_distance = 0
        # Iterate over bounding boxes, track ids and classes index
        for box, track_id, cls, conf in zip(self.boxes, self.track_ids, self.clss, self.confs):
            annotator.box_label(box, color=colors(int(cls), True), label=self.adjust_box_label(cls, conf, track_id))

            # Update selected boxes if they're being tracked
            if len(self.selected_boxes) == 2:
                for trk_id in self.selected_boxes:
                    if trk_id == track_id:
                        self.selected_boxes[track_id] = box

        if len(self.selected_boxes) == 2:
            # Calculate centroids of selected boxes
            self.centroids.extend(
                [
                    [int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)]
                    for box in map(self.get_enclosing_box, self.selected_boxes.values())
                ]
            )
            # Calculate Euclidean distance between centroids
            pixels_distance = math.sqrt(
                (self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
            )
            annotator.plot_distance_and_line(pixels_distance, self.centroids)

        self.centroids = []  # Reset centroids for next frame
        plot_im = annotator.result()
        self.display_output(plot_im)  # Display output with base class function
        if self.CFG.get("show") and self.env_check:
            cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)

        # Return SolutionResults with processed image and calculated metrics
        return SolutionResults(plot_im=plot_im, pixels_distance=pixels_distance, total_tracks=len(self.track_ids))