CADx-4/tests/fake_ceus_run.py

"""测试用假 CEUS 运行目录：布局照 `scripts/ceus_quant.py` 写出的（failed.csv、snapshot/、
`<视频号>/summary.json`、`<视频号>/keyframes/<槽位>.png`、`<视频号>/tic.png`、`<视频号>/overlay.mp4`），
数值全部编造，不调显卡。tic.png 和 overlay.mp4 只是带视频号的几个字节（核对包按字节拷贝，不打开）。

summary.json 故意带上 tic.json 不该有的字段（视频号、视频 SHA-256、分割缓存键、帧号、拟合参数、
四条曲线的拟合成败与 R²、失败原因代码、推注时间、紫圈检出率与是否可用、灌注指标的口径标签和辅助
时刻、关键帧状态与未定位原因），用来测 build 只按白名单挑字段。视频 SHA-256 和推注时间照
`fake_delivery` 的假交付包（视频内容 `VIDEO_BYTES`、推注 5 秒）。关键帧 PNG 带文本块和 EXIF。
"""

import csv
import hashlib
import json
from pathlib import Path
from typing import Any

import yaml
from PIL import Image, PngImagePlugin

from fake_delivery import REPO_ROOT, VIDEO_BYTES

SLOTS = ("B0", "D1", "D2", "D3", "D4", "F1", "F2")
CURVES = ("roi_all", "roi_perf", "roi_pan", "background_circle")
CEUS_YAML = REPO_ROOT / "config" / "ceus.yaml"
UNLOCATED_REASON = "拟合曲线上没有 D4 时刻（编造的原因）"


def perfusion(scale: float) -> dict[str, Any]:
    """一组灌注指标的 summary 段：11 个指标依次是 `scale` × 1, 2, …, 11。"""
    names = ("PE_au", "WiAUC_au_s", "RT_s", "mTTl_s", "TTP_s", "WiR_au_per_s", "WiPI_au")
    more = ("WoAUC_au_s", "WiWoAUC_au_s", "FT_s", "WoR_au_per_s")
    return {
        "signal_domain": "db_display",
        "time_origin": "video_start",
        "area_curve": "fitted_minus_baseline",
        "values": {name: scale * (i + 1) for i, name in enumerate((*names, *more))},
        "supporting_times_s": {"TI": 3.0, "mTT": 9.5, "TO": None},
    }


def _fit(success: bool, r2: float | None) -> dict[str, Any]:
    return {
        "success": success,
        "failure_reason": None if success else "credible_peak_not_found",
        "curve_model": "lognormal",
        "baseline_source": "observed",
        "valid_samples": 120,
        "duplicate_samples": 3,
        "parameters": {
            "area": 1.0,
            "mu": 2.0,
            "sigma": 0.5,
            "arrival_offset_s": 4.0,
            "baseline": 9,
        },
        "quality": {"r2": r2, "nrmse": 0.1, "residual_noise_ratio": 1.2, "multi_peak_flag": False},
        "model_arrival_time_video_s": 4.0,
        "model_landmark_times_video_s": {"D2": 6.0, "D3": 9.0, "D4": 15.0},
        "model_landmark_intensities": {"D2": 1.0, "D3": 2.0, "D4": 1.5},
    }


def summary(video_id: str, located: tuple[str, ...]) -> dict[str, Any]:
    """一例的 summary.json：`located` 里的槽位定位到（视频时间=槽位序号×2.5 秒），余下未定位。"""
    anchors = {
        slot: {
            "frame_index": index * 25 if slot in located else None,
            "frame_time_video_s": index * 2.5 if slot in located else None,
            "time_since_injection_s": index * 2.5 + 5.0 if slot in located else None,
            "status": "ok" if slot in located else "anchor_undefined",
            "reason": None if slot in located else UNLOCATED_REASON,
            "anchor_time_video_s": index * 2.5,
            "time_error_s": 0.0 if slot in located else None,
            "source_curve": "roi_perf",
        }
        for index, slot in enumerate(SLOTS)
    }
    return {
        "video_id": video_id,
        "video_sha256": hashlib.sha256(VIDEO_BYTES).hexdigest(),
        "segment_cache_key": "f" * 64,
        "injection_offset_s": 5.0,
        "reference_valid": False,
        "pan_detected_fraction": 0.4,
        "early_truncation_flag": False,
        "lesion_fit_ready": True,
        "fits": {
            "roi_all": _fit(True, 0.93),
            "roi_perf": _fit(True, 0.91),
            "roi_pan": _fit(False, None),
            "background_circle": _fit(False, 0.2),
        },
        # 紫圈拟合没过（reference_valid 为 False），整组为 null；病灶去暗区拟合过了，有指标。
        "perfusion_metrics": {"roi_perf": perfusion(1.0), "roi_pan": None},
        "anchor_frames": anchors,
        "available_dynamic_slots": [s for s in ("B0", "D1", "D2", "D3", "D4") if s in located],
    }


def keyframe_png(path: Path, shade: int) -> None:
    """一张 8 位 RGB 关键帧，带文本块和 EXIF（写进资料库前要去掉）。"""
    pixels = bytes(v for i in range(12) for v in (shade, i * 20, 200 - i * 10))
    image = Image.frombytes("RGB", (4, 3), pixels)
    info = PngImagePlugin.PngInfo()
    info.add_text("Source", "U999_roi2.mp4")
    exif = Image.Exif()
    exif[0x010E] = "U999"  # ImageDescription
    image.save(path, format="PNG", pnginfo=info, exif=exif.tobytes())


def write_run(
    run_dir: Path,
    delivery_dir: Path,
    results: dict[str, tuple[str, ...]],
    failed: list[tuple[str, int, str]],
) -> Path:
    """写一个已完成的运行目录：`delivery_dir` 记进快照里的 ceus.yaml，`results` 是视频号 → 定位到的
    槽位，`failed` 是 failed.csv 的行。"""
    (run_dir / "snapshot").mkdir(parents=True)
    (run_dir / "snapshot" / "snapshot.json").write_text("{}", encoding="utf-8")
    # 快照里的 ceus.yaml 照仓库里的那份，只把交付包路径换成这份假交付包。
    ceus_yaml = yaml.safe_load(CEUS_YAML.read_text(encoding="utf-8"))
    ceus_yaml["paths"]["delivery_dir"] = str(delivery_dir)
    (run_dir / "snapshot" / "ceus.yaml").write_text(
        yaml.safe_dump(ceus_yaml, allow_unicode=True), encoding="utf-8"
    )
    with (run_dir / "failed.csv").open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(("video_id", "step", "reason"))
        writer.writerows(failed)
    for video_id, located in results.items():
        (run_dir / video_id / "keyframes").mkdir(parents=True)
        content = json.dumps(summary(video_id, located), ensure_ascii=False, indent=2)
        (run_dir / video_id / "summary.json").write_text(content, encoding="utf-8")
        for shade, slot in enumerate(located):
            keyframe_png(run_dir / video_id / "keyframes" / f"{slot}.png", shade * 30)
        (run_dir / video_id / "tic.png").write_bytes(f"fake tic.png {video_id}".encode())
        (run_dir / video_id / "overlay.mp4").write_bytes(f"fake overlay.mp4 {video_id}".encode())
    return run_dir