pact/eval_l2_skill_description_intervention.py

#!/usr/bin/env python3
"""Skill-description intervention tests for L2 reasoning sensitivity.

The script runs paired L2 decisions on the same L1 belief states:

1. full: normal skill mechanics descriptions,
2. no_desc: skill names and cooldowns only,
3. swap_desc: skill names keep their cooldowns, but mechanics descriptions are
   deterministically swapped across skills in the same menu.

Primary intervention metrics after the 2026-07-09 metric reset:

- No-Description Decision Change: whether removing descriptions changes the
  full decision.
- Swapped-Description Decision Change: whether swapping descriptions changes
  the full decision.
- LLM Judge: reserved for the separate LLM-judge evaluator; this script records
  it as not-run rather than inventing a proxy.

The script also supports two L1-to-L2 input views:

- full: expose the current L1 handoff fields used by the L2 prompt.
- compact: expose only tracker state plus dp/dp_bin and front_cone. Angle,
  decision_zone, tactical_sector, behind, and other geometry fields are omitted
  from the L2 prompt.
"""

from __future__ import annotations

import argparse
import base64
import functools
import hashlib
import itertools
import json
import math
import os
import re
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
import urllib.request
from collections import Counter, defaultdict
from typing import Any, Dict, Iterable, List, Sequence, Tuple

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
sys.path.insert(0, os.path.dirname(__file__))
from layered_belief import CATEGORY  # noqa: E402
from eval_l2_decision import call_gemma_model, gemma_openai_payload  # noqa: E402
from model_identity import require_model_identity  # noqa: E402


ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
DEFAULT_BELIEFS = "results/vrising/beliefs_l1_main_official_test.jsonl"
DEFAULT_SKILL_LIB = "games/vrising/skill_library.json"
BELIEF_VIEWS = ("full", "compact")
COMPACT_GEOMETRY_MODES = ("distance_front", "distance_side")
MENU_POLICIES = ("natural_ready", "counterfactual_two")
INPUT_MODALITIES = ("text_only", "text_image")
DEFAULT_DATA_ROOT = "data/processed"
SAMPLE_STRATEGY_REVISION = "balanced-boss-fight-distance-v3"
PROMPT_SCHEMA_REVISION = "hp-free-v6"
FRAME_EXTRACTOR_REVISION = "ffmpeg-current-frame-v2-width-video-sha256"


