import numpy as np
import cv2
from typing import Tuple, List


BONE_NAMES = [
    "A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8",
    "J0", "J1", "J2", "J3", "J4", "J5", "J6", "J7", "J8",
    "B0", "C0", "D0", "E0", "F0", "G0", "H0", "I0",
    "B8", "C8", "D8", "E8", "F8", "G8", "H8", "I8",
]

def check_keypoints(keypoints: np.ndarray):
    """
    检查关键点坐标是否正确
    @param keypoints: 关键点坐标, shape 为 (34, 2)
    """
    if keypoints.shape != (34, 2):
        raise Exception(f"keypoints shape error: {keypoints.shape}")


def build_cells_xywh_by_cronners(corner_points: np.ndarray, padding: int = 3) -> np.ndarray:
    """
    根据 棋盘的 corner 点坐标 计算 每个位置的 xywh
    @param corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2)
    @param padding: 棋盘边框 padding

    @return: 棋盘的 xywh, shape 为 (10, 9, 4), 4 为 center_x, center_y, w, h
    """

    if corner_points.shape != (4, 2):
        raise Exception(f"corner_points shape error: {corner_points.shape}")
    
    top_left_xy = corner_points[0]
    top_right_xy = corner_points[1]
    bottom_left_xy = corner_points[2]
    bottom_right_xy = corner_points[3]
    
    # 计算 每个框的 w 和 h
    item_w = (top_right_xy[0] - top_left_xy[0]) / (9 - 1)
    item_h = (bottom_left_xy[1] - top_left_xy[1]) / (10 - 1)

    item_w = item_w
    item_h = item_h

    item_w_with_padding = item_w - padding * 2
    item_h_with_padding = item_h - padding * 2
    
    # 计算 每个框的 center 坐标
    cells_xywh = np.zeros((10, 9, 4))

    for i in range(10):
        for j in range(9):
            center_x = top_left_xy[0] + item_w * j
            center_y = top_left_xy[1] + item_h * i

            cells_xywh[i, j] = [center_x, center_y, item_w_with_padding, item_h_with_padding]

    return cells_xywh



# todo: 需要优化
def build_cells_xywh(keypoints: np.ndarray, width: int = 450, height: int = 500, padding: int = 3) -> np.ndarray:
    """
    @param keypoints: 关键点坐标, shape 为 (34, 2)
    @param width: 棋盘宽度
    @param height: 棋盘高度
    @param padding: 棋盘边框 padding
    @return: 棋盘的 xywh, shape 为 (10, 9, 4), 4 为 center_x, center_y, w, h
    """
    check_keypoints(keypoints)


    # 生成 A0 到 J8 的坐标, 如 B1 坐标 为 A1-J1 与 B0-B8 的交集点
    cells_xywh = np.zeros((10, 9, 4), dtype=np.int16)

    # 遍历 full_points 的每个点，计算其坐标
    for i in range(10):
        for j in range(9):
            # 计算 第 i 行 第 j 列 的坐标
            row_name = chr(ord('A') + i)
            col_name = str(j)
            flag_name = f"{row_name}{col_name}"
            if flag_name in BONE_NAMES:
                # 计算 第 i 行 第 j 列 的坐标
                cur_xy = keypoints[BONE_NAMES.index(flag_name)]
                cells_xywh[i, j] = [cur_xy[0], cur_xy[1], 0, 0]
            else:
                # 计算 第 i 行 第 j 列 的坐标
                row_start_name = f"{row_name}0"
                row_end_name = f"{row_name}8"
                
                col_start_name = f"A{col_name}"
                col_end_name = f"J{col_name}"

                row_start_xy = keypoints[BONE_NAMES.index(row_start_name)]
                row_end_xy = keypoints[BONE_NAMES.index(row_end_name)]

                col_start_xy = keypoints[BONE_NAMES.index(col_start_name)]
                col_end_xy = keypoints[BONE_NAMES.index(col_end_name)]

                # 计算 row_start_xy 到 row_end_xy 的直线 与 col_start_xy 到 col_end_xy 的直线 的交点
                # 使用参数方程法计算交点
                x1, y1 = row_start_xy  # 横向直线起点
                x2, y2 = row_end_xy    # 横向直线终点
                x3, y3 = col_start_xy  # 纵向直线起点
                x4, y4 = col_end_xy    # 纵向直线终点

                # 计算交点坐标
                # 使用克莱姆法则求解
                denominator = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
                
                # 计算交点的 x 坐标
                x = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denominator
                # 计算交点的 y 坐标
                y = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denominator
                
                cells_xywh[i, j] = [int(x), int(y), 0, 0]

    # 计算每个点位的 wh
    for i in range(10):
        for j in range(9):
            cur_xy = cells_xywh[i, j]
            # 获取上下左右 4 个点, 根据 4 个点计算 wh， 宽高为 4 个点 计算出来的 x1y1x2y2 的距离 的 1/2
            if i == 0:
                # [i+1, j] 的 反向点
                up_xy = 2 * cur_xy - cells_xywh[i+1, j]
            else:
                up_xy = cells_xywh[i - 1, j]
            
            if i == 9:
                # [i-1, j] 的 反向点
                down_xy = 2 * cur_xy - cells_xywh[i-1, j]
            else:
                down_xy = cells_xywh[i+1, j]

            if j == 0:
                left_xy = 2 * cur_xy - cells_xywh[i, j+1]
            else:
                left_xy = cells_xywh[i, j-1]

            if j == 8:
                right_xy = 2 * cur_xy - cells_xywh[i, j-1]
            else:
                right_xy = cells_xywh[i, j+1]

            min_x = min(up_xy[0].tolist(), down_xy[0].tolist(), left_xy[0].tolist(), right_xy[0].tolist())
            min_y = min(up_xy[1].tolist(), down_xy[1].tolist(), left_xy[1].tolist(), right_xy[1].tolist())

            min_x += padding
            min_y += padding

            # 防止 min_x 和 min_y 为 0
            min_x = max(min_x, 1)
            min_y = max(min_y, 1)

            max_x = max(up_xy[0].tolist(), down_xy[0].tolist(), left_xy[0].tolist(), right_xy[0].tolist())
            max_y = max(up_xy[1].tolist(), down_xy[1].tolist(), left_xy[1].tolist(), right_xy[1].tolist())

            max_x -= padding
            max_y -= padding

            # 防止 max_x 和 max_y 超出边界
            max_x = min(max_x, width - 1)
            max_y = min(max_y, height - 1)

            w = (max_x - min_x) / 2
            h = (max_y - min_y) / 2

            cells_xywh[i, j] = [int(cur_xy[0]), int(cur_xy[1]), int(w), int(h)]

    return cells_xywh


