CADx-4/tests/test_ceus_prompt_frames.py

"""B4 关键帧选择单元测试。

钉死 findings §3.3 / 参考 P:335-372 的关键帧语义：首帧强制（change_iou=1.0）、
形状不变不新增、``change_iou < iou_threshold``（严格小于）才记新 prompt_frame、
``min_gap_frames`` 在 IoU 判定之前抑制过密、无绿帧跳过、返回结构字段齐全。

测试用真实算法路径（``mask_iou``）而非 mock：构造形状渐变的布尔 mask 绿框观测序列，
让 IoU 在阈值上下精确落点（含恰好等于阈值的边界帧）。
"""

import numpy as np

from cadx.ceus import (
    CeusConfig,
    ColorPrompt,
    KeyframesConfig,
    PromptFrame,
    PromptObservation,
    select_prompt_frames,
)
from conftest import load_ceus_section


def _keyframes_cfg(**overrides: object) -> KeyframesConfig:
    """提示帧配置：config/ceus.yaml 的取值（旧测试用的是同值的默认值），按需改某一项。"""
    config = CeusConfig.model_validate(load_ceus_section())
    return config.keyframes.model_copy(update=overrides)


def _square_prompt(side: int, canvas: int = 100) -> ColorPrompt:
    """构造一个以原点为左上角、边长 ``side`` 的实心方块绿框 ``ColorPrompt``。

    mask 为 ``(canvas, canvas)`` 布尔图，``[0:side, 0:side]`` 为 True；
    其余几何字段按方块几何自洽填充，仅 mask 参与关键帧 IoU 判定。
    """
    mask = np.zeros((canvas, canvas), dtype=bool)
    mask[0:side, 0:side] = True
    contour = np.array(
        [[[0, 0]], [[side - 1, 0]], [[side - 1, side - 1]], [[0, side - 1]]],
        dtype=np.int32,
    )
    return ColorPrompt(
        mask=mask,
        contour=contour,
        bbox_xyxy=(0, 0, side - 1, side - 1),
        bbox_area=side * side,
        contour_area=float((side - 1) * (side - 1)),
        candidate_count=1,
    )


def _obs(frame_index: int, side: int | None) -> PromptObservation:
    """构造单帧观测：``side`` 为 None 表示无绿帧（prompt=None）。"""
    prompt = None if side is None else _square_prompt(side)
    return PromptObservation(frame_index=frame_index, prompt=prompt)


def test_first_detected_frame_is_forced_prompt_frame_with_change_iou_one() -> None:
    """首个有绿框的观测必为 prompt_frame，change_iou 恒记 1.0。"""
    observations = [_obs(0, 40)]
    prompt_frames = select_prompt_frames(observations, _keyframes_cfg())
    assert len(prompt_frames) == 1
    assert prompt_frames[0].frame_index == 0
    assert prompt_frames[0].change_iou == 1.0


def test_unchanged_shape_does_not_add_prompt_frame() -> None:
    """绿框形状完全不变（IoU=1.0 ≥ 阈值）的连续帧不新增 prompt_frame。"""
    observations = [_obs(i, 40) for i in range(5)]
    prompt_frames = select_prompt_frames(observations, _keyframes_cfg())
    assert len(prompt_frames) == 1
    assert prompt_frames[0].frame_index == 0


def test_iou_below_threshold_records_new_prompt_frame() -> None:
    """IoU 跌破阈值的帧记为新 prompt_frame，change_iou 等于实测 IoU。

    边长 40→20 的方块：IoU = 20² / 40² = 0.25 < 0.92。
    """
    observations = [_obs(0, 40), _obs(1, 20)]
    prompt_frames = select_prompt_frames(observations, _keyframes_cfg())
    assert [kf.frame_index for kf in prompt_frames] == [0, 1]
    assert prompt_frames[1].change_iou == 0.25


def test_iou_equal_to_threshold_does_not_record_boundary() -> None:
    """恰好 IoU == 阈值不记新 prompt_frame（严格 ``<``，钉死 < 非 <=）。

    取 iou_threshold = IoU(40, 20) = 0.25。第二帧 IoU 恰等于 0.25，
    ``0.25 < 0.25`` 为 False，故不新增。
    """
    config = _keyframes_cfg(iou_threshold=0.25)
    observations = [_obs(0, 40), _obs(1, 20)]
    prompt_frames = select_prompt_frames(observations, config)
    assert [kf.frame_index for kf in prompt_frames] == [0]


def test_min_gap_frames_suppresses_dense_prompt_frames() -> None:
    """``min_gap_frames`` 在 IoU 判定前门控，抑制过密 prompt_frame。

    三帧 40→20→40：默认 min_gap=1 时帧 1、帧 2 都会被记（IoU 均 0.25<0.92）；
    设 min_gap=2 后，帧 1 距首帧仅 1 < 2 被跳过，帧 2 距首帧 2 不小于 2 才比较，
    其 IoU(40,40)=1.0 ≥ 0.92 故也不记，最终仅保留首帧。
    """
    observations = [_obs(0, 40), _obs(1, 20), _obs(2, 40)]

    dense = select_prompt_frames(observations, _keyframes_cfg())
    assert [kf.frame_index for kf in dense] == [0, 1, 2]

    sparse = select_prompt_frames(observations, _keyframes_cfg(min_gap_frames=2))
    assert [kf.frame_index for kf in sparse] == [0]


def test_no_green_frames_are_skipped() -> None:
    """无绿帧（prompt=None）跳过，不占用首帧强制位也不影响 frame_index 对齐。"""
    observations = [_obs(0, None), _obs(1, None), _obs(2, 40), _obs(3, 20)]
    prompt_frames = select_prompt_frames(observations, _keyframes_cfg())
    assert [kf.frame_index for kf in prompt_frames] == [2, 3]
    # 首个有绿帧（frame 2）才是强制 prompt_frame，change_iou=1.0
    assert prompt_frames[0].change_iou == 1.0


def test_prompt_frame_result_structure_fields() -> None:
    """返回项含 frame_index / bbox_xyxy / change_iou，并保留几何溯源字段。"""
    observations = [_obs(7, 40)]
    prompt_frames = select_prompt_frames(observations, _keyframes_cfg())
    prompt_frame = prompt_frames[0]
    assert isinstance(prompt_frame, PromptFrame)
    assert prompt_frame.frame_index == 7
    assert prompt_frame.bbox_xyxy == (0, 0, 39, 39)
    assert prompt_frame.change_iou == 1.0
    assert prompt_frame.bbox_area == 1600
    assert prompt_frame.mask.dtype == bool


def test_empty_observations_returns_empty() -> None:
    """空观测序列返回空 prompt_frame 列表（无绿帧可选）。"""
    assert select_prompt_frames([], _keyframes_cfg()) == []