from .parts import *
from .utility import *
from copy import deepcopy
from .argument import ArgType
from concurrent.futures import ThreadPoolExecutor

__all__ = ["HuiMvOcr"]


class HuiMvOcr:
    __worker_count = 1

    def __init__(self, args: "ArgType"):
        self.det = Detector(args)
        self.cls = Classifier(args)
        self.rec = Recognizer(args)

        self.det_box_type = args.det_box_type
        self.drop_score = args.drop_score
        self.crop_image_res_index = 0

    def ocr_one(self, img, cls: "bool" = False, use_space: "bool" = True):
        ori_im = img.copy()
        dt_boxes = self.det(img)
        if dt_boxes is None:
            return None

        dt_boxes = sorted_boxes(dt_boxes)
        size = len(dt_boxes)
        img_crop_list = [...] * size

        for i in range(size):
            box = deepcopy(dt_boxes[i])
            if self.det_box_type == "quad":
                img_crop_list[i] = get_rotate_crop_image(ori_im, box)
            else:
                img_crop_list[i] = get_min_area_rect_crop(ori_im, box)

        if cls:
            img_crop_list = self.cls(img_crop_list)

        rec_res = self.rec(img_crop_list, use_space=use_space)

        filter_rec_res = []
        for text, score in rec_res:
            if score > self.drop_score:
                filter_rec_res.append((text, score))

        return filter_rec_res

    def ocr_multi(self, img_list, cls: "bool" = False, use_space: "bool" = True):
        pool = ThreadPoolExecutor(HuiMvOcr.__worker_count)
        loop = range(len(img_list))
        tasks = [pool.submit(self.ocr_one, img_list[i], cls, use_space) for i in loop]

        return [tasks[i].result() for i in loop]