CADx-4/src/cadx/ceus/_overlay.py

"""overlay.mp4：把病灶轮廓、暗区、背景圆叠回整幅画面逐帧渲染，供人目视检查。

各视图局部坐标的多边形加上所属视图的左上角偏移，画到整帧坐标上。颜色（BGR）和线宽照参考组件：
造影那一半（source）的 ROI 黄 (0,255,255)、分割那一半（target）的 ROI 绿 (0,255,0)、暗区橙红
(255,80,20)；背景圆画青 (255,255,0) 而不照参考组件画紫，理由见 `_BACKGROUND_CIRCLE_CYAN`。
ROI 半透明填充 alpha=0.18；线条是黑色外描边加彩色内线。

`render_overlay` 接收已解码的逐帧整帧图像（可以是迭代器）和逐帧图层，返回 mp4 字节：
`cv2.VideoWriter` 不能直写内存，所以写临时文件、读回、在 finally 里删掉。fps 由调用方给，非正报错，
不回退 25。

照搬旧项目 `overlay.py`，去掉产品层；`OverlayConfig`、`OverlayLayers` 的字段默认值去掉，调用方全部
显式给出。
"""

import tempfile
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import cv2
import numpy as np

from cadx.ceus._geometry import AnyArray, ViewRect

# 各层颜色（BGR，照旧项目 findings §3.10 / 参考组件 P:1469-1544）。
_OUTLINE_BLACK = (0, 0, 0)
_SOURCE_ROI_YELLOW = (0, 255, 255)
_TARGET_ROI_GREEN = (0, 255, 0)
_DARK_ORANGE_RED = (255, 80, 20)
_BACKGROUND_CIRCLE_CYAN = (255, 255, 0)
"""自动挑的那个背景圆画成青色。**不能是紫**——紫色是医生画在同一画面上的 ``ROI_pan``
轮廓（邻近正常胰腺实质，也就是病变旁边那块没病的胰腺腺体组织）的颜色，而这个圆是程序
自己找的一小块干净背景，两者是不同的东西（旧项目 design/0030 决策四）。参考实现
把它画成品红 (255,0,255)，而紫色轮廓已经在算了（``tic.py`` 的 ``CURVE_NAMES`` 四条里就有
``roi_pan``），两圈都是紫的话叠加在视频上人眼分不出哪个是参照区、哪个是医生画的——
而**看 overlay 是第一步唯一的人工验收手段**，分不出就等于没验。
青色和绿框、紫色轮廓在色相上都隔得开。"""
# ROI 半透明填充透明度（P:1500 硬编码 0.18）。
_FILL_ALPHA = 0.18


@dataclass(frozen=True)
class OverlayConfig:
    """各画层开关与填充：画不画造影那一半的 ROI（黄）、分割那一半的 ROI（绿）、分割那一半的暗区、
    造影那一半的暗区、背景圆（青），以及要不要对分割那一半的 ROI 做 alpha=0.18 半透明填充。"""

    draw_source_roi: bool
    draw_target_polygon: bool
    draw_dark_region: bool
    draw_source_dark_region: bool
    draw_reference_region: bool
    fill_polygon: bool


@dataclass(frozen=True)
class OverlayLayers:
    """一帧要画的各层几何（多边形都是各自视图的局部坐标）。

    - `target_polygon`、`target_rect`：分割那一半的 ROI 和它在整帧中的矩形；
    - `dark_polygons`：分割那一半的暗区多边形（可为空列表）；
    - `source_roi_polygon`、`source_rect`：造影那一半的 ROI 和矩形（不画时可为 None）；
    - `source_dark_polygons`：造影那一半的暗区多边形（可为空列表）；
    - `reference_region`：背景圆元数据（`enabled`、`center_xy`、`radius_px`，圆心是造影那一半的
      局部坐标），没有时为 None。
    """

    target_polygon: AnyArray
    target_rect: ViewRect
    dark_polygons: list[AnyArray]
    source_roi_polygon: AnyArray | None
    source_rect: ViewRect | None
    source_dark_polygons: list[AnyArray]
    reference_region: dict[str, Any] | None


