"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 4th Jul, 2024

    OBJECTIVE:

        To provide a class to detect objects in images using YOLO behaviour_models.

    REFERENCES:

        1) Code Examples: https://docs.ultralytics.com/usage/python/
        2) Models: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes

    DOWNLOADS:

        N/A

"""

#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# System-level activities:
import os

# To use YOLO architecture:
from ultralytics import YOLO

# To work with images:
from PIL import Image

# utils:
from utils_v2.string import json

# Common:
from shared import constants

# For debugging:
from icecream import IceCreamDebugger


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class YoloDetect:

    def __init__(
            self,
            model_file,
            debug = True,
            debug_prefix = "YOLO | "
    ):

        # Load the model:
        self._model = YOLO(model = model_file)

        # Initialize the dbugging tool:
        self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self._printer.disable()

    def disable_debug(self):
        self._printer.disable()

    def enable_debug(self):
        self._printer.enable()

    @staticmethod
    def _process_result(raw_result):

        # Extract the various types of results available from the inference:
        speed = {k: (v / 100.0) for k, v in raw_result.speed.items()}
        class_mapping = raw_result.names
        probabilities = raw_result.probs
        boxes = raw_result.boxes
        masks = raw_result.masks

        # Construct a default response:
        response_json = {
            "speed": speed,
            "classes": class_mapping,
            "boxes": None
        }

        # Object-Detection results:
        formatted_boxes = []
        for box in boxes:
            detected_class = int(box.cls.numpy()[0])
            x1, y1, x2, y2 = box.xyxy[0]
            formatted_boxes.append(
                {
                    "class": detected_class,
                    "className": class_mapping[detected_class],
                    "confidence": float(box.conf.numpy()[0]),
                    "x1": int(x1.numpy()),
                    "y1": int(y1.numpy()),
                    "x2": int(x2.numpy()),
                    "y2": int(y2.numpy())
                }
            )
        response_json["boxes"] = formatted_boxes

        # Done here:
        return response_json

    def predict(self, image, show = False, verbose = False):

        # Process the input image with the given task:
        results = self._model(
            source = image,
            show = show,
            verbose = verbose
        )

        # Based on the task, interpret the results:
        results_json = self._process_result(results[0])

        # Done here:
        return results_json


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    # model_file_path = os.path.join(constants.PROJECT_DIRECTORY, "ai", "yolo", "behaviour_models", "yolov8x.pt")
    model_file_path = r"/home/developer/PycharmProjects/utils/data/ai/models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt"
    # sample_image_path = r"/home/developer/Downloads/2_cats.jpg"
    sample_image_path = r"/home/developer/Downloads/flattened_image.jpg"
    my_yolo = YoloDetect(model_file = model_file_path)
    results = my_yolo.predict(image = Image.open(sample_image_path))
    print("FINAL RESULTS:", json.to_string(results))