"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Monday, 2nd Sept., 2024

    OBJECTIVE:

        To provide an easy way to assess images for blurriness.
        Tried and implemented using HuggingFace behaviour_models.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


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

# System-level activities:
import io

# For using the AI model:
import torch
from transformers import pipeline

# To downloads images from URLs:
import requests

# To read images:
from PIL import Image

# Common:
from shared import variables


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


# --- Nothing Yet


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


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  CLASSES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class AssessImageBlur:

    def __init__(
            self,
            model,
            blur_label = "Blur",
            clarity_label = "Normal",
            blur_threshold = 0.25,
            clarity_threshold = 0.65
    ):

        """
        This class assesses the input image and tells if it is blurry or clear.
        :param model: The model to use. Either the name of the HuggingFace repo, or the directory where the model is
            stored.
        :param blur_threshold: The max allowed blurriness (0 to 1 range).
        :param clarity_threshold: The minimum needed clarity (0 to 1 range).
        """

        # Note down the config:
        self.__model = model
        self.__device = "gpu" if torch.cuda.is_available() else "cpu"
        self.__blur_label = blur_label
        self.__clarity_label = clarity_label
        self.__blur_threshold = blur_threshold
        self.__clarity_threshold = clarity_threshold

        # Initialize the classifier:
        self.__classifier = pipeline(
            task = "image-classification",
            model = self.__model,
            device = self.__device
        )

    @staticmethod
    async def read_image(source):

        """
        Reads an image in whichever format it is provided and returns it as a PIL object.
        :param source: The image as either a path or a URL or a io.BytesIO object.
        :return: The image opened as a PIL object.
        """

        # If the input image is already a PIL image:
        if isinstance(source, Image.Image): return source

        # If a buffer is provided:
        elif isinstance(source, io.BytesIO):
            source.seek(0)
            return Image.open(source)

        # If a string is provided (local path or URL):
        elif isinstance(source, str):
            if source.startswith("http://") or source.startswith("https://"):
                response = await variables.http_client.get(source)
                return Image.open(io.BytesIO(response.content))
            else: return Image.open(source)

    async def classify(self, image):

        """
        To get the prediction of the model from the given input image.
        :param image: The image (as a PIL object or file path or a URL).
        :return: The dictionary of classes with their respective predictions.
        """

        image = await self.read_image(image)
        predictions = self.__classifier(image)
        classes = {p["label"]: p["score"] for p in predictions}
        return classes

    async def is_ok(
            self,
            image,
            blur_threshold = None,
            clarity_threshold = None
    ):

        """
        Checks if the image given to it can be used, or should be rejected.
        :param image: The image (as a PIL object or file path or a URL).
        :param blur_threshold: A custom threshold to test against. If not provided, the default will be taken that was
            provided when the instance was created.
        :param clarity_threshold: A custom threshold to test against. If not provided, the default will be taken that
            was provided when the instance was created.
        :return: True if the image is okay, else False.
        """

        image = await self.read_image(image)
        blur_threshold = blur_threshold or self.__blur_threshold
        clarity_threshold = clarity_threshold or self.__clarity_threshold
        classes = await self.classify(image)
        if (
                classes[self.__blur_label] <= blur_threshold and
                classes[self.__clarity_label] >= clarity_threshold
        ): return True
        return False


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


if __name__ == "__main__":

    import asyncio


    async def main():
        image = r"/home/developer/Downloads/low-res-check.png"
        assessor = AssessImageBlur(model = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/image_classification/BlurOrBokeh")
        usable = await assessor.is_ok(image)
        classes = await assessor.classify(image)
        print("IS OKAY:", usable)
        print("CLASSES:", classes)

    asyncio.run(main())