def load_json(path: str) -> Any:
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def load_jsonl(path: str) -> List[Dict[str, Any]]:
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def write_jsonl(path: str, rows: Iterable[Dict[str, Any]]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    tmp = f"{path}.tmp.{os.getpid()}"
    with open(tmp, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def append_jsonl(path: str, row: Dict[str, Any]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n")
        f.flush()
        os.fsync(f.fileno())


def write_json_atomic(path: str, value: Dict[str, Any]) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(value, f, ensure_ascii=False, indent=2, sort_keys=True, allow_nan=False)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def commit_run_outputs(
    metrics_path: str,
    rows_path: str,
    rows: Iterable[Dict[str, Any]],
    summary: Dict[str, Any],
) -> None:
    """Publish canonical rows first and metrics last as the commit marker."""
    write_jsonl(rows_path, rows)
    write_json_atomic(metrics_path, summary)


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def file_identity(path: str, hash_dir: str) -> Dict[str, Any]:
    stat = os.stat(path)
    real = os.path.realpath(path)
    size, mtime_ns = int(stat.st_size), int(stat.st_mtime_ns)
    os.makedirs(hash_dir, exist_ok=True)
    key = hashlib.sha256(real.encode("utf-8")).hexdigest()
    sidecar = os.path.join(hash_dir, f"{key}.json")
    identity = {"path": real, "size": size, "mtime_ns": mtime_ns}
    try:
        cached = load_json(sidecar)
        if cached.get("identity") == identity and len(str(cached.get("sha256", ""))) == 64:
            return {**identity, "sha256": str(cached["sha256"])}
    except (FileNotFoundError, json.JSONDecodeError, OSError):
        pass
    value = _file_identity_cached(real, size, mtime_ns)["sha256"]
    write_json_atomic(sidecar, {"identity": identity, "sha256": value})
    return {**identity, "sha256": value}


@functools.lru_cache(maxsize=256)
def _file_identity_cached(path: str, size: int, mtime_ns: int) -> Dict[str, Any]:
    return {
        "path": path,
        "size": size,
        "mtime_ns": mtime_ns,
        "sha256": sha256_file(path),
    }


def selected_video_identities(
    data_root: str, rows: Sequence[Dict[str, Any]], hash_dir: str,
) -> Dict[str, Any]:
    identities: Dict[str, Any] = {}
    for boss, fight in sorted({(str(row["boss"]), int(row["fight"])) for row in rows}):
        rel = os.path.join(boss, f"video_fight{fight}.mp4")
        path = os.path.join(data_root, rel)
        identities[rel] = file_identity(path, hash_dir)
    return identities


def selected_policy_metadata(
    data_root: str, rows: Sequence[Dict[str, Any]], hash_dir: str,
) -> Dict[str, Any]:
    policy_files: Dict[str, Any] = {}
    selected_times: List[List[Any]] = []
    cache: Dict[Tuple[str, int], List[Dict[str, Any]]] = {}
    for row in rows:
        boss, fight, index = str(row["boss"]), int(row["fight"]), int(row["index"])
        pair = (boss, fight)
        rel = os.path.join(boss, f"policy_view_fight{fight}.jsonl")
        path = os.path.join(data_root, rel)
        if rel not in policy_files:
            policy_files[rel] = file_identity(path, hash_dir)
        if pair not in cache:
            cache[pair] = [
                item for item in load_jsonl(path) if item.get("action") != "death"
            ]
        if index < 0 or index >= len(cache[pair]):
            raise IndexError(f"selected belief index {index} is outside {rel} ({len(cache[pair])} rows)")
        policy_row = cache[pair][index]
        video_t = float(
            policy_row.get("video_t", (policy_row.get("obs") or {}).get("fight_time", 0.0)) or 0.0
        )
        selected_times.append([boss, fight, index, video_t])
    return {
        "mapping_revision": "filtered-non-death-row-index-to-video-t-v1",
        "policy_files": policy_files,
        "selected_frame_times": selected_times,
    }


def row_key(row: Dict[str, Any]) -> Tuple[Any, Any, Any]:
    return (row.get("boss"), row.get("fight"), row.get("index"))


def stable_rank(seed: int, *parts: Any) -> str:
    payload = json.dumps([seed, *parts], ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def rate(values: Sequence[bool]) -> float:
    return sum(values) / max(1, len(values))


def round4(value: float) -> float:
    return round(float(value), 4)


def parse_json_object(text: str) -> Tuple[Dict[str, Any], bool]:
    """Parse one decision object and report whether conservative repair was used."""
    text = str(text or "").strip()
    try:
        return json.loads(text), False
    except json.JSONDecodeError:
        pass
    match = re.search(r"\{.*\}", text, re.S)
    if match:
        return json.loads(match.group(0)), True
    # Gemma occasionally emits exactly the requested object fields but drops
    # only the outer braces. Accept that narrow case; do not guess malformed
    # values or off-menu skills.
    try:
        repaired = json.loads("{" + text + "}")
    except json.JSONDecodeError as exc:
        raise ValueError(f"no JSON object in output: {text[:200]!r}") from exc
    if not isinstance(repaired, dict) or "skill" not in repaired:
        raise ValueError(f"repaired output is not a decision object: {text[:200]!r}")
    return repaired, True


def normalize_decision(raw: str, menu: Sequence[str]) -> Dict[str, Any]:
    try:
        obj, repaired = parse_json_object(raw)
        skill = obj.get("skill")
        reason = str(obj.get("reason", ""))
        confidence = obj.get("confidence", 0.0)
        try:
            confidence = float(confidence)
        except (TypeError, ValueError):
            confidence = 0.0
        if not math.isfinite(confidence):
            confidence = 0.0
        return {
            "skill": skill if skill in menu else None,
            "raw_skill": skill,
            "reason": reason,
            "confidence": max(0.0, min(1.0, confidence)),
            "raw_output": raw,
            "json_valid": True,
            "json_repaired": repaired,
            "schema_valid": isinstance(skill, str) and isinstance(reason, str),
            "on_menu": skill in menu,
        }
    except Exception as exc:
        return {
            "skill": None,
            "raw_skill": None,
            "reason": "",
            "confidence": 0.0,
            "raw_output": raw,
            "parse_error": str(exc),
            "json_valid": False,
            "schema_valid": False,
            "on_menu": False,
        }


def nested_belief(row: Dict[str, Any]) -> Dict[str, Any]:
    if row.get("l1_belief"):
        return row["l1_belief"]
    belief = row.get("belief") or {}
    return {
        "schema_version": "flat_l1_belief_v1",
        "geometry": {
            "distance_value": belief.get("player_distance_value"),
            "distance_bin": belief.get("player_distance_bin"),
            "dp_bin": belief.get("dp_bin"),
            "angle_value": belief.get("player_angle_value"),
            "angle_bin": belief.get("player_angle_bin"),
            "relative_side": belief.get("player_relative_side", belief.get("player_angle_bin")),
            "angle_36bin": belief.get("angle_36bin"),
            "front_cone": belief.get("front_cone"),
            "decision_zone": belief.get("decision_zone"),
            "tactical_sector": belief.get("tactical_sector"),
            "behind": belief.get("behind"),
        },
        "boss_state": {
            "prev_skill": belief.get("prev_boss_skill"),
            "skill_phase": belief.get("skill_phase", "decision"),
            "skill_finished": belief.get("skill_finished", True),
        },
        "resource_state": {
            # hp_phase 已从 L2 全面移除（不进 belief、不进 prompt）。
            "cooldown_ready": belief.get("cooldown_ready", {}),
            "cooldown_seconds": belief.get("cooldown_seconds", {}),
        },
    }


def ready_menu(
    boss: str,
    row: Dict[str, Any],
    belief: Dict[str, Any],
    skill_lib: Dict[str, Any],
    keep_prev_nonmelee: bool,
    menu_policy: str = "natural_ready",
) -> List[str]:
    legal = list(row.get("legal_skills") or sorted(skill_lib.get(boss, {})))
    if menu_policy == "counterfactual_two":
        pairs = list(itertools.combinations(sorted(legal), 2))
        if not pairs:
            return sorted(legal)
        rank = int(stable_rank(0, "counterfactual_two", *row_key(row)), 16)
        return list(pairs[rank % len(pairs)])
    if menu_policy != "natural_ready":
        raise ValueError(f"unknown menu_policy: {menu_policy}")
    resource = belief.get("resource_state", {}) or {}
    cooldown_ready = resource.get("cooldown_ready", {}) or {}
    prev = (belief.get("boss_state", {}) or {}).get("prev_skill")
    melee = {
        s for s in legal
        if (skill_lib.get(boss, {}).get(s, {}) or {}).get("function") == "melee"
    }
    menu = [
        s for s in legal
        if cooldown_ready.get(s, True) and (keep_prev_nonmelee or s != prev or s in melee)
    ]
    if not menu:
        menu = [s for s in legal if cooldown_ready.get(s, True)] or legal
    return sorted(menu)


def choose_rows(
    rows: List[Dict[str, Any]], limit: int | None, strategy: str, seed: int = 0
) -> List[Dict[str, Any]]:
    if limit is None or limit >= len(rows):
        return rows
    if strategy == "first":
        return rows[:limit]
    # Stratify by fight as well as boss and distance bin. Without fight in the key,
    # every bucket is drawn from the earliest fight (rows arrive fight-ordered), so
    # the whole sample collapses onto one fight per boss and adjacent frames within
    # it -- far fewer effective samples than the row count suggests.
    buckets: Dict[Tuple[str, str, str], List[Dict[str, Any]]] = defaultdict(list)
    for row in rows:
        belief = nested_belief(row)
        geom = belief.get("geometry", {}) or {}
        buckets[(
            row.get("boss", ""),
            str(row.get("fight", "na")),
            str(geom.get("dp_bin") or geom.get("distance_bin") or "na"),
        )].append(row)
    # Hash order prevents the old temporal bias (bucket.pop(0) always selected the
    # earliest adjacent frames). First round-robin within each boss across fights/
    # distance strata, then round-robin across bosses. Therefore every prefix is
    # boss-balanced and an n=60 image sample is an exact prefix of the n=200 text sample.
    by_boss: Dict[str, List[Dict[str, Any]]] = {}
    bosses = sorted({key[0] for key in buckets})
    for boss in bosses:
        boss_keys = [key for key in buckets if key[0] == boss]
        boss_keys.sort(key=lambda key: stable_rank(seed, "stratum", *key))
        for key in boss_keys:
            buckets[key].sort(key=lambda row: stable_rank(seed, "row", *row_key(row)))
        sequence: List[Dict[str, Any]] = []
        active = list(boss_keys)
        while active:
            remaining = []
            for key in active:
                if buckets[key]:
                    sequence.append(buckets[key].pop(0))
                if buckets[key]:
                    remaining.append(key)
            active = remaining
        by_boss[boss] = sequence

    selected: List[Dict[str, Any]] = []
    positions = {boss: 0 for boss in bosses}
    while len(selected) < limit:
        progressed = False
        for boss in bosses:
            pos = positions[boss]
            if pos < len(by_boss[boss]) and len(selected) < limit:
                selected.append(by_boss[boss][pos])
                positions[boss] += 1
                progressed = True
        if not progressed:
            break
    return selected


def build_resume_manifest(
    args: argparse.Namespace,
    beliefs_path: str,
    skill_lib_path: str,
    selected_rows: List[Dict[str, Any]],
    server_identity: Dict[str, Any] | None,
) -> Dict[str, Any]:
    selected_keys = [row_key(row) for row in selected_rows]
    data_root = os.path.realpath(args.data_root if os.path.isabs(args.data_root) else os.path.join(ROOT, args.data_root))
    frame_dir = os.path.realpath(args.frame_dir if os.path.isabs(args.frame_dir) else os.path.join(ROOT, args.frame_dir))
    consumes_frames = args.input_modality == "text_image" and not args.dry_run_prompts
    manifest: Dict[str, Any] = {
        "resume_schema_revision": 2,
        "prompt_schema_revision": PROMPT_SCHEMA_REVISION,
        "sample_strategy_revision": SAMPLE_STRATEGY_REVISION,
        "beliefs_sha256": sha256_file(beliefs_path),
        "skill_library_sha256": sha256_file(skill_lib_path),
        "selected_keys_sha256": hashlib.sha256(
            json.dumps(selected_keys, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        ).hexdigest(),
        "selected_n": len(selected_rows),
        "backend": args.backend,
        "requested_model": os.environ.get("GEMMA_MODEL") if args.backend == "gemma" else args.backend,
        "server_identity": server_identity,
        "endpoint_used": os.environ.get("GEMMA_OPENAI_BASE_URL") if args.backend == "gemma" else None,
        "temperature": float(os.environ.get("GEMMA_TEMPERATURE", "0")) if args.backend == "gemma" else None,
        "belief_view": args.belief_view,
        "compact_geometry_mode": getattr(args, "compact_geometry_mode", "distance_front"),
        "menu_policy": getattr(args, "menu_policy", "natural_ready"),
        "input_modality": args.input_modality,
        "setting_id": args.setting_id,
        "limit": args.limit,
        "sample_strategy": args.sample_strategy,
        "sample_seed": args.sample_seed,
        "anonymous_names": args.anonymous_names,
        "keep_prev_nonmelee": args.keep_prev_nonmelee,
        "frame_width": args.frame_width if args.input_modality == "text_image" else None,
        "dry_run_prompts": bool(args.dry_run_prompts),
        "data_root": data_root if consumes_frames else None,
        "source_videos": selected_video_identities(
            data_root, selected_rows, os.path.join(frame_dir, "_source_hashes")
        ) if consumes_frames else None,
        "policy_view_metadata": selected_policy_metadata(
            data_root, selected_rows, os.path.join(frame_dir, "_source_hashes")
        ) if consumes_frames else None,
        "frame_extractor_revision": FRAME_EXTRACTOR_REVISION if consumes_frames else None,
    }
    manifest["signature"] = hashlib.sha256(
        json.dumps(manifest, sort_keys=True, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    return manifest


def swap_mapping(menu: Sequence[str]) -> Dict[str, str]:
    """Return displayed_skill -> description_source_skill."""
    menu = sorted(menu)
    n = len(menu)
    if n <= 1:
        return {s: s for s in menu}
    best_offset = 1
    best_score = -1
    for offset in range(1, n):
        score = 0
        for i, skill in enumerate(menu):
            src = menu[(i + offset) % n]
            score += int(CATEGORY.get(skill) != CATEGORY.get(src))
        if score > best_score:
            best_score = score
            best_offset = offset
    return {skill: menu[(i + best_offset) % n] for i, skill in enumerate(menu)}


def card_for_skill(
    boss: str,
    display_skill: str,
    desc_source_skill: str,
    belief: Dict[str, Any],
    skill_lib: Dict[str, Any],
    variant: str,
    force_ready: bool = False,
) -> Dict[str, Any]:
    resource = belief.get("resource_state", {}) or {}
    seconds = resource.get("cooldown_seconds", {}) or {}
    lib = skill_lib.get(boss, {})
    own = lib.get(display_skill, {}) or {}
    src = lib.get(desc_source_skill, {}) or {}
    cooldown = own.get("cooldown_s_est")
    seconds_since = seconds.get(display_skill)
    if force_ready and cooldown is not None:
        seconds_since = max(float(seconds_since or 0.0), float(cooldown))
    card = {
        "skill": display_skill,
        "cooldown_s": cooldown,
        "seconds_since_last_use": seconds_since,
    }
    if variant != "no_desc":
        card.update({
            "function": src.get("function"),
            "effect_on_distance": src.get("effect_on_distance"),
            "role": src.get("role"),
        })
    return card


def distance_text(geom: Dict[str, Any]) -> Any:
    dp = geom.get("dp_bin") or geom.get("distance_bin")
    return {
        "0-2": "very close (point-blank)",
        "2-4": "close",
        "4-6": "far",
        "6+": "very far",
    }.get(dp, dp)


def tracker_payload(
    resource: Dict[str, Any],
    boss_state: Dict[str, Any],
    shown: Dict[str, str],
    anonymous: bool = True,
) -> Dict[str, Any]:
    # hp_phase 已从 L2 移除：L1 无 hp 训练头、闭环无来源，且它不进决策/菜单逻辑（见 CHECK_RECORD）。
    # previous_skill 一致匿名：在菜单内 -> 其 skill_XX；不在菜单内且匿名 -> 中性占位，
    # 绝不回退真名（旧代码 shown.get(prev, prev) 会在 prev 被排出菜单时泄露真技能名）。
    prev_skill = boss_state.get("prev_skill")
    if prev_skill is None:
        prev_display = None
    elif prev_skill in shown:
        prev_display = shown[prev_skill]
    else:
        prev_display = "a_previous_skill_not_in_menu" if anonymous else prev_skill
    return {
        "previous_skill": prev_display,
        "skill_phase": boss_state.get("skill_phase", "decision"),
        "skill_finished": boss_state.get("skill_finished", True),
    }


def situation_payload(
    boss: str,
    geom: Dict[str, Any],
    resource: Dict[str, Any],
    boss_state: Dict[str, Any],
    shown: Dict[str, str],
    belief_view: str,
    anonymous: bool = True,
    compact_geometry_mode: str = "distance_front",
) -> Dict[str, Any]:
    if belief_view not in BELIEF_VIEWS:
        raise ValueError(f"unknown belief_view: {belief_view}")
    base = {
        "boss": boss,
        "player_distance": distance_text(geom),
        "player_distance_value": geom.get("distance_value"),
        "player_distance_bin": geom.get("dp_bin") or geom.get("distance_bin"),
        "tracker": tracker_payload(resource, boss_state, shown, anonymous),
    }
    if belief_view == "compact":
        if compact_geometry_mode == "distance_front":
            base["player_in_front_cone"] = geom.get("front_cone")
        elif compact_geometry_mode == "distance_side":
            base["player_relative_side"] = geom.get("relative_side") or geom.get("angle_bin")
        else:
            raise ValueError(f"unknown compact_geometry_mode: {compact_geometry_mode}")
        return base
    base["player_in_front_cone"] = geom.get("front_cone")
    base.update(
        {
            "player_angle_value": geom.get("angle_value"),
            "player_angle_bin": geom.get("angle_bin"),
            "angle_36bin": geom.get("angle_36bin"),
            "decision_zone": geom.get("decision_zone"),
            "tactical_sector": geom.get("tactical_sector"),
            "behind": geom.get("behind"),
        }
    )
    return base


def build_prompt(
    row: Dict[str, Any],
    belief: Dict[str, Any],
    skill_lib: Dict[str, Any],
    variant: str,
    mapping: Dict[str, str],
    anonymous_names: bool,
    belief_view: str,
    compact_geometry_mode: str = "distance_front",
    menu_policy: str = "natural_ready",
) -> Tuple[str, Dict[str, str], Dict[str, str]]:
    boss = row["boss"]
    menu = sorted(mapping)
    geom = belief.get("geometry", {}) or {}
    resource = belief.get("resource_state", {}) or {}
    boss_state = belief.get("boss_state", {}) or {}

    if anonymous_names:
        shown = {skill: f"skill_{i:02d}" for i, skill in enumerate(menu)}
        hidden_to_real = {v: k for k, v in shown.items()}
    else:
        shown = {skill: skill for skill in menu}
        hidden_to_real = {skill: skill for skill in menu}

    cards = []
    for skill in menu:
        card = card_for_skill(
            boss, skill, mapping[skill], belief, skill_lib, variant,
            force_ready=menu_policy == "counterfactual_two",
        )
        card["skill"] = shown[skill]
        cards.append(card)

    payload = {
        "task": (
            "Choose the next boss skill for an action game. Use the current situation, "
            "skill menu, mechanics, and cooldown fields. Return JSON only."
        ),
        "belief_view": belief_view,
        "situation": situation_payload(
            boss, geom, resource, boss_state, shown, belief_view, anonymous_names,
            compact_geometry_mode,
        ),
        "skill_library": cards,
        "cooldown_rule": (
            "A skill is usable only if seconds_since_last_use >= cooldown_s. Pick exactly one skill "
            "from skill_library whose mechanics fit the current situation."
        ),
        "output_format": {
            "skill": "one of the listed skill ids",
            "reason": "one sentence citing situation and mechanics",
            "confidence": "0..1",
        },
    }
    return json.dumps(payload, ensure_ascii=False), shown, hidden_to_real


def load_policy_rows(data_root: str, boss: str, fight: int) -> List[Dict[str, Any]]:
    path = os.path.join(data_root, boss, f"policy_view_fight{fight}.jsonl")
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                row = json.loads(line)
                if row.get("action") != "death":
                    rows.append(row)
    return rows


def policy_row_for(
    cache: Dict[Tuple[str, int], List[Dict[str, Any]]],
    data_root: str,
    row: Dict[str, Any],
) -> Dict[str, Any]:
    k = (str(row["boss"]), int(row["fight"]))
    if k not in cache:
        cache[k] = load_policy_rows(data_root, k[0], k[1])
    return cache[k][int(row["index"])]


def extract_frame(row: Dict[str, Any], policy_row: Dict[str, Any], data_root: str,
                  frame_dir: str, width: int) -> str:
    """Decode current RGB; cache key binds width, video content, and extractor revision."""
    boss, fight, idx = str(row["boss"]), int(row["fight"]), int(row["index"])
    video_t = float(
        policy_row.get("video_t", (policy_row.get("obs") or {}).get("fight_time", 0.0)) or 0.0
    )
    video = os.path.join(data_root, boss, f"video_fight{fight}.mp4")
    video_id = file_identity(video, os.path.join(frame_dir, "_source_hashes"))
    rev = hashlib.sha256(FRAME_EXTRACTOR_REVISION.encode("utf-8")).hexdigest()[:8]
    out = os.path.join(
        frame_dir,
        f"{boss}_fight{fight}_idx{idx}_t{video_t:.3f}_w{width}_{video_id['sha256'][:12]}_{rev}.jpg",
    )
    if os.path.exists(out):
        from PIL import Image
        with Image.open(out) as image:
            if image.width != width:
                raise RuntimeError(f"cached frame width mismatch: {out}: {image.width} != {width}")
        return out
    os.makedirs(frame_dir, exist_ok=True)
    tmp = f"{out}.tmp.{os.getpid()}.jpg"
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-ss", f"{video_t:.3f}",
         "-i", video, "-frames:v", "1", "-vf", f"scale={width}:-1", "-q:v", "3", "-y", tmp],
        check=True,
    )
    from PIL import Image
    with Image.open(tmp) as image:
        if image.width != width:
            raise RuntimeError(f"extracted frame width mismatch: {tmp}: {image.width} != {width}")
    os.replace(tmp, out)
    return out


def image_data_uri(path: str) -> str:
    with open(path, "rb") as f:
        return "data:image/jpeg;base64," + base64.b64encode(f.read()).decode("ascii")


def call_gemma_vision(prompt: str, image_path: str, timeout: float, max_tokens: int) -> str:
    base_url = os.environ.get("GEMMA_OPENAI_BASE_URL")
    if not base_url:
        raise RuntimeError("text_image modality requires GEMMA_OPENAI_BASE_URL")
    body = gemma_openai_payload([
        {"type": "text", "text": prompt},
        {"type": "image_url", "image_url": {"url": image_data_uri(image_path)}},
    ], max_tokens=max_tokens)
    effective_timeout = float(os.environ.get("GEMMA_TIMEOUT", str(timeout)))
    attempts = max(1, int(os.environ.get("GEMMA_RETRIES", "4")))
    last_exc: Exception | None = None
    for attempt in range(attempts):
        req = urllib.request.Request(
            base_url.rstrip("/") + "/chat/completions",
            data=json.dumps(body).encode("utf-8"),
            headers={"Content-Type": "application/json",
                     "Authorization": f"Bearer {os.environ.get('GEMMA_API_KEY', 'EMPTY')}"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(req, timeout=effective_timeout) as resp:
                obj = json.loads(resp.read().decode("utf-8"))
            return obj["choices"][0]["message"]["content"]
        except Exception as exc:  # timeout / reset / transient server failure
            last_exc = exc
            if attempt < attempts - 1:
                time.sleep(2.0 * (attempt + 1))
    raise RuntimeError(f"gemma vision endpoint failed after {attempts} attempts: {last_exc}")


def run_model(prompt: str, backend: str, sleep_s: float,
              image_path: str | None = None, timeout: float = 180.0,
              max_tokens: int = 256) -> str:
    if backend == "gemma":
        raw = (
            call_gemma_vision(prompt, image_path, timeout, max_tokens)
            if image_path else call_gemma_model(prompt)
        )
    elif backend == "echo":
        # Deterministic debug backend for validating the pipeline only.
        obj = json.loads(prompt)
        skill = obj["skill_library"][0]["skill"] if obj.get("skill_library") else None
        raw = json.dumps({"skill": skill, "reason": "debug echo backend", "confidence": 0.0})
    else:
        raise ValueError(backend)
    if sleep_s > 0:
        time.sleep(sleep_s)
    return raw


def evaluate_sample(
    row: Dict[str, Any],
    skill_lib: Dict[str, Any],
    backend: str,
    dry_run_prompts: bool,
    anonymous_names: bool,
    keep_prev_nonmelee: bool,
    belief_view: str,
    sleep_s: float,
    compact_geometry_mode: str = "distance_front",
    menu_policy: str = "natural_ready",
    image_path: str | None = None,
) -> Dict[str, Any]:
    belief = nested_belief(row)
    menu = ready_menu(
        row["boss"], row, belief, skill_lib, keep_prev_nonmelee, menu_policy
    )
    full_mapping = {s: s for s in menu}
    swapped = swap_mapping(menu)
    variants = {
        "full": full_mapping,
        "no_desc": full_mapping,
        "swap_desc": swapped,
    }
    decisions: Dict[str, Any] = {}
    prompts: Dict[str, Any] = {}
    aliases: Dict[str, Any] = {}
    for variant, mapping in variants.items():
        prompt, shown, hidden_to_real = build_prompt(
            row, belief, skill_lib, variant, mapping, anonymous_names, belief_view,
            compact_geometry_mode,
            menu_policy,
        )
        prompts[variant] = json.loads(prompt)
        aliases[variant] = {"shown": shown, "hidden_to_real": hidden_to_real}
        if dry_run_prompts:
            decisions[variant] = {
                "skill": None,
                "raw_skill": None,
                "reason": "",
                "confidence": 0.0,
                "json_valid": False,
                "schema_valid": False,
                "on_menu": False,
                "dry_run": True,
            }
            continue
        raw = run_model(prompt, backend, sleep_s, image_path=image_path)
        parsed = normalize_decision(raw, list(hidden_to_real))
        if parsed["skill"] is not None:
            parsed["skill"] = hidden_to_real[parsed["skill"]]
        if parsed["raw_skill"] in hidden_to_real:
            parsed["raw_skill_real"] = hidden_to_real[parsed["raw_skill"]]
        decisions[variant] = parsed

    full_skill = decisions["full"].get("skill")
    no_desc_skill = decisions["no_desc"].get("skill")
    swap_skill = decisions["swap_desc"].get("skill")
    full_valid = bool(decisions["full"].get("on_menu"))
    no_desc_valid = bool(decisions["no_desc"].get("on_menu"))
    swap_valid = bool(decisions["swap_desc"].get("on_menu"))
    # menu<2 时干预无定义：无描述/交换描述都不可能改变"唯一可选"的决策，
    # 恒等 swap_mapping 会让 follow 与 name-bias 同时恒真（灌水且互相矛盾）。
    # 修复后的正确 prev 会被排出菜单，Isaac 的 2 技能 boss 因此出现大量单技能菜单。
    menu_interventionable = len(menu) >= 2
    no_desc_pair_valid = full_valid and no_desc_valid and menu_interventionable
    swap_pair_valid = full_valid and swap_valid and menu_interventionable
    metrics = {
        "full_valid": full_valid,
        "no_desc_valid": no_desc_valid,
        "swap_desc_valid": swap_valid,
        "no_desc_pair_valid": no_desc_pair_valid,
        "swap_pair_valid": swap_pair_valid,
        "no_desc_decision_changed": bool(no_desc_pair_valid and no_desc_skill != full_skill),
        "no_desc_decision_same": bool(no_desc_pair_valid and no_desc_skill == full_skill),
        "swap_desc_decision_changed": bool(swap_pair_valid and swap_skill != full_skill),
        "swap_desc_follow": bool(swap_pair_valid and swapped.get(swap_skill) == full_skill),
        "swap_name_bias": bool(swap_pair_valid and swap_skill == full_skill),
    }
    out = {
        "boss": row.get("boss"),
        "fight": row.get("fight"),
        "index": row.get("index"),
        "belief_view": belief_view,
        "compact_geometry_mode": compact_geometry_mode,
        "menu_policy": menu_policy,
        "input_modality": "text_image" if image_path else "text_only",
        "frame_path": image_path,
        "target_skill": row.get("target_skill"),
        "menu": menu,
        "geometry": (belief.get("geometry", {}) or {}),
        "l2_input_situation": prompts["full"]["situation"] if prompts.get("full") else None,
        "swap_mapping": swapped,
        "decisions": decisions,
        "metrics": metrics,
    }
    if dry_run_prompts:
        out["prompts"] = prompts
        out["aliases"] = aliases
    return out


def summarize(rows: List[Dict[str, Any]], args: argparse.Namespace) -> Dict[str, Any]:
    # 聚合层强制 menu>=2（resume 复用行内旧 metrics，evaluate_sample 层的
    # menu_interventionable 修复对已存行不生效，这里兜底）：单技能菜单上
    # 抹除/交换干预无定义，恒等 swap_mapping 会让 follow 与 name-bias 同时恒真。
    metrics = [dict(r["metrics"], _menu_ok=len(r.get("menu") or []) >= 2) for r in rows]
    n_singleton = sum(1 for m in metrics if not m["_menu_ok"])
    no_desc_pairs = [m for m in metrics if m["no_desc_pair_valid"] and m["_menu_ok"]]
    swap_pairs = [m for m in metrics if m["swap_pair_valid"] and m["_menu_ok"]]
    no_desc_change = (
        round4(rate([m["no_desc_decision_changed"] for m in no_desc_pairs])) if no_desc_pairs else None
    )
    swap_desc_change = (
        round4(rate([m["swap_desc_decision_changed"] for m in swap_pairs])) if swap_pairs else None
    )
    no_desc_same = (
        round4(rate([m.get("no_desc_decision_same", False) for m in no_desc_pairs])) if no_desc_pairs else None
    )
    swap_desc_follow = (
        round4(rate([m.get("swap_desc_follow", False) for m in swap_pairs])) if swap_pairs else None
    )
    swap_name_bias = (
        round4(rate([m.get("swap_name_bias", False) for m in swap_pairs])) if swap_pairs else None
    )

    by_boss: Dict[str, Dict[str, Any]] = {}
    for boss in sorted({str(r.get("boss")) for r in rows}):
        boss_rows = [r for r in rows if str(r.get("boss")) == boss]
        # 与主指标同口径:单技能菜单不计入干预配对(否则 follow/nb 恒真灌水)
        boss_metrics = [r["metrics"] for r in boss_rows if len(r.get("menu") or []) >= 2]
        boss_no_desc = [m for m in boss_metrics if m["no_desc_pair_valid"]]
        boss_swap = [m for m in boss_metrics if m["swap_pair_valid"]]
        by_boss[boss] = {
            "n": len(boss_rows),
            "no_desc_pairs": len(boss_no_desc),
            "swap_pairs": len(boss_swap),
            "no_desc_decision_change_rate": (
                round4(rate([m["no_desc_decision_changed"] for m in boss_no_desc]))
                if boss_no_desc else None
            ),
            "swap_desc_decision_change_rate": (
                round4(rate([m["swap_desc_decision_changed"] for m in boss_swap]))
                if boss_swap else None
            ),
            "swap_desc_follow_rate": (
                round4(rate([m.get("swap_desc_follow", False) for m in boss_swap]))
                if boss_swap else None
            ),
            "swap_name_bias_rate": (
                round4(rate([m.get("swap_name_bias", False) for m in boss_swap]))
                if boss_swap else None
            ),
        }

    return {
        "protocol": "pact-eval-v1",
        "layer": "L2",
        "setting_id": args.setting_id,
        # v2 stratifies by fight as well as boss and distance bin; v1 collapsed onto
        # one fight per boss, so the two sample sets are not comparable.
        "sample_set_id": (
            f"L2-SAMPLE-{args.limit}-v3-balanced" if args.sample_strategy == "stratified" else None
        ),
        "sample_strategy_revision": SAMPLE_STRATEGY_REVISION,
        "sample_seed": args.sample_seed,
        "prompt_schema_revision": PROMPT_SCHEMA_REVISION,
        "input_modality": args.input_modality,
        "input_modality_detail": (
            "text_only_l1_belief_plus_skill_descriptions"
            if args.input_modality == "text_only"
            else "same_text_payload_plus_current_rgb_frame_image_url"
        ),
        "belief_view": args.belief_view,
        "compact_geometry_mode": getattr(args, "compact_geometry_mode", "distance_front"),
        "menu_policy": getattr(args, "menu_policy", "natural_ready"),
        "required_metrics": {
            "no_desc_decision_change_rate": no_desc_change,
            "swap_desc_decision_change_rate": swap_desc_change,
            "llm_judge_selected_acceptable_rate": None,
            "llm_judge_status": "not_run_by_this_script",
        },
        "no_desc_flip_rate": no_desc_change,
        "no_desc_same_rate": no_desc_same,
        "swap_flip_rate": swap_desc_change,
        "swap_desc_follow_rate": swap_desc_follow,
        "swap_name_bias_rate": swap_name_bias,
        "backend": args.backend,
        "model": (
            (args.server_identity or {}).get("model_id") if args.backend == "gemma" else args.backend
        ),
        "requested_model": os.environ.get("GEMMA_MODEL") if args.backend == "gemma" else None,
        "server_identity": args.server_identity,
        "resume_signature": (getattr(args, "resume_manifest", {}) or {}).get("signature"),
        "resume_schema_revision": (getattr(args, "resume_manifest", {}) or {}).get("resume_schema_revision"),
        "temperature": float(os.environ.get("GEMMA_TEMPERATURE", "0")) if args.backend == "gemma" else None,
        "gemma_endpoint": os.environ.get("GEMMA_OPENAI_BASE_URL") if args.backend == "gemma" else None,
        "beliefs": args.beliefs,
        "skill_library": args.skill_library,
        "limit": args.limit,
        "sample_strategy": args.sample_strategy,
        "anonymous_names": args.anonymous_names,
        "dry_run_prompts": args.dry_run_prompts,
        "n": len(rows),
        "sanity_checks": {
            "full_valid_rate": round4(rate([m["full_valid"] for m in metrics])),
            "no_desc_valid_rate": round4(rate([m["no_desc_valid"] for m in metrics])),
            "swap_desc_valid_rate": round4(rate([m["swap_desc_valid"] for m in metrics])),
            "no_desc_pair_n": len(no_desc_pairs),
            "swap_pair_n": len(swap_pairs),
            "menu_singleton_excluded": n_singleton,
        },
        "by_boss": by_boss,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--beliefs", default=DEFAULT_BELIEFS)
    ap.add_argument("--skill_library", default=DEFAULT_SKILL_LIB)
    ap.add_argument("--backend", choices=["gemma", "echo"], default="gemma")
    ap.add_argument("--limit", type=int, default=30)
    ap.add_argument("--sample_strategy", choices=["stratified", "first"], default="stratified")
    ap.add_argument("--sample_seed", type=int, default=0)
    ap.add_argument("--anonymous_names", action="store_true")
    ap.add_argument("--keep_prev_nonmelee", action="store_true")
    ap.add_argument("--belief_view", choices=BELIEF_VIEWS, default="full")
    ap.add_argument(
        "--compact_geometry_mode", choices=COMPACT_GEOMETRY_MODES,
        default="distance_front",
    )
    ap.add_argument("--menu_policy", choices=MENU_POLICIES, default="natural_ready")
    ap.add_argument("--input_modality", choices=INPUT_MODALITIES, default="text_only")
    ap.add_argument(
        "--setting_id",
        default="L2-SEMANTIC-INTERVENTION",
        help="Experiment block recorded in metrics; model ablations use L2-MODEL-ABLATION.",
    )
    ap.add_argument("--data_root", default=DEFAULT_DATA_ROOT)
    ap.add_argument("--frame_dir", default="out/eval_runs/l2_frames_cache")
    ap.add_argument("--frame_width", type=int, default=512)
    ap.add_argument("--dry_run_prompts", action="store_true")
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--sleep_s", type=float, default=0.0)
    ap.add_argument("--workers", type=int, default=1,
                    help="并发发起 LLM 请求的线程数；默认 1 = 原来的串行行为")
    ap.add_argument("--out", default=None)
    ap.add_argument("--rows_out", default=None)
    args = ap.parse_args()

    args.server_identity = None
    if args.backend == "gemma" and not args.dry_run_prompts:
        endpoint = os.environ.get("GEMMA_OPENAI_BASE_URL")
        requested = os.environ.get("GEMMA_MODEL", "gemma-4-E2B-it")
        if not endpoint:
            ap.error("Gemma runs require GEMMA_OPENAI_BASE_URL")
        args.server_identity = require_model_identity(endpoint, requested)

    tag = f"n{args.limit}" if args.limit is not None else "all"
    anon = "_anon" if args.anonymous_names else ""
    view = f"_{args.belief_view}"
    mod = "" if args.input_modality == "text_only" else "_img"
    dry = "_dry" if args.dry_run_prompts else ""
    out = args.out or f"out/layered/l2_skill_desc_intervention_{args.backend}{anon}{view}{mod}_{tag}{dry}.json"
    rows_out = args.rows_out or os.path.splitext(out)[0] + "_rows.jsonl"

    beliefs_path = args.beliefs if os.path.isabs(args.beliefs) else os.path.join(ROOT, args.beliefs)
    skill_lib_path = (
        args.skill_library if os.path.isabs(args.skill_library) else os.path.join(ROOT, args.skill_library)
    )
    skill_lib = load_json(skill_lib_path)
    rows = choose_rows(load_jsonl(beliefs_path), args.limit, args.sample_strategy, args.sample_seed)
    resume_manifest = build_resume_manifest(
        args, beliefs_path, skill_lib_path, rows, args.server_identity
    )
    args.resume_manifest = resume_manifest
    resume_manifest_path = rows_out + ".resume.json"

    results: List[Dict[str, Any]] = []
    done = set()
    if args.resume and os.path.exists(rows_out):
        if not os.path.exists(resume_manifest_path):
            raise RuntimeError(
                f"refusing legacy resume without {resume_manifest_path}; use a new output directory"
            )
        saved_manifest = load_json(resume_manifest_path)
        if saved_manifest.get("signature") != resume_manifest["signature"]:
            raise RuntimeError(
                "resume configuration/source/sample mismatch; use a new output directory instead of "
                "mixing rows from different runs"
            )
        allowed = {row_key(row) for row in rows}
        for result in load_jsonl(rows_out):
            key = row_key(result)
            if key not in allowed:
                raise RuntimeError(f"resume row {key!r} is outside the selected sample")
            if key in done:
                raise RuntimeError(f"duplicate resume row {key!r}")
            if result.get("belief_view") != args.belief_view:
                raise RuntimeError(f"resume row {key!r} has a different belief_view")
            if result.get("compact_geometry_mode", "distance_front") != args.compact_geometry_mode:
                raise RuntimeError(f"resume row {key!r} has a different compact_geometry_mode")
            if result.get("menu_policy", "natural_ready") != args.menu_policy:
                raise RuntimeError(f"resume row {key!r} has a different menu_policy")
            if result.get("input_modality") != args.input_modality:
                raise RuntimeError(f"resume row {key!r} has a different input_modality")
            results.append(result)
            done.add(key)
        print(f"resuming from {rows_out}: {len(results)} completed rows", file=sys.stderr)
    elif args.resume and os.path.exists(resume_manifest_path):
        saved_manifest = load_json(resume_manifest_path)
        if saved_manifest.get("signature") != resume_manifest["signature"]:
            raise RuntimeError("orphan resume manifest does not match this run")
    write_json_atomic(resume_manifest_path, resume_manifest)

    data_root = args.data_root if os.path.isabs(args.data_root) else os.path.join(ROOT, args.data_root)
    frame_dir = args.frame_dir if os.path.isabs(args.frame_dir) else os.path.join(ROOT, args.frame_dir)
    policy_cache: Dict[Tuple[str, int], List[Dict[str, Any]]] = {}

    pending = [row for row in rows if row_key(row) not in done]

    def evaluate_one(row: Dict[str, Any]) -> Dict[str, Any]:
        image_path = None
        if args.input_modality == "text_image" and not args.dry_run_prompts:
            with io_lock:
                pol = policy_row_for(policy_cache, data_root, row)
                image_path = extract_frame(row, pol, data_root, frame_dir, args.frame_width)
        return evaluate_sample(
            row,
            skill_lib,
            args.backend,
            args.dry_run_prompts,
            args.anonymous_names,
            args.keep_prev_nonmelee,
            args.belief_view,
            args.sleep_s,
            args.compact_geometry_mode,
            args.menu_policy,
            image_path=image_path,
        )

    def record(result: Dict[str, Any]) -> None:
        results.append(result)
        done.add(row_key(result))
        if args.resume:
            append_jsonl(rows_out, result)
        n = len(results)
        if n == 1 or n % 10 == 0 or n == len(rows):
            print(f"completed {n}/{len(rows)} paired states", file=sys.stderr)

    io_lock = threading.Lock()
    if args.workers > 1:
        # 每个 state 是 3 次独立的 LLM 调用；server 是多线程的，串行发请求会把 GPU 闲置。
        # 用 as_completed 而不是 map：map 按提交顺序回收，一个慢请求会挡住所有已完成的结果落盘。
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            futures = [pool.submit(evaluate_one, row) for row in pending]
            for fut in as_completed(futures):
                with io_lock:
                    record(fut.result())
    else:
        for row in pending:
            record(evaluate_one(row))

    # Persist canonical sample order, independent of worker completion order.
    result_by_key = {row_key(result): result for result in results}
    results = [result_by_key[row_key(row)] for row in rows]
    summary = summarize(results, args)
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    # Canonical rows and their order are part of the result. metrics.json is the
    # final commit marker consumed by wrappers/gates.
    commit_run_outputs(out, rows_out, results, summary)
    print(f"wrote {out}")
    print(f"wrote {rows_out}")


if __name__ == "__main__":
    main()