import torch
import numpy as np

# 工具函数：密度常量、剩余密度计算、掩码生成、邻接检测

MAP_W = 13 # 地图宽度
MAP_H = 13 # 地图高度
DENSITY_DIM = 5 # [wall, door, monster, entrance, resource]
MAP_SIZE = MAP_W * MAP_H # 地图大小

WALL_DENSITY_IDX = 0
DOOR_DENSITY_IDX = 1
MONSTER_DENSITY_IDX = 2
ENTRANCE_DENSITY_IDX = 3
RESOURCE_DENSITY_IDX = 4

def print_memory(device, tag=""):
    if torch.cuda.is_available():
        print(f"{tag} | 当前显存: {torch.cuda.memory_allocated(device) / 1024**2:.2f} MB, 最大显存: {torch.cuda.max_memory_allocated(device) / 1024**2:.2f} MB")
    else:
        print("当前设备不支持 cuda.")

def compute_remaining(
    current: torch.Tensor,
    target_density: torch.Tensor,
    stage: int
) -> torch.Tensor:
    remain = torch.zeros(current.size(0), DENSITY_DIM, device=current.device)

    visible_wall = (current == 1).sum(dim=1).float() / MAP_SIZE
    visible_door = ((current == 2) | (current == 6)).sum(dim=1).float() / MAP_SIZE
    visible_monster = (current == 4).sum(dim=1).float() / MAP_SIZE
    visible_entrance = (current == 5).sum(dim=1).float() / MAP_SIZE
    visible_resource = (current == 3).sum(dim=1).float() / MAP_SIZE

    if stage == 1:
        remain[:, WALL_DENSITY_IDX] = (
            target_density[:, WALL_DENSITY_IDX] - visible_wall
        ).clamp(min=0.0, max=1.0)
    elif stage == 2:
        remain[:, DOOR_DENSITY_IDX] = (
            target_density[:, DOOR_DENSITY_IDX] - visible_door
        ).clamp(min=0.0, max=1.0)
        remain[:, MONSTER_DENSITY_IDX] = (
            target_density[:, MONSTER_DENSITY_IDX] - visible_monster
        ).clamp(min=0.0, max=1.0)
        remain[:, ENTRANCE_DENSITY_IDX] = (
            target_density[:, ENTRANCE_DENSITY_IDX] - visible_entrance
        ).clamp(min=0.0, max=1.0)
    elif stage == 3:
        remain[:, RESOURCE_DENSITY_IDX] = (
            target_density[:, RESOURCE_DENSITY_IDX] - visible_resource
        ).clamp(min=0.0, max=1.0)

    return remain

def rect_mask(
    ratio: float, h_range: tuple[int, int] = (2, 7),
    w_range: tuple[int, int] = (2, 7)
) -> np.ndarray:
    # 纯矩形分块掩码，反复放置随机矩形直到掩码格数达标
    target = int(MAP_SIZE * ratio)
    mask = np.zeros((MAP_H, MAP_W), dtype=bool)
    while mask.sum() < target:
        bh = np.random.randint(h_range[0], h_range[1])
        bw = np.random.randint(w_range[0], w_range[1])
        x = np.random.randint(0, MAP_H - bh + 1)
        y = np.random.randint(0, MAP_W - bw + 1)
        mask[x:x + bh, y:y + bw] = True
    return mask

def compute_adjacency_mask(flat_state: torch.Tensor) -> torch.Tensor:
    # 返回与输入同形状的 bool tensor，True 表示该位置与任意墙壁 4-邻接
    # 支持 [MAP_SIZE] 和 [B, MAP_SIZE] 两种输入
    was_1d = flat_state.dim() == 1
    if was_1d:
        flat_state = flat_state.unsqueeze(0)
    state_2d = flat_state.reshape(flat_state.size(0), MAP_H, MAP_W)
    wall = (state_2d == 1)
    adj = torch.zeros_like(wall, dtype=torch.bool)
    adj[:, :, 1:] |= wall[:, :, :-1]
    adj[:, :, :-1] |= wall[:, :, 1:]
    adj[:, 1:, :] |= wall[:, :-1, :]
    adj[:, :-1, :] |= wall[:, 1:, :]
    result = adj.reshape(flat_state.size(0), MAP_SIZE)
    if was_1d:
        result = result.squeeze(0)
    return result

def summarize_codebook_hits(code_hits):
    # code_hits 为 tuple of 3 tensors（各量器不同 K）
    # 返回各阶段独立统计 + 汇总统计
    names = ["q1(stage1)", "q2(stage2)", "q3(stage3)"]
    result = {}
    for hits, name in zip(code_hits, names):
        total = hits.sum()
        if total.item() <= 0:
            result[name] = {"ppl": 0.0, "usage": 0.0, "active": 0, "K": int(hits.numel())}
            continue
        probs = hits / total
        ppl = torch.exp(
            -(probs * torch.log(probs.clamp_min(1e-10))).sum()
        ).item()
        active = int((hits > 0).sum().item())
        usage = active / hits.numel()
        result[name] = {"ppl": ppl, "usage": usage, "active": active, "K": int(hits.numel())}

    # 汇总统计
    combined = torch.cat([h.flatten() for h in code_hits], dim=0)
    total_hits = combined.sum()
    if total_hits.item() > 0:
        probs = combined / total_hits
        result["combined"] = {
            "ppl": float(torch.exp(
                -(probs * torch.log(probs.clamp_min(1e-10))).sum()
            ).item()),
            "active": int((combined > 0).sum().item()),
            "K": int(combined.numel()),
        }
    else:
        result["combined"] = {"ppl": 0.0, "active": 0, "K": int(combined.numel())}
    return result