#!/usr/bin/env python3
"""Stage 1 legacy PPO inference for Karman and Illusion data collection.

Outputs are isolated under ``data/runs/<run_id>/<scene_id>/<scene>/``.  The
module deliberately keeps LegacyCelerisLab imports inside runtime functions so
CLI validation and unit tests do not require CUDA.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shlex
import subprocess
import sys
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_ROOT = REPO_ROOT / "src"
for import_root in (REPO_ROOT, SRC_ROOT):
    if str(import_root) not in sys.path:
        sys.path.insert(0, str(import_root))

from SR_analysis.configs import (  # noqa: E402
    CONV_LEN,
    FIFO_LEN,
    LEGACY_CFG_DIR,
    get_scene,
    get_scene_list,
    model_path_for_scene,
)
from SR_analysis.utils.harmonics import (  # noqa: E402
    analyze_harmonics,
    gen_target_states_at,
    harmonics_document,
)

DATA_TYPE = np.float32
SCHEMA_VERSION = "sr-stage1-npz-v2"
HARMONICS_SCHEMA_VERSION = "sr-target-harmonics-v2"
METRIC_ID = "legacy_similarity_v1"
ILLUSION_TARGET_CHANNELS = [
    "target_cylinder_fx", "target_cylinder_fy",
    "sensor_top_ux", "sensor_top_uy",
    "sensor_center_ux", "sensor_center_uy",
    "sensor_bottom_ux", "sensor_bottom_uy",
]
SENSOR_CHANNELS = [
    "sensor_top_ux", "sensor_top_uy", "sensor_center_ux",
    "sensor_center_uy", "sensor_bottom_ux", "sensor_bottom_uy",
]
FORCE_CHANNELS = ["front_fx", "front_fy", "top_fx", "top_fy", "bottom_fx", "bottom_fy"]
ACTION_CHANNELS = ["front_alpha", "top_alpha", "bottom_alpha"]

SCENE_GROUPS = {
    "karman_all": lambda: get_scene_list("karman"),
    "karman_trained": lambda: [s for s in get_scene_list("karman") if get_scene(s).get("model_name")],
    "illusion_trained": lambda: ["illusion_0.75L", "illusion_1L", "illusion_1.5L"],
    "illusion_generalization": lambda: ["illusion_0.5L", "illusion_0.6L", "illusion_0.8L", "illusion_1.2L", "illusion_2L"],
}


def _jsonable(value: Any) -> Any:
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, tuple):
        return list(value)
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_jsonable(v) for v in value]
    if isinstance(value, np.generic):
        return value.item()
    return value


def _write_json(path: Path, value: Any) -> None:
    path.write_text(json.dumps(_jsonable(value), indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _sha256(path: Path | None) -> str | None:
    if path is None or not path.is_file():
        return None
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _git_sha() -> str | None:
    try:
        return subprocess.run(
            ["git", "-C", str(REPO_ROOT), "rev-parse", "HEAD"],
            check=True, capture_output=True, text=True,
        ).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return None


def default_run_id() -> str:
    return datetime.now(timezone.utc).strftime("legacy-%Y%m%dT%H%M%SZ")


def gpu_identity(logical_device: int, cuda_visible_devices: str | None) -> dict[str, str | None]:
    """Resolve identity only when visibility provides an explicit mapping."""
    visible_tokens = None if cuda_visible_devices is None else [
        token.strip() for token in cuda_visible_devices.split(",") if token.strip()
    ]
    if visible_tokens is not None and not 0 <= logical_device < len(visible_tokens):
        return {"uuid": None, "name": None}
    selector = visible_tokens[logical_device] if visible_tokens is not None else None
    try:
        command = ["nvidia-smi"]
        if selector is not None:
            command.append(f"--id={selector}")
        command.extend(["--query-gpu=uuid,name", "--format=csv,noheader"])
        output = subprocess.run(
            command, check=True, capture_output=True, text=True, timeout=5,
        ).stdout.strip().splitlines()
        if selector is None and len(output) != 1:
            return {"uuid": None, "name": None}
        if len(output) != 1:
            return {"uuid": None, "name": None}
        uuid, name = (part.strip() for part in output[0].split(",", 1))
        return {"uuid": uuid or None, "name": name or None}
    except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError):
        return {"uuid": None, "name": None}


def cfd_device_manifest(logical_device: int) -> dict[str, Any]:
    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
    identity = gpu_identity(logical_device, cuda_visible_devices)
    return {
        "cfd_device": logical_device,
        "cfd_device_logical": logical_device,
        "cuda_visible_devices": cuda_visible_devices,
        "cfd_gpu_uuid": identity["uuid"],
        "cfd_gpu_name": identity["name"],
    }


def resolve_output_root(output_root: str | os.PathLike[str] | None) -> Path:
    if output_root is None:
        return REPO_ROOT / "src" / "SR_analysis" / "data" / "runs"
    path = Path(output_root).expanduser()
    return path.resolve() if path.is_absolute() else (REPO_ROOT / path).resolve()


def scene_output_dir(output_root: Path, run_id: str, scene_name: str, cfg: Mapping[str, Any]) -> Path:
    if not run_id or run_id in {".", ".."} or Path(run_id).name != run_id:
        raise ValueError("--run-id must be one non-empty path component")
    return output_root / run_id / str(cfg["scene_id"]) / scene_name


def prepare_output_dir(path: Path, overwrite: bool) -> None:
    if path.exists() and any(path.iterdir()) and not overwrite:
        raise FileExistsError(f"Refusing to overwrite non-empty output directory: {path}")
    path.mkdir(parents=True, exist_ok=True)


def frozen_scene_dir(scene_name: str, cfg: Mapping[str, Any]) -> Path:
    return REPO_ROOT / "src" / "SR_analysis" / "data" / str(cfg["scene_id"]) / scene_name


def validate_norm(norm: Mapping[str, Any]) -> dict[str, Any]:
    required = {"force_norm_fact", "sens_deviation", "sens_norm_fact"}
    missing = required - set(norm)
    if missing:
        raise ValueError(f"norm missing keys: {sorted(missing)}")
    force = float(norm["force_norm_fact"])
    deviation = np.asarray(norm["sens_deviation"], dtype=np.float32)
    scale = np.asarray(norm["sens_norm_fact"], dtype=np.float32)
    if not np.isfinite(force) or force <= 0:
        raise ValueError("force_norm_fact must be finite and positive")
    if deviation.shape != (6,) or scale.shape != (6,):
        raise ValueError("sens_deviation and sens_norm_fact must each contain 6 values")
    if not np.all(np.isfinite(deviation)) or not np.all(np.isfinite(scale)) or np.any(scale <= 0):
        raise ValueError("sensor normalization values must be finite and scales positive")
    result = dict(norm)
    result.update(force_norm_fact=force, sens_deviation=deviation.tolist(), sens_norm_fact=scale.tolist())
    return result


def load_existing_norm(scene_name: str, cfg: Mapping[str, Any]) -> tuple[dict[str, Any], Path]:
    path = frozen_scene_dir(scene_name, cfg) / "norm.json"
    if not path.is_file():
        raise FileNotFoundError(f"Frozen normalization not found: {path}")
    return validate_norm(json.loads(path.read_text(encoding="utf-8"))), path


def normalize_raw_observation(raw: np.ndarray, norm: Mapping[str, Any]) -> np.ndarray:
    raw = np.asarray(raw, dtype=np.float32)
    forces = raw[6:12] / float(norm["force_norm_fact"])
    sensors = (raw[0:6] - np.asarray(norm["sens_deviation"], dtype=np.float32)) / np.asarray(norm["sens_norm_fact"], dtype=np.float32)
    return np.clip(np.hstack((forces, sensors)), -1.0, 1.0).astype(np.float32)


def extract_illusion_target(raw_obs: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return canonical 8 channels, force[0:2], and sensor[2:8] channels."""
    canonical = np.asarray(raw_obs, dtype=np.float32).reshape(-1)[:8]
    if canonical.shape != (8,):
        raise ValueError(f"Illusion target observation requires 8 channels, got {canonical.size}")
    return canonical, canonical[:2].copy(), canonical[2:8].copy()


