import cv2
import numpy as np
import torch

def blend_alpha(bg, fg, alpha):
    """ 使用 alpha 通道混合前景图块和背景图 """
    for c in range(3):  # 只混合 RGB 三个通道
        bg[:, :, c] = (1 - alpha) * bg[:, :, c] + alpha * fg[:, :, c]
    return bg

def matrix_to_image_cv(map_matrix, tile_set, tile_size=32):
    """
    使用OpenCV加速的版本（适合大尺寸地图）
    :param map_matrix: [H, W] 的numpy数组
    :param tile_set: 字典 {tile_id: cv2图像（BGR格式）}
    :param tile_size: 图块边长（像素）
    """
    H, W = map_matrix.shape  # 获取地图尺寸
    canvas = np.zeros((H * tile_size, W * tile_size, 3), dtype=np.uint8)  # 画布（黑色背景）
    
    # 遍历地图矩阵
    for row in range(H):
        for col in range(W):
            tile_index = str(map_matrix[row, col])  # 获取当前坐标的图块类型
            x, y = col * tile_size, row * tile_size  # 计算像素位置

            # 先绘制地面（0）
            if '0' in tile_set:
                canvas[y:y+tile_size, x:x+tile_size] = tile_set['0'][:, :, :3]  # 仅填充 RGB

            # 叠加其他透明图块
            if tile_index in tile_set and tile_index != 0:
                tile_rgba = tile_set[tile_index]
                tile_rgb = tile_rgba[:, :, :3]  # 提取 RGB
                alpha = tile_rgba[:, :, 3] / 255.0  # 归一化 alpha

                # 混合当前图块到背景
                canvas[y:y+tile_size, x:x+tile_size] = blend_alpha(
                    canvas[y:y+tile_size, x:x+tile_size], tile_rgb, alpha
                )

    return canvas

def annotate(img: np.ndarray, text: str, y: int = 14) -> np.ndarray:
    # 在图片左上角叠加文字标注（黑色描边 + 白色填充，确保任意背景下可读）
    img = img.copy()
    cv2.putText(img, text, (2, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 0), 2)
    cv2.putText(img, text, (2, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
    return img

def annotate_labels(
    img: np.ndarray,
    struct: torch.Tensor,
    target_density: torch.Tensor
) -> np.ndarray:
    # 三行标注：第一行结构标签，后两行显示五维目标密度
    s = struct.tolist()
    d = target_density.tolist()
    line1 = f"sym:{s[0]} outer:{s[1]}"
    line2 = f"wall:{d[0]:.2f} door:{d[1]:.2f}"
    line3 = f"enemy:{d[2]:.2f} ent:{d[3]:.2f} res:{d[4]:.2f}"
    img = img.copy()
    for text, y in [(line1, 12), (line2, 24), (line3, 36)]:
        cv2.putText(img, text, (2, y), cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 0), 2)
        cv2.putText(img, text, (2, y), cv2.FONT_HERSHEY_SIMPLEX, 0.35, (255, 255, 255), 1)
    return img