"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 19th Sept., 2024

    OBJECTIVE:

        To provide a way to easily scan documents like printouts and visiting cards.

    REFERENCES:

        01. https://medium.com/@victorolufemi/build-a-document-scanner-with-opencv-ff9f645a4085
        02. https://github.com/JaidedAI/EasyOCR

    DOWNLOADS:

        N/A

"""


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


# To make sibling directories accessible for imports:
import sys

from diffusers.utils.import_utils import candidates

sys.path.append(".")
sys.path.append("..")

# System-level activities:
import io

# To work with image processing:
import cv2
import imutils
import numpy as np

# To run OCR:
import easyocr

# To download images from the web:
import requests

# My utils:
from utils_v2.string import json
from utils_v2.ai.object_detection.yolo import YoloDetect


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class DocumentScanner:

    def __init__(
            self,
            layout_detection_yolo = None,
            whitelisted_yolo_classes = None,
            ocr_languages = None
    ):

        self.__yolo = YoloDetect(
            model_file = layout_detection_yolo,
            debug = False
        )
        self.__whitelisted_yolo_classes = whitelisted_yolo_classes or []
        self.__ocr_engine = easyocr.Reader(ocr_languages or ["en"])
        pass

    @staticmethod
    def open_image(image):

        """
        Opens an image from various sources.
        :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
        :return: The image as a CV2 object.
        """

        # When the image is provided as a string,
        # it could be either as a URL or a path to a local file:
        if isinstance(image, str):
            if image.startswith("https://") or image.startswith("http://"):
                image = io.BytesIO(requests.get(image).content)
            else: image = cv2.imread(image)

        # If the image is provided as a io.BytesIO buffer:
        if isinstance(image, io.BytesIO):
            image.seek(0)
            image = np.asarray(bytearray(image.read()), dtype = np.uint8)
            image = cv2.imdecode(image, cv2.IMREAD_COLOR)

        # Done here:
        return image

    @staticmethod
    def polygon_area(points):

        """
        Computes the area occupied by a shape described by the array of points.
        Example input: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
        :return: The area of the shape.
        """

        n = len(points)
        area = 0.0
        for i in range(n):
            j = (i + 1) % n  # next vertex
            area += points[i][0] * points[j][1]
            area -= points[j][0] * points[i][1]
        return abs(area) / 2.0

    @staticmethod
    def show(image, title = "Preview", wait = True):

        """
        Just a quick wrapper to show the image in a window.
        :param image: The image that you want to show.
        :param title: The title of the window.
        :param wait: Set this to True when you want to show the window(s). This is useful when you want to show many
            windows at once. Suppose you want to show a lot of windows, you set this to False for all calls except the
            very last one.
        :return: None.
        """

        cv2.imshow(title, image)
        if wait:
            cv2.waitKey()
            cv2.destroyAllWindows()

    @staticmethod
    def enhance(image):

        """
        Applies a form of contrast boost to make the edges more easily visible.
        :param image: The OpenCV image that needs to be enhanced.
        :return: The enhanced CV2 image.
        """

        return cv2.detailEnhance(
            src = image,
            sigma_s = 20,
            sigma_r = 0.15
        )

    @staticmethod
    def get_edges(image):

        """
        Gets the edges in an image.
        :param image: A CV2 image.
        :return: The image with the edges detected.
        """

        # Denoise the image:
        processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        processed_image = cv2.GaussianBlur(
            processed_image,
            ksize = (5, 5),
            sigmaX = 0
        )

        # Detect edges in the image:
        processed_image = cv2.Canny(
            processed_image,
            threshold1 = 50,
            threshold2 = 200
        )

        # Close small gaps in the edges:
        kernel = np.ones((5, 5), np.uint8)
        processed_image = cv2.dilate(processed_image, kernel, iterations = 1)
        # processed_image = cv2.morphologyEx(processed_image, cv2.MORPH_CLOSE, kernel)
        # processed_image = cv2.erode(processed_image, kernel, iterations = 1)

        # Done here:
        return processed_image

    def scan(self, image, min_area = 0.125):

        """
        Looks for a rectangle in an image and flattens it out. No colour filters are applied here.
        :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
        :param min_area: The minimum area occupied by the document in the image.
        :return: The flattened image if any, else None.
        """

        # Open the image and read the data:
        image = self.open_image(image)
        if image is None: return None
        height, width = image.shape[:2]
        image_area = height * width

        # Enhance the image to have better visibility of edges:
        # processed_image = self.enhance(image.copy())

        # Detect edges in the image:
        processed_image = self.get_edges(image.copy())

        # Find the contours in the edges,
        # and sort them in ascending order:
        contours = cv2.findContours(processed_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        contours = imutils.grab_contours(contours)
        contours = sorted(contours, key = cv2.contourArea, reverse = True)

        # We try to get clean and passable contours.
        # For our purposes, 'clean' is when it can be approximated to exactly 4 sides,
        # and 'passable' is when the approximation comes close to it:
        clean_rects = []
        passable_rects = []
        for contour in contours:
            perimeter = cv2.arcLength(contour, closed = True)
            approximation = cv2.approxPolyDP(contour, 0.025 * perimeter, closed = True)
            if len(approximation) == 4: clean_rects.append(approximation)
            elif len(approximation) <= 32:
                candidate_contour = cv2.minAreaRect(approximation)
                candidate_contour = cv2.boxPoints(candidate_contour)
                candidate_contour = np.intp(candidate_contour)
                passable_rects.append(candidate_contour)
        clean_rects = sorted(clean_rects, key = cv2.contourArea, reverse = True)
        passable_rects = sorted(passable_rects, key = cv2.contourArea, reverse = True)

        # Now we decide between the best candidate:
        document_outline = None
        if len(clean_rects) > 0:
            candidate_rect = np.array([p[0] for p in clean_rects[0]])
            rectangle_area = cv2.contourArea(candidate_rect)
            if rectangle_area / image_area >= min_area: document_outline = candidate_rect
        if document_outline is None and len(passable_rects) > 0:
            candidate_rect = cv2.minAreaRect(passable_rects[0])
            candidate_rect = cv2.boxPoints(candidate_rect)
            candidate_rect = np.intp(candidate_rect)
            rectangle_area = cv2.contourArea(candidate_rect)
            if rectangle_area / image_area >= min_area: document_outline = candidate_rect

        # If there is no rectangular contour,
        # we exit with failure:
        if document_outline is None: return None

        # Since we did get the best candidate for the document,
        # we figure out which point n the contour is which corner of the document:
        document_rectangle = np.zeros((4, 2), dtype = "float32")
        sum_points = document_outline.sum(axis = 1)
        document_rectangle[0] = document_outline[np.argmin(sum_points)]
        document_rectangle[2] = document_outline[np.argmax(sum_points)]
        diff_points = np.diff(document_outline, axis = 1)
        document_rectangle[1] = document_outline[np.argmin(diff_points)]
        document_rectangle[3] = document_outline[np.argmax(diff_points)]
        (top_left, top_right, bottom_right, bottom_left) = document_rectangle

        # let's compute the dimensions of the document in the image:
        left_height = np.sqrt(((top_left[0] - bottom_left[0]) ** 2) + ((top_left[1] - bottom_left[1]) ** 2))
        right_height = np.sqrt(((top_right[0] - bottom_right[0]) ** 2) + ((top_right[1] - bottom_right[1]) ** 2))
        top_width = np.sqrt(((top_right[0] - top_left[0]) ** 2) + ((top_right[1] - top_left[1]) ** 2))
        bottom_width = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + ((bottom_right[1] - bottom_left[1]) ** 2))
        max_height = max(int(left_height), int(right_height))
        max_width = max(int(top_width), int(bottom_width))

        # We compute the destination of the transform:
        destination = np.array(
            object = [
                [0, 0],
                [max_width - 1, 0],
                [max_width - 1, max_height - 1],
                [0, max_height - 1]
            ],
            dtype = "float32"
        )

        # We apply the transform that flattens out the document:
        matrix = cv2.getPerspectiveTransform(document_rectangle, destination)
        flattened_image = cv2.warpPerspective(image, matrix, dsize = (max_width, max_height))

        # Done here:
        return flattened_image

    def extract_text(
            self,
            image,
            min_confidence = 0.5,
            margin = 0.05
    ):

        """
        To run OCR on the input image.
        :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer.
        :param min_confidence: The minimum amount of confidence for detected text to be considered.
        :param margin: The extra pixels to include when cropping into a section for OCR.
        :return: The extracted text.
        """

        # Make a variable that will hold the final result:
        extracted_text = {}

        # Open the image and read the data:
        image = self.open_image(image)
        height, width = image.shape[:2]

        # Run OCR on the entire canvas:
        ocr_result = [r[1] for r in self.__ocr_engine.readtext(image) if r[-1] >= min_confidence]
        extracted_text["fullDoc"] = " ".join(ocr_result)

        # Get the doc layout in the image:
        document_sections = self.__yolo.predict(image)
        document_sections = document_sections["boxes"]

        # For each section, we run the OCR process individually:
        section_results = []
        for section in document_sections:
            if section["class"] in self.__whitelisted_yolo_classes:
                section_width = section["x2"] - section["x1"]
                section_height = section["y2"] - section["y1"]
                x1 = max(int(section["x1"] - (margin * section_width)), 0)
                y1 = max(int(section["y1"] - (margin * section_height)), 0)
                x2 = min(int(section["x2"] + (margin * section_width)), width)
                y2 = min(int(section["y2"] + (margin * section_height)), height)
                sub_image = image[y1:y2, x1:x2]
                ocr_result = [r[1] for r in self.__ocr_engine.readtext(sub_image) if r[-1] >= min_confidence]
                section_results.append({
                    "text": " ".join(ocr_result),
                    "class": section["class"],
                    "className": section["className"],
                    "x1": x1,
                    "y1": y1,
                    "x2": x2,
                    "y2": y2
                })
        extracted_text["bySection"] = section_results

        # Done here:
        return extracted_text


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


# --- Nothing Yet


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


if __name__ == "__main__":

    import time

    my_scanner = DocumentScanner(
        layout_detection_yolo = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt",
        whitelisted_yolo_classes = [0, 1, 3, 4, 5, 7, 9, 10],
        # whitelisted_yolo_classes = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10],
        ocr_languages = ["en"]
    )

    image = my_scanner.open_image(r"/home/developer/Downloads/talkaholics_card.jpg")
    document_image = my_scanner.scan(image)
    if document_image is not None:
        cv2.imwrite(r"/home/developer/Downloads/flattened_image.jpg", document_image)
        start_time = time.time()
        document_text = my_scanner.extract_text(
            document_image,
            min_confidence = 0.5
        )
        print("OCR RESULT:")
        print(document_text)

        print(f"FINISHED IN {time.time() - start_time} SECONDS!")

    else: print("No image")