def target_phase_index(step: int, *, after_action: bool = False) -> int:
    if step < 0:
        raise ValueError("step must be non-negative")
    return step + int(after_action)


def illusion_observation(raw: np.ndarray, norm: Mapping[str, Any], target_force: np.ndarray) -> np.ndarray:
    base = normalize_raw_observation(raw, norm)
    target = np.asarray(target_force, dtype=np.float32).reshape(-1)
    if target.shape != (2,):
        raise ValueError("target_force must contain fx and fy")
    return np.clip(np.hstack((base, target / float(norm["force_norm_fact"]))), -1.0, 1.0).astype(np.float32)


def _npz_metadata(scene_name: str, cfg: Mapping[str, Any], layout: Sequence[str], alignment: str) -> dict[str, np.ndarray]:
    dt_c = float(cfg["sample_interval"]) * float(cfg["u0"]) / 20.0
    return {
        "schema_version": np.array(SCHEMA_VERSION),
        "layout": np.asarray(layout),
        "units_system": np.array("LegacyCelerisLab lattice/SI-normalized metadata"),
        "sample_interval": np.array(int(cfg["sample_interval"]), dtype=np.int64),
        "SI": np.array(int(cfg["sample_interval"]), dtype=np.int64),
        "dt_c": np.array(dt_c, dtype=np.float64),
        "alignment": np.array(alignment),
        "scene": np.array(scene_name),
    }