def perspective_transform(
        image: cv2.UMat, 
        src_points: np.ndarray, 
        keypoints: np.ndarray,
        dst_size=(450, 500)) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
    """
    透视变换
    @param image: 图片
    @param src_points: 源点坐标
    @param keypoints: 关键点坐标
    @param dst_size: 目标尺寸 (width, height) 10 行 9 列

    @return:
        result: 透视变换后的图片
        transformed_keypoints: 透视变换后的关键点坐标
        corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
    """

    check_keypoints(keypoints)


    # 源点和目标点
    src = np.float32(src_points)
    padding = 50
    corner_points = np.float32([
        # 左上角
        [padding, padding], 
        # 右上角
        [dst_size[0]-padding, padding], 
        # 左下角
        [padding, dst_size[1]-padding], 
        # 右下角
        [dst_size[0]-padding, dst_size[1]-padding]])

    # 计算透视变换矩阵
    matrix = cv2.getPerspectiveTransform(src, corner_points)

    # 执行透视变换
    result = cv2.warpPerspective(image, matrix, dst_size)

    # 重塑数组为要求的格式 (N,1,2)
    keypoints_reshaped = keypoints.reshape(-1, 1, 2).astype(np.float32)
    transformed_keypoints = cv2.perspectiveTransform(keypoints_reshaped, matrix)
    # 转回原来的形状
    transformed_keypoints = transformed_keypoints.reshape(-1, 2)

    return result, transformed_keypoints, corner_points



def get_board_corner_points(keypoints: np.ndarray) -> np.ndarray:
    """
    计算棋局四个边角的 points
    @param keypoints: 关键点坐标, shape 为 (34, 2)
    @return: 边角的坐标, shape 为 (4, 2)
    """
    check_keypoints(keypoints)

    # 找到 A0 A8 J0 J8 的坐标 以及 A4 和 J4 的坐标
    a0_index = BONE_NAMES.index("A0")
    a8_index = BONE_NAMES.index("A8")
    j0_index = BONE_NAMES.index("J0")
    j8_index = BONE_NAMES.index("J8")

    a0_xy = keypoints[a0_index]
    a8_xy = keypoints[a8_index]
    j0_xy = keypoints[j0_index]
    j8_xy = keypoints[j8_index]

    # 计算新的四个角点坐标
    dst_points = np.array([
        a0_xy,
        a8_xy,
        j0_xy,
        j8_xy
    ], dtype=np.float32)

    return dst_points

def extract_chessboard(img: cv2.UMat, keypoints: np.ndarray) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
    """
    提取棋盘信息
    @param img: 图片
    @param keypoints: 关键点坐标, shape 为 (34, 2)
    @return:
        transformed_image: 透视变换后的图片
        transformed_keypoints: 透视变换后的关键点坐标
        transformed_corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
    """

    check_keypoints(keypoints)

    source_corner_points = get_board_corner_points(keypoints)

    transformed_image, transformed_keypoints, transformed_corner_points = perspective_transform(img, source_corner_points, keypoints)

    return transformed_image, transformed_keypoints, transformed_corner_points


def collect_cells_images(image: cv2.UMat, cells_xywh: np.ndarray) -> List[List[np.ndarray]]:
    """
    收集 棋盘的 cells_xywh 对应的图片集合
    """
    width = image.shape[1]
    height = image.shape[0]
    crop_cells: List[List[np.ndarray]] = []

    for i in range(10):
        row_cells = []
        for j in range(9):
            x, y, w, h = cells_xywh[i, j]

            x_0 = max(int(x-w/2), 0)
            y_0 = max(int(y-h/2), 0)
            x_1 = min(int(x+w/2), width-1)
            y_1 = min(int(y+h/2), height-1)

            crop_img = image[y_0:y_1, x_0:x_1]
            row_cells.append(crop_img)
        crop_cells.append(row_cells)

    return crop_cells

def draw_cells_box(image: cv2.UMat, cells_xywh: np.ndarray) -> cv2.UMat:
    """
    绘制 棋盘的 cells_xywh 对应的 矩形框
    """
    width = image.shape[1]
    height = image.shape[0]
    for i in range(10):
        for j in range(9):
            x, y, w, h = cells_xywh[i, j]

            x_0 = max(int(x-w/2), 0)
            y_0 = max(int(y-h/2), 0)
            x_1 = min(int(x+w/2), width-1)
            y_1 = min(int(y+h/2), height-1)

            cv2.rectangle(image,(x_0, y_0), (x_1, y_1), (0, 0, 255), 1)

    return image