# Import machine learning
from ultralytics import YOLO

# Import computer vision
import cv2
from PIL import Image


def predict(detection_model: YOLO, image_path: str) -> Image:
    """
    The function predicts with the given model and returns the image with the bounding boxes

    Args:
    detection_model (YOLO): The YOLO model
    image_path (str): The path to the image

    Returns:
    Image: The image with the bounding boxes
    """
    # Read the image
    img = cv2.imread(image_path)

    # Predict with the model
    output = detection_model.predict(image_path, verbose=False)

    # If output is not None then draw the bounding boxes
    if output:
        # Loop through the output
        for bbox in output:
            for i, p in enumerate(bbox):

                # Converting the p to cpu
                p = p.to("cpu")

                # Conveting to numpy
                p = p.numpy()

                # Extracting the coords
                coords = p.boxes.xyxy[0]

                # Extracting the coords
                xmin = coords[0]
                ymin = coords[1]
                xmax = coords[2]
                ymax = coords[3]

                # Converting to int
                xmin = int(xmin)
                ymin = int(ymin)
                xmax = int(xmax)
                ymax = int(ymax)

                # Extracting the prob
                prob = p.boxes.conf[0]

                # Drawing the bounding box
                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

                # Drawing the text with the probability
                cv2.putText(
                    img,
                    f"{prob:.2f}",
                    (xmin, ymin - 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.9,
                    (36, 255, 12),
                    2,
                )

        # Converting the image to RGB
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Converting the image to PIL Image object
        pil_image = Image.fromarray(img_rgb)

        # Return the image
        return pil_image

    # If output is None then return the original image
    return Image.open(image_path)