def save_target_npz(path: Path, scene_name: str, cfg: Mapping[str, Any], target: np.ndarray, *, illusion: bool) -> None:
    target = np.asarray(target, dtype=np.float32)
    layout = ILLUSION_TARGET_CHANNELS if illusion else SENSOR_CHANNELS
    payload: dict[str, Any] = {"target_states": target, **_npz_metadata(scene_name, cfg, layout, "sample k is post-step k target state")}
    if illusion:
        payload["target_forces"] = target[:, :2]
        payload["target_sensors"] = target[:, 2:8]
    np.savez_compressed(path, **payload)


def save_controlled_npz(
    path: Path, scene_name: str, cfg: Mapping[str, Any], *, sensors: np.ndarray,
    forces: np.ndarray, actions_norm: np.ndarray, actions_commanded: np.ndarray,
    rewards: np.ndarray | None = None, target_forces: np.ndarray | None = None,
    phase_indices: np.ndarray | None = None,
) -> None:
    sensors = np.asarray(sensors, dtype=np.float32)
    forces = np.asarray(forces, dtype=np.float32)
    actions_norm = np.asarray(actions_norm, dtype=np.float32)
    payload: dict[str, Any] = {
        "sensors": sensors, "forces": forces, "actions": actions_norm,
        "sensors_post": sensors, "forces_post": forces,
        "actions_norm": actions_norm,
        "actions_alpha_commanded": np.asarray(actions_commanded, dtype=np.float32) / float(cfg["u0"]),
        "actions_omega_commanded": np.asarray(actions_commanded, dtype=np.float32),
        **_npz_metadata(scene_name, cfg, SENSOR_CHANNELS + FORCE_CHANNELS + ACTION_CHANNELS,
                        "action[k] produces sensors_post[k]/forces_post[k]"),
    }
    if rewards is not None:
        payload["rewards"] = np.asarray(rewards, dtype=np.float32)
    if target_forces is not None:
        target_forces = np.asarray(target_forces, dtype=np.float32)
        payload["target_forces_for_action"] = target_forces
        payload["target_forces"] = target_forces
    if phase_indices is not None:
        payload["target_phase_index"] = np.asarray(phase_indices, dtype=np.int64)
    np.savez_compressed(path, **payload)


