from typing import Literal

import cv2

from mower.static import Digital
from mower.utils import config
from mower.utils import typealias as tp
from mower.utils.image import cropimg, thres2
from mower.utils.log import logger


class NumberRecognizer:
    def segment(self, img: tp.GrayImage) -> list[tuple[int, int, int, int]]:
        """
        传入二值化的图像，分割图像中的连通区域
        """
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        rect = [cv2.boundingRect(c) for c in contours]
        rect.sort(key=lambda c: c[0])  # 按照x坐标排序
        return rect

    def filter_rectangles(
        self,
        rect: list[tuple[int, int, int, int]],
        st: int | None = None,
        ed: int | None = None,
        limits: list[dict] = [],
    ) -> list:
        """
        过滤并替换不满足条件的矩形
        limits: **大的数字写前面**,如[{"w":20,"h":10,"char":"."},{"w":10,"h":3,"char":""}]则将w<10,h<3的矩形替换为"."，w<20,h<10的矩形替换为""
        """
        rect = rect[st:ed]
        if limits:
            for i in range(len(rect)):
                x, y, w, h = rect[i]
                for lim in limits:
                    if w < lim["w"] and h < lim["h"]:
                        rect[i] = lim["char"]
                        break
        return rect

    def number_match(
        self,
        img,
        rect: list,
        font: Literal["riic_base", "noto", "secret_front"] = "noto",
        target_range: list = range(10),
    ) -> str:
        if font == "riic_base":
            templates = Digital().riic_base_digits
        elif font == "noto":
            templates = Digital().noto_sans
        elif font == "secret_front":
            templates = Digital().secret_front
        value = ""
        for i in range(len(rect)):
            if len(rect[i]) != 4:
                value += rect[i]
                continue
            x, y, w, h = rect[i]
            digit = cropimg(img, ((x, y), (x + w, y + h)))
            digit = cv2.copyMakeBorder(
                digit, 10, 10, 10, 10, cv2.BORDER_CONSTANT, None, (0,)
            )
            score = {}
            for i in target_range:
                im = templates[i]
                result = cv2.matchTemplate(digit, im, cv2.TM_SQDIFF_NORMED)
                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
                score[min_val] = i
            value += str(score[(min(score))])

        return value

    def number(
        self,
        font: Literal["riic_base", "noto", "secret_front"] = "noto",
        scope: tp.Scope | None = None,
        height: int = 25,
        thres: int = 127,
        img: tp.GrayImage | None = None,
        rect_st: int | None = None,
        rect_ed: int | None = None,
        rect_limits: list[dict] = [],
        target_range: list = range(10),
        res_if_error: str = "",
    ) -> str:
        try:
            if img is None:
                img = cropimg(config.recog.gray, scope)
                default_height = 25
                if height != default_height:
                    scale = 25 / height
                    img = cv2.resize(img, None, None, scale, scale)
                img = thres2(img, thres)
            rect = self.segment(img)
            rect = self.filter_rectangles(rect, rect_st, rect_ed, rect_limits)
            return self.number_match(img, rect, font, target_range)
        except Exception as e:
            logger.exception(e)
            return res_if_error

    def number_int(
        self,
        font: Literal["riic_base", "noto", "secret_front"] = "noto",
        scope: tp.Scope | None = None,
        height: int = 25,
        thres: int = 127,
        img: tp.GrayImage | None = None,
        rect_st: int | None = None,
        rect_ed: int | None = None,
        rect_limits: list[dict] = [],
        target_range: list = range(10),
        res_if_error: str = "0",
    ) -> int:
        res = self.number(
            font,
            scope,
            height,
            thres,
            img,
            rect_st,
            rect_ed,
            rect_limits,
            target_range,
            res_if_error,
        )
        return int(res) if res else int(res_if_error)