def _to_global_points(polygon: AnyArray, rect: ViewRect) -> AnyArray:
    """把视图局部坐标多边形抬到整帧坐标，并量化为 cv2 折线点 ``(M, 1, 2)`` int32。

    参数:
        polygon: 形如 ``(N, 2)`` 的局部坐标顶点。
        rect: 该多边形所属视图在整帧中的矩形（提供左上角偏移）。

    返回:
        ``cv2.polylines`` / ``cv2.fillPoly`` 直接可用的 ``(N, 1, 2)`` int32 点阵。

    关键实现:
        加 ``rect.x1`` / ``rect.y1`` 偏移后四舍五入取整（参考 P:1487-1489），先 copy
        避免就地污染入参多边形。
    """
    global_poly = polygon.copy().astype(np.float64)
    global_poly[:, 0] += rect.x1
    global_poly[:, 1] += rect.y1
    return np.round(global_poly).astype(np.int32).reshape(-1, 1, 2)


def _draw_outlined_polygon(
    frame: AnyArray,
    points: AnyArray,
    color: tuple[int, int, int],
    outline_thickness: int,
    color_thickness: int,
) -> None:
    """就地画"黑色外描边 + 彩色内线"的闭合多边形（参考 P:1490-1491 等的双描边手法）。

    参数:
        frame: 待绘制的整帧（就地修改）。
        points: ``(N, 1, 2)`` int32 折线点阵（整帧坐标）。
        color: 彩色内线 BGR。
        outline_thickness: 黑色外描边线宽（> 彩色线宽，作描边效果）。
        color_thickness: 彩色内线线宽。
    """
    cv2.polylines(
        frame,
        [points],
        isClosed=True,
        color=_OUTLINE_BLACK,
        thickness=outline_thickness,
    )
    cv2.polylines(frame, [points], isClosed=True, color=color, thickness=color_thickness)


def _fill_target_polygon(frame: AnyArray, points: AnyArray) -> AnyArray:
    """对 target ROI 内部做 alpha=0.18 绿色半透明填充（参考 P:1497-1500）。

    参数:
        frame: 当前帧。
        points: target ROI 整帧坐标折线点阵。

    返回:
        混合后的新帧（``addWeighted(fill, 0.18, frame, 0.82, 0)``）。

    关键实现:
        在帧副本上实心填充绿色后与原帧按 0.18 / 0.82 加权混合，得到半透明效果。
    """
    fill = frame.copy()
    cv2.fillPoly(fill, [points], _TARGET_ROI_GREEN)
    return cv2.addWeighted(fill, _FILL_ALPHA, frame, 1.0 - _FILL_ALPHA, 0)


def _draw_source_roi_layer(frame: AnyArray, layers: OverlayLayers) -> None:
    """画左 / source ROI 轮廓（黄，线宽 5→2，参考 P:1481-1491）。需 source_rect。"""
    if layers.source_roi_polygon is None or layers.source_rect is None:
        return
    points = _to_global_points(layers.source_roi_polygon, layers.source_rect)
    _draw_outlined_polygon(frame, points, _SOURCE_ROI_YELLOW, 5, 2)


def _draw_dark_layer(
    frame: AnyArray,
    polygons: Sequence[AnyArray],
    rect: ViewRect,
    outline_thickness: int,
    color_thickness: int,
) -> None:
    """画一组暗区轮廓（橙红，参考 P:1504-1526）；逐多边形抬到整帧再双描边。

    参数:
        frame: 当前帧（就地修改）。
        polygons: 暗区局部坐标多边形列表。
        rect: 这些暗区所属视图矩形（target 或 source）。
        outline_thickness: 黑色外描边线宽（target 6 / source 5，参考 P:1510/1525）。
        color_thickness: 橙红内线线宽（target 3 / source 2，参考 P:1511/1526）。
    """
    for polygon in polygons:
        points = _to_global_points(polygon, rect)
        _draw_outlined_polygon(frame, points, _DARK_ORANGE_RED, outline_thickness, color_thickness)