def _runtime_imports():
    from LegacyCelerisLab import FlowField
    from SR_analysis.utils.cfd_interface import (
        add_pinball, build_karman_cloak_env, compute_similarity, load_legacy_configs,
        load_ppo_model, scale_action,
    )
    return FlowField, add_pinball, build_karman_cloak_env, compute_similarity, load_legacy_configs, load_ppo_model, scale_action


def _compute_norm(samples: np.ndarray) -> dict[str, Any]:
    samples = np.asarray(samples, dtype=np.float32)
    force_scale = 6.0 * float(np.max(np.abs(samples[:, 6:12])))
    deviation = np.mean(samples[:, 0:6], axis=0)
    sensor_scale = 5.0 * np.max(np.abs(samples[:, 0:6] - deviation), axis=0)
    return validate_norm({"force_norm_fact": force_scale, "sens_deviation": deviation, "sens_norm_fact": sensor_scale})


def infer_karman(
    scene_name: str, device_id: int, out_dir: Path, n_steps: int,
    norm_source: str, model_device: str = "cpu",
) -> dict[str, Any]:
    FlowField, add_pinball, build_env, similarity, load_cfg, load_model, scale_action = _runtime_imports()
    cfg = get_scene(scene_name)
    cuda_cfg, field_cfg = load_cfg(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target, _ = build_env(ff, u0=cfg["u0"], l0=20.0, sample_interval=cfg["sample_interval"], fifo_len=FIFO_LEN, data_type=DATA_TYPE)
    save_target_npz(out_dir / "target.npz", scene_name, cfg, target, illusion=False)
    computed_norm = add_pinball(ff, l0=20.0, u0=cfg["u0"], sample_interval=cfg["sample_interval"], fifo_len=FIFO_LEN,
                                data_type=DATA_TYPE, action_bias=cfg["action_bias"], pinball_front_x=cfg["pinball_front_x"],
                                pinball_rear_x=cfg["pinball_rear_x"], obs_slice_start=cfg["obs_slice"][0], obs_slice_end=cfg["obs_slice"][1])
    norm, norm_path = (load_existing_norm(scene_name, cfg) if norm_source == "existing" else (validate_norm(computed_norm), None))
    _write_json(out_dir / "norm.json", norm)

    ff.restore_ddf(); ff.apply_ddf()
    sensors_u, forces_u = [], []
    for _ in range(n_steps):
        ff.run(cfg["sample_interval"], np.zeros(cfg["n_objects_env"], dtype=DATA_TYPE))
        raw = ff.obs.copy()[2:14]; sensors_u.append(raw[:6]); forces_u.append(raw[6:12])
    np.savez_compressed(out_dir / "uncontrolled.npz", sensors=np.asarray(sensors_u, dtype=np.float32), forces=np.asarray(forces_u, dtype=np.float32),
                        **_npz_metadata(scene_name, cfg, SENSOR_CHANNELS + FORCE_CHANNELS, "all samples are post-step states"))

    model_path = model_path_for_scene(scene_name)
    result: dict[str, Any] = {"scene": scene_name, "controlled": False, "metric_id": METRIC_ID, "norm_source": norm_source,
                              "norm_source_path": str(norm_path) if norm_path else None}
    if model_path:
        model = load_model(model_path, device=model_device, s_dim=cfg.get("s_dim", 12))
        ff.restore_ddf(); ff.apply_ddf()
        fifo: deque[np.ndarray] = deque(
            (np.asarray(row, dtype=np.float32) for row in computed_norm["save_states"]),
            maxlen=FIFO_LEN,
        )
        obs = normalize_raw_observation(fifo[-1], norm)
        sensors, forces, actions, commanded, rewards = [], [], [], [], []
        for _ in range(n_steps):
            action, _ = model.predict(obs, deterministic=True); action = np.asarray(action, dtype=np.float32).flatten()
            alpha = (action * float(cfg["action_scale"]) + np.asarray(cfg["action_bias"], dtype=np.float32)) * float(cfg["u0"])
            omega = scale_action(action, scale=cfg["action_scale"], bias=cfg["action_bias"], u0=cfg["u0"], n_total_bodies=cfg["n_objects_env"])
            ff.context.push(); ff.run(cfg["sample_interval"], omega); ff.context.pop()
            raw = ff.obs.copy()[2:14]; fifo.append(raw); sensors.append(raw[:6]); forces.append(raw[6:12]); actions.append(action); commanded.append(alpha)
            obs = normalize_raw_observation(raw, norm)
            if len(fifo) >= CONV_LEN:
                sa = np.asarray(fifo, dtype=np.float32)
                cd = float(np.mean(sa[-1, [0, 2, 4]]) / norm["force_norm_fact"])
                cl = float(np.mean(sa[-1, [1, 3, 5]]) / norm["force_norm_fact"])
                sim = similarity(target, sa[:, :6], CONV_LEN)
                rewards.append(min(0.3*np.exp(-abs(cd*20))+0.4*np.exp(-abs(cl*80))+0.3*np.exp(-10*abs(sim-1)), 1.0))
        save_controlled_npz(out_dir / "controlled.npz", scene_name, cfg, sensors=np.asarray(sensors), forces=np.asarray(forces),
                            actions_norm=np.asarray(actions), actions_commanded=np.asarray(commanded), rewards=np.asarray(rewards))
        result.update(controlled=True, similarity=float(similarity(target, np.asarray(sensors), CONV_LEN)))
        if rewards: result["avg_reward_last100"] = float(np.mean(rewards[-100:]))
    del ff
    return result


def infer_illusion(
    scene_name: str, device_id: int, out_dir: Path, n_steps: int,
    target_only: bool, norm_source: str, model_device: str = "cpu",
) -> dict[str, Any]:
    FlowField, _, _, similarity, load_cfg, load_model, _ = _runtime_imports()
    cfg = get_scene(scene_name); l0 = 20.0; u0 = float(cfg["u0"]); si = int(cfg["sample_interval"])
    cuda_cfg, field_cfg = load_cfg(LEGACY_CFG_DIR); field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id); cy = (ff.FIELD_SHAPE[1] - 1) / 2
    # Legacy add_cylinder's final argument is radius; target_diameter is expressed in L0 units historically.
    ff.add_cylinder((20.0 * l0, cy, 0.0), float(cfg["target_diameter"]) * l0)
    sensor_y = [cy + 2*l0, cy, cy - 2*l0]
    for y in sensor_y: ff.add_sensor((float(cfg["sensor_x"]) * l0, y, 0.0), 5.0)
    warmup = int(4 * ff.FIELD_SHAPE[0] / u0); ff.run(warmup, np.zeros(4, dtype=DATA_TYPE))
    rows = []
    for _ in range(FIFO_LEN):
        ff.run(si, np.zeros(4, dtype=DATA_TYPE)); rows.append(extract_illusion_target(ff.obs.copy())[0])
    target = np.asarray(rows, dtype=np.float32)
    save_target_npz(out_dir / "target.npz", scene_name, cfg, target, illusion=True)
    harmonics = analyze_harmonics(target, n_harmonics=5)
    _write_json(out_dir / "target_harmonics.json", harmonics_document(harmonics, ILLUSION_TARGET_CHANNELS, schema_version=HARMONICS_SCHEMA_VERSION))
    result: dict[str, Any] = {"scene": scene_name, "target_only": target_only, "controlled": False, "metric_id": METRIC_ID, "norm_source": norm_source}
    if target_only:
        del ff
        return result

    del ff; ff = FlowField(field_cfg, cuda_cfg, device_id=device_id); cy = (ff.FIELD_SHAPE[1] - 1) / 2
    for y in [cy + 2*l0, cy, cy - 2*l0]: ff.add_sensor((float(cfg["sensor_x"]) * l0, y, 0.0), 5.0)
    ff.add_cylinder((float(cfg["pinball_front_x"])*l0, cy, 0.0), 10.0)
    span = float(cfg.get("pinball_y_span", 0.75*l0))
    ff.add_cylinder((float(cfg["pinball_rear_x"])*l0, cy+span, 0.0), 10.0)
    ff.add_cylinder((float(cfg["pinball_rear_x"])*l0, cy-span, 0.0), 10.0)
    n_obj = ff.obs.size // 2; ff.run(warmup, np.zeros(n_obj, dtype=DATA_TYPE)); ff.get_ddf(); ff.save_ddf()
    if norm_source == "existing":
        norm, norm_path = load_existing_norm(scene_name, cfg)
    else:
        samples = []
        for _ in range(FIFO_LEN): ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE)); samples.append(ff.obs.copy()[slice(*cfg["obs_slice"])])
        norm, norm_path = _compute_norm(np.asarray(samples)), None
    result["norm_source_path"] = str(norm_path) if norm_path else None; _write_json(out_dir / "norm.json", norm)

    ff.apply_ddf(); fifo: deque[np.ndarray] = deque(maxlen=FIFO_LEN)
    bias_arr = np.zeros(n_obj, dtype=DATA_TYPE); bias_arr[-3:] = np.asarray([0.0, -u0, u0], dtype=DATA_TYPE)
    for _ in range(FIFO_LEN): ff.run(si, bias_arr); fifo.append(ff.obs.copy()[slice(*cfg["obs_slice"])])
    ff.get_ddf(); ff.save_ddf()
    model_path = model_path_for_scene(scene_name)
    if model_path:
        model = load_model(model_path, device=model_device, s_dim=cfg.get("s_dim", 14))
        raw = fifo[-1].copy()
        force0 = gen_target_states_at(target_phase_index(0), harmonics)[:2]
        obs = illusion_observation(raw, norm, force0)
        sensors, forces, actions, commanded, targets, phases = [], [], [], [], [], []
        for step in range(n_steps):
            phase = target_phase_index(step); target_force = gen_target_states_at(phase, harmonics)[:2]
            action, _ = model.predict(obs, deterministic=True); action = np.asarray(action, dtype=np.float32).flatten()
            alpha = (action * float(cfg["action_scale"]) + np.asarray(cfg["action_bias"], dtype=np.float32)) * u0
            omega = np.zeros(n_obj, dtype=DATA_TYPE); omega[-3:] = alpha
            ff.context.push(); ff.run(si, omega); ff.context.pop()
            raw = ff.obs.copy()[slice(*cfg["obs_slice"])]; fifo.append(raw)
            sensors.append(raw[:6]); forces.append(raw[6:12]); actions.append(action); commanded.append(alpha); targets.append(target_force); phases.append(phase)
            next_force = gen_target_states_at(target_phase_index(step, after_action=True), harmonics)[:2]
            obs = illusion_observation(raw, norm, next_force)
        save_controlled_npz(out_dir / "controlled.npz", scene_name, cfg, sensors=np.asarray(sensors), forces=np.asarray(forces),
                            actions_norm=np.asarray(actions), actions_commanded=np.asarray(commanded), target_forces=np.asarray(targets),
                            phase_indices=np.asarray(phases))
        result.update(controlled=True, similarity=float(similarity(target[:, 2:8], np.asarray(sensors), cfg.get("conv_len", 36))))
    del ff
    return result


