from pathlib import Path
from typing import Dict, Iterable, List

import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.yolo.engine.results import Boxes as YOLOBoxes
from ultralytics.yolo.engine.results import Results as YOLOResults

from app import constants as c
from app.data_structures.object_detection import (
    DetectedObject,
    DetectedObjectsData,
    ObjectDetectionResult,
)
from app.exceptions import ObjectDetectionError
from app.loggers import object_detector as logger


class ObjectDetector:
    def __init__(self, model_path: Path) -> None:
        self.net = YOLO(str(model_path))
        self.classes_by_id = self.net.names

    def detect(self, file_path: Path, timestamp: float) -> ObjectDetectionResult:
        """
        Detects objects in the image and returns the result.
        :param file_path: path to the video file
        :param timestamp: timestamp of the frame in the video
        :return: ObjectDetectionResult object
        """
        logger.info(f"Detecting objects in {file_path}")
        # load the image from video
        image = self._load_image(file_path, timestamp)
        objects_data = DetectedObjectsData(image=None, objects=[])
        if image is not None:
            try:
                objects_data = self.process_frame(image)
            except Exception as ex:
                raise ObjectDetectionError(
                    f"Error while processing video {file_path}"
                    f" at time {timestamp}: {ex}"
                ) from ex
            logger.info(f"Detected {len(objects_data.objects)} objects")
        return ObjectDetectionResult(file_path=file_path, objects_data=objects_data)

    def process_frame(self, image: np.ndarray) -> DetectedObjectsData:
        # get the output from the network
        net_output = self._process_net(image)
        # extract detected objects
        result = self._extract_results(net_output)
        return result

    def _process_net(self, image: np.ndarray) -> List[YOLOResults]:
        """
        Processes the image through the network and returns the output.
        :param image: numpy array with image data
        :return: list of outputs from the network
        """
        # preprocess the image for network input
        results = self.net.predict(
            image,
            conf=c.OBJECT_DETECTION_CONFIDENCE_THRESHOLD,
            device="cpu",
        )
        return results

    def _extract_results(self, results: List[YOLOResults]) -> DetectedObjectsData:
        """
        Extracts detected objects from the output of the network.
        :param results: output of the network
        :return: information about detected objects
        """
        result = results[0]
        objects = []
        need_boxes = self._filter_boxes_by_class(
            result.boxes, allowed_classes=c.OBJECT_DETECTION_DETECT_ONLY_CLASSES
        )
        for box in need_boxes:
            class_id = int(box.cls)
            confidence = float(box.conf)
            objects.append(
                DetectedObject(
                    label=self.classes_by_id[class_id],
                    confidence=confidence,
                )
            )
        result.boxes = need_boxes
        rendered_image = result.plot()
        return DetectedObjectsData(
            image=rendered_image,
            objects=objects,
        )

    def _filter_boxes_by_class(self, boxes: YOLOBoxes, allowed_classes: Iterable[str]):
        """
        Filters boxes by list of allowed classes.
        :param boxes: list of boxes
        :param allowed_classes: list of allowed classes (if empty, do not filter)
        :return: list of filtered boxes
        """
        if not allowed_classes:
            return boxes
        allowed_class_ids = self._extract_class_ids_from_map(
            self.classes_by_id, need_classes=allowed_classes
        )
        return [box for box in boxes if box.cls in allowed_class_ids]

    @staticmethod
    def _load_image(file_path: Path, timestamp: float) -> np.ndarray:
        """
        Loads image from the video file at given timestamp.
        :param file_path: path to the video file
        :param timestamp: timestamp of the frame in the video to be loaded
        :return: numpy array with image data
        """
        cap = cv2.VideoCapture(str(file_path))
        cap.set(cv2.CAP_PROP_POS_MSEC, timestamp)
        ret, image = cap.read()
        cap.release()
        return image

    @staticmethod
    def _extract_class_ids_from_map(
        classes_by_id: Dict[int, str], need_classes: Iterable[str]
    ) -> List[int]:
        """
        Extracts class ids from the classes_by_id map by class names.
        :param classes_by_id: map of class ids and class names
        :param need_classes: list of class names which ids should be extracted
        :return: list of class ids
        """
        return [
            class_id
            for class_id, class_name in classes_by_id.items()
            if class_name in need_classes
        ]
