import matplotlib.pyplot as plt
import numpy as np
import cv2
import easyocr


class OCR:
    def __init__(self, image) -> None:
        self.image = image
        self.reader = easyocr.Reader(
            lang_list=['en'],
            gpu=False,
            model_storage_directory='./EasyOCR/model/',
            download_enabled=False,
            user_network_directory='./EasyOCR/user_network/'
        )

    def detection(self):
        img_arr = np.array(self.image, dtype=np.uint8)
        response = self.reader.readtext(image=img_arr)
        plot_inputs = []

        for box, text, conf in response:
            plot_inputs.append((text, np.array(box, dtype=np.float32)))

        fig, ax = plt.subplots(nrows=1, ncols=1)
        plot = self.drawAnnotations(
            image=img_arr, predictions=plot_inputs, ax=ax)
        return fig

    def drawAnnotations(self, image, predictions, ax=None):
        """Draw text annotations onto image.

        Args:
            image: The image on which to draw
            predictions: The predictions as provided by `pipeline.recognize`.
            ax: A matplotlib axis on which to draw.
        """
        if ax is None:
            _, ax = plt.subplots()
        ax.imshow(self.drawBoxes(image=image, boxes=predictions,
                                 boxes_format="predictions"))
        predictions = sorted(predictions, key=lambda p: p[1][:, 1].min())
        left = []
        right = []
        for word, box in predictions:
            if box[:, 0].min() < image.shape[1] / 2:
                left.append((word, box))
            else:
                right.append((word, box))
        ax.set_yticks([])
        ax.set_xticks([])
        for side, group in zip(["left", "right"], [left, right]):
            for index, (text, box) in enumerate(group):
                y = 1 - (index / len(group))
                xy = box[0] / np.array([image.shape[1], image.shape[0]])
                xy[1] = 1 - xy[1]
                ax.annotate(
                    text=text,
                    xy=xy,
                    xytext=(-0.05 if side == "left" else 1.05, y),
                    xycoords="axes fraction",
                    arrowprops={"arrowstyle": "->", "color": "r"},
                    color="r",
                    fontsize=14,
                    horizontalalignment="right" if side == "left" else "left",
                )
        return ax

    def drawBoxes(self, image, boxes, color=(255, 0, 0), thickness=1, boxes_format="boxes"):
        """Draw boxes onto an image.

        Args:
            image: The image on which to draw the boxes.
            boxes: The boxes to draw.
            color: The color for each box.
            thickness: The thickness for each box.
            boxes_format: The format used for providing the boxes. Options are
                "boxes" which indicates an array with shape(N, 4, 2) where N is the
                number of boxes and each box is a list of four points) as provided
                by `keras_ocr.detection.Detector.detect`, "lines" (a list of
                lines where each line itself is a list of (box, character) tuples) as
                provided by `keras_ocr.data_generation.get_image_generator`,
                or "predictions" where boxes is by itself a list of (word, box) tuples
                as provided by `keras_ocr.pipeline.Pipeline.recognize` or
                `keras_ocr.recognition.Recognizer.recognize_from_boxes`.
        """
        if len(boxes) == 0:
            return image
        canvas = image.copy()
        if boxes_format == "lines":
            revised_boxes = []
            for line in boxes:
                for box, _ in line:
                    revised_boxes.append(box)
            boxes = revised_boxes
        if boxes_format == "predictions":
            revised_boxes = []
            for _, box in boxes:
                revised_boxes.append(box)
            boxes = revised_boxes
        for box in boxes:
            cv2.polylines(
                img=canvas,
                pts=box[np.newaxis].astype("int32"),
                color=color,
                thickness=thickness,
                isClosed=True,
            )
        return canvas