yes_cmdr/common/renderer.py

import numpy as np
import pygame
import math
from .utils import *
from .agent import TileNode, Agent, AgentType

terrain_colors = [
    (245, 245, 220), # 米色 Plain
    (210, 180, 140), # 褐色 Road
    (224, 255, 255), # 青色 Water
    (173, 216, 230)  # 蓝色 Subwater
]

team_colors = [
    (255, 0, 0),    # 红色
    (0, 0, 255),    # 蓝色
    (0, 255, 0),    # 绿色
    (255, 255, 0),  # 黄色
    (255, 165, 0),  # 橙色
    (128, 0, 128),  # 紫色
    (0, 255, 255),  # 青色
    (255, 192, 203) # 粉色
]

icons_path = "./tianqiong/envs/icons"

unit_icons: Dict[str, Dict[int, pygame.Surface]] = {
    # AgentType.AntiAir: {0: pygame.image.load(f"{icons_path}/AntiAir_b.png"), 1: pygame.image.load(f"{icons_path}/AntiAir_r.png")},
    # AgentType.Airport: {0: pygame.image.load(f"{icons_path}/Airport_b.png"), 1: pygame.image.load(f"{icons_path}/Airport_r.png")},
    # AgentType.Artillery: {0: pygame.image.load(f"{icons_path}/Artillery_b.png"), 1: pygame.image.load(f"{icons_path}/Artillery_r.png")},
    # AgentType.Bomber: {0: pygame.image.load(f"{icons_path}/Bomber_b.png"), 1: pygame.image.load(f"{icons_path}/Bomber_r.png")},
    # AgentType.Carrier: {0: pygame.image.load(f"{icons_path}/Carrier_b.png"), 1: pygame.image.load(f"{icons_path}/Carrier_r.png")},
    # AgentType.Fighter: {0: pygame.image.load(f"{icons_path}/Fighter_b.png"), 1: pygame.image.load(f"{icons_path}/Fighter_r.png")},
    # AgentType.Helicopter: {0: pygame.image.load(f"{icons_path}/Helicopter_b.png"), 1: pygame.image.load(f"{icons_path}/Helicopter_r.png")},
    # AgentType.Infantry: {0: pygame.image.load(f"{icons_path}/Infantry_b.png"), 1: pygame.image.load(f"{icons_path}/Infantry_r.png")},
    # AgentType.Tank: {0: pygame.image.load(f"{icons_path}/Tank_b.png"), 1: pygame.image.load(f"{icons_path}/Tank_r.png")},
    # AgentType.TransportHelicopter: {0: pygame.image.load(f"{icons_path}/TransportHelicopter_b.png"), 1: pygame.image.load(f"{icons_path}/TransportHelicopter_r.png")},
    # AgentType.TransportShip: {0: pygame.image.load(f"{icons_path}/TransportShip_b.png"), 1: pygame.image.load(f"{icons_path}/TransportShip_r.png")},
    # AgentType.CombatShip: {0: pygame.image.load(f"{icons_path}/CombatShip_b.png"), 1: pygame.image.load(f"{icons_path}/CombatShip_r.png")},
}

class Renderer:
    def __init__(self, nodes: List[TileNode], hex_size=20):
        self.nodes = nodes
        self.hex_size = hex_size
        self.window_size, (self.offset_x, self.offset_y) = get_window_size_and_offsets(nodes, hex_size)

    def render(self, player_agents: List[Agent], spotted_enemy_agents: List[Agent], attack_agent: Agent=None, defend_agent: Agent=None) -> np.ndarray:
        screen = pygame.Surface(self.window_size)
        offset_x = self.offset_x
        offset_y = self.offset_y
        hex_size = self.hex_size
        nodes = self.nodes
        draw_hex_grid(screen, offset_x, offset_y, hex_size, nodes)
        for agent in player_agents:
            pos = agent.pos
            alpha = agent.endurance / agent.max_endurance * 255
            draw_unit(screen, offset_x, offset_y, hex_size, axial_to_pixel(pos, hex_size), alpha=alpha, agent_type=agent.agent_type, team_id=agent.team_id)
        for agent in spotted_enemy_agents:
            pos = agent.pos
            alpha = agent.endurance / agent.max_endurance * 255
            draw_unit(screen, offset_x, offset_y, hex_size, axial_to_pixel(pos, hex_size), alpha=alpha, agent_type=agent.agent_type, team_id=agent.team_id)

        if attack_agent and defend_agent:
            start_pos = axial_to_pixel(attack_agent.pos, hex_size)
            end_pos = axial_to_pixel(defend_agent.pos, hex_size)
            draw_hexagon(screen, hex_size, (start_pos[0] + offset_x, start_pos[1] + offset_y), team_colors[attack_agent.team_id], True, 5)
            draw_hexagon(screen, hex_size, (end_pos[0] + offset_x, end_pos[1] + offset_y), team_colors[attack_agent.team_id], True, 5)

        rgb_array = pygame.surfarray.array3d(screen)

        return np.transpose(rgb_array, (1, 0, 2))

