import numpy as np
from PIL import Image
import cv2


# __all__ = []

def ver_code_rect_split(rect, max_width=50):
    """
    验证码中的矩形框切分
    :param rect: 矩形框(left, upper, width, height)
    :param max_width: 当矩形框超过此值，则会被切分
    :return: list(rect)
    """
    rect_num = rect[2] // max_width
    if rect_num <= 1:
        return [rect]

    rect_lst = []
    rect_width = rect[2] // rect_num
    for i in range(rect_num):
        rect_temp = list(rect)
        rect_temp[0] = rect_temp[0] + i * rect_width
        rect_temp[2] = rect_width
        rect_lst.append(tuple(rect_temp))
    return rect_lst


def ver_code_split(img, out_size=(32, 32), channel_first: bool = True, mode=None):
    if isinstance(img, str):
        img = Image.open(img)
    img_height = 132
    img_width = int(img.size[0] * (img_height / img.size[1]))
    # 先将图像数据放大，再做高斯滤波和中值滤波，最后二值化处理
    img = img.resize((img_width, img_height), resample=Image.BICUBIC).convert('RGB')
    img_pro = cv2.GaussianBlur(np.array(img.convert('L')), (3, 3), 0)
    img_pro = cv2.medianBlur(img_pro, 11)
    ret, img_pro = cv2.threshold(img_pro, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 先做边缘检测，再做轮廓检测(只需要外轮廓)
    img_pro = cv2.Canny(img_pro, 80, 200, L2gradient=True)
    cnts = cv2.findContours(img_pro.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
    # 可过滤一些非常小的轮廓，轮廓小于50的过滤，！注意：这个参数应当随图像尺寸大小变化面变化
    cnts = [c for c in cnts if len(c) >= 50]
    rects = [cv2.boundingRect(cf) for cf in cnts]  # 找出轮廓的外接矩形

    # 有的字母连接在一起，矩形会很宽，将其分割开来
    filter_rects = []
    for rt in rects:
        if rt[2] * rt[3] <= 1500:  # 矩形面积小于1500的过滤掉
            continue
        filter_rects.extend(ver_code_rect_split(rt))  # 将大的矩形框分割开

    # 将所有矩形内的字母提取出来
    letter_lst = []
    for x, y, w, h in sorted(filter_rects, key=lambda px: px[0]):
        img_temp = img.crop((x, y, x + w, y + h))
        if out_size is not None:  # 更改尺寸
            img_temp = img_temp.resize(out_size)
        if mode == "L":
            img_array = np.array(img_temp.convert(mode))
        else:
            img_array = np.array(img_temp)
            if channel_first:
                img_array = img_array.transpose((2, 0, 1))
        letter_lst.append(img_array)
    return letter_lst