import cv2
import numpy as np

class Visualizer:
    """可视化：绘制灯条、装甲板、预测位置等"""
    def __init__(self):
        self.colors = {
            "red": (0, 0, 255),
            "blue": (255, 0, 0),
            "light_bar": (0, 255, 255),
            "armor": (0, 255, 0),
            "prediction": (0, 165, 255),
            "text": (255, 255, 255)
        }

    def draw_light_bars(self, frame, light_bars, target_color):
        """绘制灯条（矩形+中心线+标签）"""
        bar_color = self.colors[target_color]
        line_color = self.colors["light_bar"]

        for lb in light_bars:
            cv2.drawContours(frame, [lb["box"]], 0, bar_color, 2)
            end1 = tuple(map(int, lb["center_line"][0]))
            end2 = tuple(map(int, lb["center_line"][1]))
            cv2.line(frame, end1, end2, line_color, 2)
            center = tuple(map(int, lb["center"]))
            cv2.putText(
                frame, target_color, (center[0], center[1] - 8),
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, bar_color, 1
            )

        return frame

    def draw_armor_plate(self, frame, armor):
        """绘制装甲板（2D框+3D信息）"""
        corners = armor.get("corners_2d")
        if corners is not None:
            # 关键修改：np.int0 替换为 np.int32
            cv2.polylines(frame, [np.int32(corners)], True, self.colors["armor"], 2)

        center = tuple(map(int, armor["center"]))
        cv2.circle(frame, center, 3, self.colors["armor"], -1)
        cv2.putText(
            frame, "Armor", (center[0] + 5, center[1] - 5),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors["armor"], 2
        )

        return frame

    def draw_prediction(self, frame, predicted_pos, label="Predicted"):
        """绘制预测位置（橙色圆点+标签）"""
        if predicted_pos is None:
            return frame

        pos = tuple(map(int, predicted_pos))
        cv2.circle(frame, pos, 5, self.colors["prediction"], -1)
        cv2.putText(
            frame, label, (pos[0] + 10, pos[1] - 10),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.colors["prediction"], 2
        )

        return frame

    def draw_offset_text(self, frame, display_center, target_color, is_predicted):
        """绘制装甲板相对图像中心的偏移量"""
        if display_center is None:
            text = "Armor: unknown | Offset: (0, 0)"
        else:
            img_center = (frame.shape[1] // 2, frame.shape[0] // 2)
            offset_x = int(display_center[0] - img_center[0])
            offset_y = int(display_center[1] - img_center[1])
            color_text = f"{target_color} (predicted)" if is_predicted else target_color
            text = f"Armor: {color_text} | Offset: ({offset_x}, {offset_y})"

        cv2.putText(
            frame, text, (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, self.colors["text"], 2
        )

        return frame