def _draw_reference_layer(frame: AnyArray, layers: OverlayLayers) -> None:
    """画 source 视图那个自动挑的背景圆（青，参考 P:1528-1543 画的是紫）。

    关键实现:
        仅当 ``reference_region.enabled`` 且 ``source_rect`` / ``center_xy`` /
        ``radius_px`` 齐备时绘制；圆心由 source 局部坐标加 ``source_rect`` 偏移抬到整帧。
        黑色外描边线宽 6、青色内线线宽 3（参考 P:1542-1543）。
    """
    region = layers.reference_region
    if region is None or layers.source_rect is None or not region.get("enabled", False):
        return
    center = region.get("center_xy")
    radius = region.get("radius_px")
    if center is None or radius is None:
        return
    center_global = (
        round(float(center[0]) + layers.source_rect.x1),
        round(float(center[1]) + layers.source_rect.y1),
    )
    radius_int = round(float(radius))
    cv2.circle(frame, center_global, radius_int, color=_OUTLINE_BLACK, thickness=6)
    cv2.circle(frame, center_global, radius_int, color=_BACKGROUND_CIRCLE_CYAN, thickness=3)


def draw_polygon_overlay(
    frame: AnyArray,
    layers: OverlayLayers,
    cfg: OverlayConfig,
) -> AnyArray:
    """在整帧坐标系上叠加 ROI / 暗区 / reference 各层，返回画好的新帧（不改入参）。

    参数:
        frame: 整帧 BGR 图像 ``(H, W, 3)``。
        layers: 单帧各层几何（局部坐标多边形 + 视图矩形）。
        cfg: 各画层开关与填充配置。

    返回:
        叠加后的整帧副本（尺寸 / dtype 与入参一致）。

    关键实现:
        先 copy 帧避免污染入参；按参考 P:1469-1544 的层序绘制——source ROI（黄）→
        target ROI 填充 + 轮廓（绿）→ target 暗区（橙红）→ source 暗区（橙红）→
        背景圆（青）。各层的局部→整帧坐标变换在 ``_to_global_points`` 内统一完成。
        填充层（``fill_polygon``）须先于轮廓绘制，否则填充会盖住轮廓。
    """
    out = frame.copy()

    # Phase 1: source ROI 轮廓（黄）。
    if cfg.draw_source_roi:
        _draw_source_roi_layer(out, layers)

    # Phase 2: target ROI——先半透明填充，再画绿色轮廓（主产物）。
    target_points = _to_global_points(layers.target_polygon, layers.target_rect)
    if cfg.fill_polygon:
        out = _fill_target_polygon(out, target_points)
    if cfg.draw_target_polygon:
        _draw_outlined_polygon(out, target_points, _TARGET_ROI_GREEN, 7, 3)

    # Phase 3: target 暗区轮廓（橙红）。
    if cfg.draw_dark_region and layers.dark_polygons:
        _draw_dark_layer(out, layers.dark_polygons, layers.target_rect, 6, 3)

    # Phase 4: source 暗区轮廓（橙红）。
    if (
        cfg.draw_source_dark_region
        and layers.source_dark_polygons
        and layers.source_rect is not None
    ):
        _draw_dark_layer(out, layers.source_dark_polygons, layers.source_rect, 5, 2)

    # Phase 5: 背景圆（青）。
    if cfg.draw_reference_region:
        _draw_reference_layer(out, layers)

    return out


