"""
This class implements the YoloDetector. It is used to detect elements in a
given image or batch of images. It uses YOLOv6 as the backbone of the detector.

Author: Eric Canas.
Github: https://github.com/Eric-Canas
Email: eric@ericcanas.com
Date: 09-07-2022
"""
from __future__ import annotations
import numpy as np
import torch
from ultralytics import YOLO

import os

class YoloV8:
    def __init__(self, pt_path: str, inference_device:str = 'cuda', half_mode: bool = False):#, dataset_yml: str):

        """
        Initialize the YoloDetector (YoloV6).

        :param weights: String. Path to the weights file. It must be the .pt file generated by the YOLOv6 tool.
        :param dataset_yml: String. Path to the dataset .yaml file used for training the model in the YOLOv6 tool. It
                            must contain the 'nc' and 'names' fields.
        :param confidence_th: Float. The confidence threshold to use for assuming that a detection is valid. Take
                                into account that, when the detector is used within the pipeline together with the tracker,
                                this threshold will be used only for filtering high confidence detections with non-max
                                suppression, but it will also internally return everything below it. (To let the tracker
                                make its magic).
        :param nms_th: Float. The IoU threshold to use for non-max suppression.
        :param max_dets: Integer (greater than 0). The maximum number of detections to return by frame.
        :param inference_device: String ('cpu', 'cuda' or 'cuda:[int]'). The device to use for inference.
        :param image_size: Integer. The size of the images to use for inference.
        :param half_mode: Boolean. Whether to use half precision or not. Half precision only works with CUDA.
        :param merge_batches: Boolean. Whether to merge batches or not. If it is True, and the input is a batch
                                of images, predictions will be merged as if they came from a single image. It is only
                                recommended for very static scenes.
        :param agnostic_nms: Boolean. Whether to use agnostic nms or not. If it is True, the detector will use the
                                agnostic nms algorithm.
        """
        if inference_device.startswith('cuda'):
            assert torch.cuda.is_available(), 'CUDA is not available.'
        else:
            assert inference_device == 'cpu', f'Inference device must be cpu or cuda. Got {inference_device}.'
            assert half_mode is False, f'Half mode must be False for CPU inference.'
        assert os.path.isfile(pt_path), f'Weights file {pt_path} does not exist.'
        self.model = YOLO(model=pt_path, task='segment')

    def predict(self, image: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """
        Predicts the bounding boxes and classes of the given image.

        :param image: Numpy array. The image to predict the bounding boxes and classes of.
        :return: Tuple of two numpy arrays. The first one contains the bounding boxes in the format (x, y, w, h) and the
                    second one contains the classes.
        """
        # Predict
        results = self.model.predict(source=image)
        return results