# 轴坐标系到像素坐标的转换函数（尖角朝上）
def axial_to_pixel(pos: Tuple[int, int], hex_size):
    q, r = pos
    x = hex_size * math.sqrt(3) * q
    y = hex_size * 3/2 * r
    x += hex_size * math.sqrt(3) / 2 * r
    return x, y

# 绘制六边形的函数
def draw_hexagon(surface, hex_size, center, color=(255, 255, 255), only_frame=False, line_width=1):
    line_color = (169, 169, 169)  # 灰色
    angle_offset = math.pi / 6    # 顶点指向上方
    points = [
        (
            center[0] + hex_size * math.cos(angle_offset + math.pi / 3 * i),
            center[1] + hex_size * math.sin(angle_offset + math.pi / 3 * i)
        )
        for i in range(6)
    ]
    if not only_frame:
        pygame.draw.polygon(surface, color, points)
        pygame.draw.polygon(surface, line_color, points, width=1)  # 绘制六边形边框
    else:
        pygame.draw.polygon(surface, color, points, width=line_width)  # 绘制六边形边框

def get_window_size_and_offsets(nodes: List[TileNode], hex_size: int) -> Tuple[Tuple[int, int], Tuple[int, int]]:
    min_x = min_y = float('inf')
    max_x = max_y = float('-inf')
    offset_x = offset_y = float('-inf')

    for node in nodes:
        q, r = node.pos
        x, y = axial_to_pixel((q, r), hex_size)
        min_x = min(min_x, x)
        max_x = max(max_x, x)
        min_y = min(min_y, y)
        max_y = max(max_y, y)
        offset_x = max(offset_x, -x)
        offset_y = max(offset_y, -y)
    
    # 加一些边距
    margin = hex_size * 2
    offset_x += margin
    offset_y += margin
    width = int(max_x - min_x + 2 * margin)
    height = int(max_y - min_y + 2 * margin)

    width = (width - 15) // 16 * 16 + 16
    height = (height - 15) // 16 * 16 + 16

    return (width, height), (int(offset_x), int(offset_y))

# 绘制单位的函数，带透明度
def draw_unit(surface, offset_x, offset_y, hex_size, center, alpha=128, agent_type=None, team_id=None):
    # print(agent_type, team_id)
    if agent_type in unit_icons and team_id in unit_icons[agent_type]:
        # 如果有图标，用图标表示
        icon = unit_icons[agent_type][team_id]
        icon_size = hex_size * 1.2
        icon = pygame.transform.scale(icon, (icon_size, icon_size))  # 缩放图标至合适尺寸
        icon.set_alpha(alpha)  # 设置透明度
        surface.blit(icon, (center[0] + offset_x - icon_size // 2, center[1] + offset_y - icon_size // 2))
    else:
        # 如果没有图标，用圆圈表示单位
        unit_surface = pygame.Surface((hex_size, hex_size), pygame.SRCALPHA)
        unit_surface.set_alpha(alpha)  # 设置透明度
        color = team_colors[team_id]
        pygame.draw.circle(unit_surface, color, (hex_size // 2, hex_size // 2), hex_size // 3)
        surface.blit(unit_surface, (center[0] + offset_x - hex_size // 2, center[1] + offset_y - hex_size // 2))

# 主绘制功能
def draw_hex_grid(surface, offset_x, offset_y, hex_size, nodes: List[TileNode]):
    background_color = (255, 255, 255)  # 白色
    surface.fill(background_color)

    for node in nodes:
        q, r = node.pos
        color = terrain_colors[node.terrain_type.value]
        pixel_x, pixel_y = axial_to_pixel((q, r), hex_size)
        center = (pixel_x + offset_x, pixel_y + offset_y)
        draw_hexagon(surface, hex_size, center, color)