def render_overlay(
    frames: Iterable[AnyArray],
    per_frame_layers: Sequence[OverlayLayers],
    fps: float,
    cfg: OverlayConfig,
    *,
    frame_size: tuple[int, int],
) -> bytes:
    """逐帧渲染 overlay 并编码为 mp4 字节（fourcc mp4v、原帧尺寸、调用方 fps）。

    参数:
        frames: 已解码的逐帧整帧 BGR 图像，非空且尺寸一致。**逐帧消费、不留**，
            所以可以直接喂一个解码器的迭代器。
        per_frame_layers: 与 ``frames`` 一一对应的逐帧图层；长度须等于帧数。
        fps: 视频帧率，**必须 > 0**（调用方已探测）；非正直接报错，不静默回退 25.0。
        cfg: 各画层开关与填充配置。
        frame_size: 输出视频的 ``(width, height)``。**由调用方给，不再从首帧读**——
            编码器要在写第一帧之前就建好，而流式之后「先看一眼首帧再决定尺寸」意味着
            要么把首帧留住、要么把迭代器消费掉一帧。调用方本来就从探测拿到了这个尺寸。

    返回:
        mp4 视频字节（mp4v 编码），由调用方写文件。

    异常:
        ValueError: ``frames`` 为空、``frames`` 与 ``per_frame_layers`` 长度不一致
            （由 ``zip`` 的 ``strict`` 执行——流式之后帧数要遍历完才知道）、
            ``fps`` 非正，或 ``frame_size`` 非正。
        RuntimeError: ``cv2.VideoWriter`` 无法打开（编码器不可用 / 尺寸非法）。

    关键实现:
        ``cv2.VideoWriter`` 不能直写内存，故写到临时 ``.mp4`` 文件、读回字节、``finally``
        ``unlink`` 删除（仿 ``roi_inference/frames.py`` 临时文件纪律，确保异常路径也清理）；
        本模块不落本地、不碰对象存储。逐帧 ``draw_polygon_overlay`` 保持帧尺寸不变，
        故无需逐帧重校尺寸。
    """
    # Phase 1: 防御校验（不静默兜底，缺陷直接报错）。
    if float(fps) <= 0.0:
        raise ValueError(f"fps 必须为正（调用方应已探测），得到 {fps}。")
    width, height = frame_size
    if width <= 0 or height <= 0:
        raise ValueError(f"frame_size 必须为正，得到 {frame_size}。")
    if not per_frame_layers:
        raise ValueError("per_frame_layers 不能为空：一帧都不写产不出视频。")

    # Phase 2: 写临时 mp4，结束（含异常）后必删。
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as handle:
        temp_path = Path(handle.name)
    try:
        writer = cv2.VideoWriter(
            str(temp_path),
            cv2.VideoWriter.fourcc(*"mp4v"),
            float(fps),
            (width, height),
        )
        if not writer.isOpened():
            raise RuntimeError(f"cv2.VideoWriter 无法打开（尺寸 {width}x{height}、fps {fps}）。")
        try:
            # strict：帧数要遍历完才知道，长度不一致在这里当场抛，不静默按较短的截断。
            for index, (frame, layers) in enumerate(zip(frames, per_frame_layers, strict=True)):
                if index == 0 and frame.shape[:2] != (height, width):
                    # **尺寸不符的帧会被 VideoWriter 静默丢掉**：实测写 6 帧读回 3 帧,
                    # 只在 stderr 留一行警告，不抛异常也不改返回值。结果是 overlay 视频
                    # 短一截甚至是空的，而它照样有字节、照样被写成文件。
                    # 尺寸以前由首帧决定所以结构上不可能不符，现在由调用方给，这条得自己守。
                    raise ValueError(
                        f"帧尺寸 {frame.shape[:2]} 与 frame_size {(width, height)} 不符；"
                        "尺寸不符的帧会被编码器静默丢掉"
                    )
                writer.write(draw_polygon_overlay(frame, layers, cfg))
        finally:
            writer.release()

        return temp_path.read_bytes()
    finally:
        temp_path.unlink(missing_ok=True)