import cv2
from time import time
import numpy as np
from compress import compress_image

BorderWidth = 5


def isWhiteBorder(img: "np.ndarray") -> "bool":
    row = np.concatenate([img[:BorderWidth], img[-BorderWidth:]], axis=0).reshape((BorderWidth * 2, -1))
    col = np.concatenate([img[:, :BorderWidth], img[:, -BorderWidth:]], axis=1).reshape((BorderWidth * 2, -1))
    full = np.concatenate([row, col], axis=1)
    return np.average(full) > 127


def resize(img: "np.ndarray", height: "int" = None, width: "int" = None) -> "np.ndarray":
    if height is None and width is None:
        return img
    h, w = img.shape[:2]
    if height is not None:
        dim = int(w * height / h), height
    else:
        dim = width, int(h * width / w)

    return cv2.resize(img, dim, interpolation=cv2.INTER_AREA)  # noqa


def distance(pa: "list", pb: "list") -> "float":
    return np.sqrt((pa[0] - pb[0]) ** 2 + (pa[1] - pb[1]) ** 2)


def rect_area(points: "np.ndarray") -> "float":
    return distance(points[0], points[1]) * distance(points[0], points[3])


def transform(img: "np.ndarray", points: "np.ndarray") -> "np.ndarray":
    tl, bl, br, tr = points
    width = int(max(distance(tl, tr), distance(bl, br)))
    height = int(max(distance(tl, bl), distance(tr, br)))
    rect = np.array([tl, tr, bl, br], dtype="float32")
    dest = np.array(
        [[0, 0], [width - 1, 0],
         [0, height - 1], [width - 1, height - 1]],
        dtype="float32"
    )

    matrix = cv2.getPerspectiveTransform(rect, dest)
    return cv2.warpPerspective(img, matrix, (width, height))


def preprocess(img: "np.ndarray") -> "np.ndarray":
    # img = compress_image(img, 1000 * 600)
    # img = resize(img, height=500)
    # RGB -> gray
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 二值化
    thresh_type = cv2.THRESH_BINARY
    if isWhiteBorder(gray):
        thresh_type = cv2.THRESH_BINARY_INV
    _, thresh = cv2.threshold(gray, 158, 255, thresh_type)
    # 边缘检测
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return img
    max_contour = max(contours, key=cv2.contourArea)  # noqa 找到最大的轮廓
    # 图像校正
    approx = cv2.approxPolyDP(max_contour, 0.02 * cv2.arcLength(max_contour, True), True)
    if len(approx) == 4:  # 较为规则，可以校正
        return transform(gray, approx.reshape(4, 2))
    else:  # 不太规则，尝试裁剪
        rect = cv2.minAreaRect(max_contour)  # noqa 计算最小外接矩形
        box = np.intp(cv2.boxPoints(rect))  # noqa 获取矩形的四个角点
        if rect_area(box) * 5 < img.shape[0] * img.shape[1]:
            return img
        return img[min(box[:, 1]):max(box[:, 1]), min(box[:, 0]):max(box[:, 0])]


def main():
    # TODO 10  18
    import os
    dd = r"C:\Users\huimv\Pictures\Saved Pictures"
    for name in os.listdir(dd):
        if name.endswith(".ini"):
            continue
        img = cv2.imread(rf"{dd}\{name}")
        st = time()
        deal = preprocess(img)
        print(name, time() - st)
        cv2.imshow(name, deal)


if __name__ == '__main__':
    main()