def _scene_names(args: argparse.Namespace, parser: argparse.ArgumentParser) -> list[str]:
    if args.scene:
        return [args.scene]
    if args.group in SCENE_GROUPS:
        return list(SCENE_GROUPS[args.group]())
    return [item.strip() for item in args.group.split(",") if item.strip()]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Stage 1 legacy PPO inference")
    selection = parser.add_mutually_exclusive_group(required=True)
    selection.add_argument("--scene")
    selection.add_argument("--group")
    parser.add_argument("--device", type=int, default=2, help="Legacy CFD GPU device")
    parser.add_argument(
        "--model-device", default="cpu",
        help="PyTorch PPO device; CPU avoids conflicts with the active PyCUDA context",
    )
    parser.add_argument("--steps", type=int, default=200)
    parser.add_argument("--target-only", action="store_true")
    parser.add_argument("--run-id", default=None)
    parser.add_argument("--output-root", default=None)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--norm-source", choices=("existing", "recompute"), default="existing")
    parser.add_argument("--dry-run", action="store_true")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser(); args = parser.parse_args(argv)
    if args.steps <= 0: parser.error("--steps must be positive")
    run_id = args.run_id or default_run_id(); output_root = resolve_output_root(args.output_root)
    command = [sys.executable, str(Path(__file__).resolve()), *(list(argv) if argv is not None else sys.argv[1:])]
    try:
        scenes = _scene_names(args, parser)
        plans = []
        for scene in scenes:
            cfg = get_scene(scene)
            if args.target_only and cfg["scene_id"] != "illusion":
                raise ValueError("--target-only is only valid for Illusion scenes")
            model = Path(model_path_for_scene(scene)) if model_path_for_scene(scene) else None
            out_dir = scene_output_dir(output_root, run_id, scene, cfg)
            if out_dir.exists() and any(out_dir.iterdir()) and not args.overwrite:
                raise FileExistsError(f"Refusing to overwrite non-empty output directory: {out_dir}")
            if args.norm_source == "existing" and not args.target_only:
                load_existing_norm(scene, cfg)
            plans.append((scene, cfg, model, out_dir))
        if args.dry_run:
            for scene, cfg, model, out_dir in plans:
                print(f"{scene}: scene_id={cfg['scene_id']} model={model or 'NONE'} SI={cfg['sample_interval']} output={out_dir}")
            return 0
        device_manifest = cfd_device_manifest(args.device)
        for scene, cfg, model, out_dir in plans:
            prepare_output_dir(out_dir, args.overwrite)
            _write_json(out_dir / "config.json", cfg)
            manifest = {
                "schema_version": "sr-stage1-manifest-v2", "run_id": run_id, "scene": scene,
                "scene_id": cfg["scene_id"], "config": cfg, "model_path": str(model) if model else None,
                "model_sha256": _sha256(model), "metric_id": METRIC_ID, "command": shlex.join(command),
                "git_sha": _git_sha(), "created_utc": datetime.now(timezone.utc).isoformat(),
                "output_dir": str(out_dir), "norm_source": args.norm_source,
                **device_manifest, "model_device": args.model_device,
            }
            _write_json(out_dir / "manifest.json", manifest)
            if cfg["scene_id"] == "karman":
                result = infer_karman(
                    scene, args.device, out_dir, args.steps, args.norm_source, args.model_device,
                )
            elif cfg["scene_id"] == "illusion":
                result = infer_illusion(
                    scene, args.device, out_dir, args.steps, args.target_only,
                    args.norm_source, args.model_device,
                )
            else: raise ValueError(f"Unsupported legacy scene_id: {cfg['scene_id']}")
            result.update({k: manifest[k] for k in ("run_id", "model_path", "model_sha256", "metric_id", "command", "git_sha", "config")})
            result["schema_version"] = "sr-stage1-result-v2"
            _write_json(out_dir / "result.json", result)
        return 0
    except (FileNotFoundError, FileExistsError, KeyError, ValueError, NotImplementedError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
