#!/usr/bin/env python3
"""Stage 3 legacy closed-loop validation for Karman and Illusion.

The module is deliberately CUDA-free at import time.  Runtime CFD dependencies
are imported only by the environment factories after planning has completed.
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import subprocess
import sys
import traceback
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Mapping, Protocol, 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
    FIFO_LEN,
    LEGACY_CFG_DIR,
    get_scene,
    get_scene_list,
    model_path_for_scene,
)
from SR_analysis.stage_1_infer import (  # noqa: E402
    illusion_observation,
    load_existing_norm,
    target_phase_index,
)
from SR_analysis.utils.data_contracts import CAUSAL_ALIGNMENT  # noqa: E402
from SR_analysis.utils.feature_state import StatefulFeatureBuilder  # noqa: E402
from SR_analysis.utils.formula_schema import evaluate_formula, load_formula  # noqa: E402
from SR_analysis.utils.g_operator import apply_G_alpha, apply_G_raw12, apply_G_target_force  # noqa: E402
from SR_analysis.utils.harmonics import gen_target_states_at  # noqa: E402
from SR_analysis.utils.metrics import LEGACY_DTW_VERSION, legacy_cycle_similarity  # noqa: E402
from SR_analysis.utils.provenance import atomic_write_json, hash_file, hash_json  # noqa: E402

DATA_TYPE = np.float32
METRIC_VERSION = LEGACY_DTW_VERSION
VALIDATION_VERSION = "sr-stage3-validation-v2"
TELEMETRY_VERSION = "sr-stage3-telemetry-v2"
ACTION_ORDER = ("front", "upper", "lower")
GROUPS = {
    "karman_all": lambda: get_scene_list("karman"),
    "karman_trained": lambda: [
        scene for scene in get_scene_list("karman") if get_scene(scene).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"],
    "all": lambda: get_scene_list(),
}


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


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("validate-%Y%m%dT%H%M%SZ")


def _gpu_provenance(requested_logical_device: int | None) -> dict[str, Any]:
    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
    result: dict[str, Any] = {
        "requested_logical_device": requested_logical_device,
        "CUDA_VISIBLE_DEVICES": visible,
        "uuid": None,
        "name": None,
    }
    if requested_logical_device is None:
        return result
    selector = str(requested_logical_device)
    if visible is not None:
        entries = [entry.strip() for entry in visible.split(",") if entry.strip()]
        if requested_logical_device >= len(entries):
            return result
        selector = entries[requested_logical_device]
    try:
        output = subprocess.run(
            [
                "nvidia-smi", f"--id={selector}", "--query-gpu=uuid,name",
                "--format=csv,noheader,nounits",
            ],
            check=True,
            capture_output=True,
            text=True,
            timeout=5,
        ).stdout.strip().splitlines()
        if output:
            uuid, name = (part.strip() for part in output[0].split(",", 1))
            result.update({"uuid": uuid or None, "name": name or None})
    except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired, ValueError):
        pass
    return result


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


def validate_run_id(run_id: str) -> str:
    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 run_id


def _load_json(path: Path) -> Any:
    if not path.is_file():
        raise FileNotFoundError(f"Required frozen artifact not found: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def _hash_optional(path: Path | None) -> str | None:
    return hash_file(path) if path is not None and path.is_file() else None


@dataclass(frozen=True)
class FormulaPair:
    front: dict[str, Any]
    rear: dict[str, Any]
    front_path: Path
    rear_path: Path
    pair_hash: str
    compatibility: dict[str, Any]
    lower: dict[str, Any] | None = None
    lower_path: Path | None = None
    deployment_semantics: str = "mapped_shared"


def _validate_formula_role(formula: Mapping[str, Any], expected: str) -> dict[str, Any]:
    legacy = bool(formula.get("legacy"))
    role = formula.get("role")
    anchor = formula.get("anchor")
    parity = formula.get("parity", formula.get("g_parity"))
    channel = formula.get("channel")
    notes: dict[str, Any] = {"legacy": legacy, "channel": channel}
    if expected == "front":
        if legacy and channel == "front":
            notes["legacy_role_mapping"] = "channel=front -> role=front_odd"
        elif role == "front_odd":
            notes["canonical_role"] = "front_odd"
        elif role == "front" and parity == "odd":
            notes["compatibility_role_mapping"] = "role=front, parity=odd -> front_odd"
        else:
            raise ValueError("front formula requires canonical role='front_odd'")
    else:
        if legacy and channel == "top":
            notes["legacy_role_mapping"] = "channel=top -> role=rear_shared, anchor=upper"
        elif role != "rear_shared" or anchor != "upper":
            raise ValueError("rear formula requires role='rear_shared' and anchor='upper'")
    return notes


def load_formula_pair(
    front_path: str | os.PathLike[str],
    rear_path: str | os.PathLike[str],
    lower_path: str | os.PathLike[str] | None = None,
) -> FormulaPair:
    front_file = Path(front_path).expanduser().resolve()
    rear_file = Path(rear_path).expanduser().resolve()
    lower_file = None if lower_path is None else Path(lower_path).expanduser().resolve()
    front = load_formula(front_file)
    rear = load_formula(rear_file)
    lower = None if lower_file is None else load_formula(lower_file)
    semantics = "three_head_independent" if lower is not None else str(
        front.get("g_metadata", {}).get("deployment_semantics", "mapped_shared")
    )
    compatibility = {
        "deployment_semantics": semantics,
        "policy_mapping": (
            "exact_front_odd_projection_and_rear_g_mapping"
            if lower is None
            else "direct_three_head_evaluation"
        ),
    }
    if lower is None:
        compatibility.update({
            "front": _validate_formula_role(front, "front"),
            "rear": _validate_formula_role(rear, "rear"),
        })
        if bool(front.get("legacy")) or bool(rear.get("legacy")):
            compatibility["historical_formula_compatibility"] = (
                "legacy formula schema loaded; mapped_shared deployment still applies "
                "the exact current symmetry projection"
            )
    else:
        expected_roles = ("front_independent", "upper_independent", "lower_independent")
        for name, formula, expected_role in zip(("front", "upper", "lower"), (front, rear, lower), expected_roles):
            if formula.get("role") != expected_role:
                raise ValueError(f"{name} formula requires role={expected_role!r}")
            compatibility[name] = {"canonical_role": expected_role}
    for formula in (front, rear) if lower is None else (front, rear, lower):
        probe = {name: np.array([0.0]) for name in formula["feature_names"]}
        value = evaluate_formula(formula, probe, expression_kind="deployment")
        if not np.all(np.isfinite(value)):
            raise ValueError("formula deployment expression is non-finite at the zero probe")
    hashes = {"front": front["source_hash"], "rear": rear["source_hash"]}
    if lower is not None:
        hashes["lower"] = lower["source_hash"]
    return FormulaPair(
        front, rear, front_file, rear_file, hash_json(hashes), compatibility,
        lower=lower, lower_path=lower_file, deployment_semantics=semantics,
    )


@dataclass(frozen=True)
class ValidationPlan:
    scene: str
    cfg: dict[str, Any]
    mode: str
    n_steps: int
    run_id: str
    validation_dir: Path
    telemetry_dir: Path
    stem: str
    validation_path: Path
    telemetry_path: Path
    formula_pair: FormulaPair | None
    model_path: Path | None
    norm: dict[str, Any] | None
    norm_path: Path | None
    harmonics: Any | None
    harmonics_path: Path | None


def _artifact_stem(scene: str, mode: str, formula_hash: str, sample_interval: int) -> str:
    safe_scene = scene.replace(os.sep, "_").replace(".", "p")
    return f"{safe_scene}__{mode}__f-{formula_hash[:12]}__si-{sample_interval}"


def prepare_plan(
    *,
    scene: str,
    mode: str,
    n_steps: int,
    run_id: str,
    output_root: Path,
    formula_pair: FormulaPair | None,
) -> ValidationPlan:
    cfg = get_scene(scene)
    scene_id = cfg["scene_id"]
    if scene_id not in {"karman", "illusion", "steady"}:
        raise NotImplementedError(f"scene_id={scene_id!r} is not supported")
    if tuple(cfg.get("action_layout", ())) != ACTION_ORDER:
        raise ValueError(f"scene action_layout must be native {ACTION_ORDER}")
    if mode == "pysr" and formula_pair is None:
        raise ValueError("pysr mode requires both front and rear formulas")

    model_path = None
    norm = None
    norm_path = None
    if mode == "ppo":
        resolved = model_path_for_scene(scene)
        if resolved is None:
            raise FileNotFoundError(f"No PPO model found for scene {scene}")
        model_path = Path(resolved).resolve()
        norm, norm_path = load_existing_norm(scene, cfg)
    elif scene_id == "illusion":
        candidate_norm = REPO_ROOT / "src" / "SR_analysis" / "data" / "illusion" / scene / "norm.json"
        if candidate_norm.is_file():
            norm, norm_path = load_existing_norm(scene, cfg)
        elif cfg.get("source") != "generalization":
            raise FileNotFoundError(f"Frozen normalization not found: {candidate_norm}")

    harmonics = None
    harmonics_path = None
    if scene_id == "illusion":
        harmonics_path = REPO_ROOT / "src" / "SR_analysis" / "data" / "illusion" / scene / "target_harmonics.json"
        harmonics = _load_json(harmonics_path)

    formula_hash = formula_pair.pair_hash if formula_pair is not None else ("ppo" if mode == "ppo" else "none")
    validation_dir = output_root / validate_run_id(run_id) / "validations"
    telemetry_dir = output_root / validate_run_id(run_id) / "telemetry"
    stem = _artifact_stem(scene, mode, formula_hash, int(cfg["sample_interval"]))
    validation_path = validation_dir / f"{stem}.json"
    telemetry_path = telemetry_dir / f"{stem}.npz"
    for path in (validation_path, telemetry_path):
        if path.exists():
            raise FileExistsError(f"Refusing to overwrite existing artifact: {path}")
    return ValidationPlan(
        scene, cfg, mode, n_steps, run_id, validation_dir, telemetry_dir, stem,
        validation_path, telemetry_path, formula_pair, model_path, norm, norm_path,
        harmonics, harmonics_path,
    )


class RolloutEnvironment(Protocol):
    current_raw: np.ndarray
    target_sensors: np.ndarray

    def step(self, omega_native: np.ndarray) -> np.ndarray: ...
    def close(self) -> None: ...


class Policy(Protocol):
    def action(self, raw: np.ndarray, step: int) -> tuple[np.ndarray, np.ndarray | None, int | None]: ...
    def observe(self, omega_native: np.ndarray) -> None: ...


class UncontrolledPolicy:
    def action(self, raw: np.ndarray, step: int) -> tuple[np.ndarray, None, None]:
        return np.zeros(3, dtype=np.float64), None, None

    def observe(self, omega_native: np.ndarray) -> None:
        return None


class ConstantRotationPolicy:
    def __init__(self, alpha_native: Sequence[float], u0: float):
        self.omega_native = np.asarray(alpha_native, dtype=np.float64).reshape(3) * float(u0)

    def action(self, raw: np.ndarray, step: int) -> tuple[np.ndarray, None, None]:
        return self.omega_native.copy(), None, None

    def observe(self, omega_native: np.ndarray) -> None:
        return None


class PPOPolicy:
    def __init__(self, model: Any, cfg: Mapping[str, Any], norm: Mapping[str, Any], harmonics: Any | None):
        self.model = model
        self.cfg = cfg
        self.norm = norm
        self.harmonics = harmonics

    def action(self, raw: np.ndarray, step: int) -> tuple[np.ndarray, np.ndarray | None, int | None]:
        target_force = None
        phase = None
        if self.cfg["scene_id"] == "illusion":
            phase = target_phase_index(step)
            target_force = gen_target_states_at(phase, self.harmonics)[:2]
            observation = illusion_observation(raw, self.norm, target_force)
        else:
            forces = raw[6:12] / float(self.norm["force_norm_fact"])
            sensors = (raw[:6] - np.asarray(self.norm["sens_deviation"])) / np.asarray(self.norm["sens_norm_fact"])
            observation = np.clip(np.hstack((forces, sensors)), -1, 1).astype(np.float32)
        normalized, _ = self.model.predict(observation, deterministic=True)
        normalized = np.asarray(normalized, dtype=np.float64).reshape(3)
        alpha = normalized * float(self.cfg["action_scale"]) + np.asarray(self.cfg["action_bias"], dtype=np.float64)
        return alpha * float(self.cfg["u0"]), target_force, phase

    def observe(self, omega_native: np.ndarray) -> None:
        return None


class SymbolicPolicy:
    def __init__(self, pair: FormulaPair, cfg: Mapping[str, Any], harmonics: Any | None):
        self.pair = pair
        self.cfg = cfg
        self.harmonics = harmonics
        self.dt_c = float(cfg["sample_interval"]) * float(cfg["u0"]) / 20.0
        common = dict(
            mu=float(cfg["mu"]), dt_c=self.dt_c, u0=float(cfg["u0"]), d=20.0,
            action_mode="physical", include_mu=True, include_cos_sin=True,
            require_target=cfg["scene_id"] == "illusion",
        )
        self.builder = StatefulFeatureBuilder(**common)
        self.g_builder = StatefulFeatureBuilder(**common)
        initial_alpha = np.asarray(cfg["policy_init_action"], dtype=np.float64)
        initial_omega = initial_alpha * float(cfg["u0"])
        # The first evaluated state is the current post-bias FIFO observation;
        # seed both action lags with the command that produced that FIFO.
        for _ in range(2):
            self.builder.observe_action(initial_omega)
            self.g_builder.observe_action(apply_G_alpha(initial_omega))

    @staticmethod
    def _evaluate(formula: Mapping[str, Any], features: Mapping[str, Any]) -> float:
        value = np.asarray(evaluate_formula(formula, features, expression_kind="deployment"), dtype=float)
        if value.size != 1:
            raise ValueError("controller formula must produce one scalar")
        scalar = float(value.reshape(-1)[0])
        if not np.isfinite(scalar):
            raise FloatingPointError("controller formula produced a non-finite action")
        return scalar

    def action(self, raw: np.ndarray, step: int) -> tuple[np.ndarray, np.ndarray | None, int | None]:
        phase = target_phase_index(step) if self.harmonics is not None else None
        target_force = None if phase is None else gen_target_states_at(phase, self.harmonics)[:2]
        features, _ = self.builder.step(raw[:6], raw[6:12], target_force=target_force)
        g_raw = apply_G_raw12(raw)
        g_target = None if target_force is None else apply_G_target_force(target_force)
        g_features, _ = self.g_builder.step(g_raw[:6], g_raw[6:12], target_force=g_target)
        if self.pair.deployment_semantics == "three_head_independent" and self.pair.lower is not None:
            front = self._evaluate(self.pair.front, features)
            upper = self._evaluate(self.pair.rear, features)
            lower = self._evaluate(self.pair.lower, features)
        else:
            front = 0.5 * (
                self._evaluate(self.pair.front, features)
                - self._evaluate(self.pair.front, g_features)
            )
            upper = self._evaluate(self.pair.rear, features)
            lower = -self._evaluate(self.pair.rear, g_features)
        alpha_native = np.asarray((front, upper, lower), dtype=np.float64)
        return alpha_native * float(self.cfg["u0"]), target_force, phase

    def observe(self, omega_native: np.ndarray) -> None:
        self.builder.observe_action(omega_native)
        self.g_builder.observe_action(apply_G_alpha(omega_native))


def build_policy(plan: ValidationPlan, *, model_loader: Callable[..., Any] | None = None) -> Policy:
    if plan.mode == "uncontrolled":
        return UncontrolledPolicy()
    if plan.mode == "constant":
        return ConstantRotationPolicy(plan.cfg["constant_alpha"], plan.cfg["u0"])
    if plan.mode == "pysr":
        assert plan.formula_pair is not None
        return SymbolicPolicy(plan.formula_pair, plan.cfg, plan.harmonics)
    if model_loader is None:
        from SR_analysis.utils.cfd_interface import load_ppo_model
        model_loader = load_ppo_model
    assert plan.model_path is not None and plan.norm is not None
    model = model_loader(
        str(plan.model_path),
        device=str(plan.cfg.get("model_device", "cpu")),
        s_dim=int(plan.cfg["s_dim"]),
    )
    return PPOPolicy(model, plan.cfg, plan.norm, plan.harmonics)


class _LegacyEnvironment:
    def __init__(self, ff: Any, cfg: Mapping[str, Any], current_raw: np.ndarray, target_sensors: np.ndarray):
        self.ff = ff
        self.cfg = cfg
        self.current_raw = np.asarray(current_raw, dtype=np.float64)
        self.target_sensors = np.asarray(target_sensors, dtype=np.float64)

    def step(self, omega_native: np.ndarray) -> np.ndarray:
        command = np.zeros(int(self.cfg["n_objects_env"]), dtype=DATA_TYPE)
        command[-3:] = np.asarray(omega_native, dtype=DATA_TYPE)
        self.ff.context.push()
        self.ff.run(int(self.cfg["sample_interval"]), command)
        self.ff.context.pop()
        self.current_raw = self.ff.obs.copy()[slice(*self.cfg["obs_slice"])].astype(np.float64)
        return self.current_raw.copy()

    def close(self) -> None:
        del self.ff


def _runtime_cfd() -> tuple[Any, Any, Any, Any]:
    from LegacyCelerisLab import FlowField
    from SR_analysis.utils.cfd_interface import add_pinball, build_karman_cloak_env, load_legacy_configs
    return FlowField, add_pinball, build_karman_cloak_env, load_legacy_configs


def build_karman_environment(plan: ValidationPlan, device_id: int) -> RolloutEnvironment:
    FlowField, add_pinball, build_karman, load_configs = _runtime_cfd()
    cfg = plan.cfg
    cuda_cfg, field_cfg = load_configs(LEGACY_CFG_DIR)
    ff = FlowField(field_cfg._replace(viscosity=float(cfg["nu"])), cuda_cfg, device_id=device_id)
    target, _ = build_karman(ff, u0=cfg["u0"], l0=20.0, sample_interval=cfg["sample_interval"], fifo_len=FIFO_LEN, data_type=DATA_TYPE)
    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],
    )
    current = np.asarray(norm["save_states"])[-1]
    return _LegacyEnvironment(ff, cfg, current, target)


def build_illusion_environment(plan: ValidationPlan, device_id: int) -> RolloutEnvironment:
    FlowField, _, _, load_configs = _runtime_cfd()
    cfg = plan.cfg
    cuda_cfg, field_cfg = load_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))
    l0 = 20.0

    # Record the canonical target sensor orbit with the exact Stage 1 phase
    # construction. Frozen harmonics remain the sole force/phase source; they
    # are never recomputed from this validation rollout.
    target_ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target_cy = (target_ff.FIELD_SHAPE[1] - 1) / 2
    target_ff.add_cylinder((20.0 * l0, target_cy, 0.0), float(cfg["target_diameter"]) * l0)
    for y in (target_cy + 2 * l0, target_cy, target_cy - 2 * l0):
        target_ff.add_sensor((float(cfg["sensor_x"]) * l0, y, 0.0), 5.0)
    target_count = target_ff.obs.size // 2
    target_ff.run(int(4 * target_ff.FIELD_SHAPE[0] / float(cfg["u0"])), np.zeros(target_count, dtype=DATA_TYPE))
    target_rows = []
    for _ in range(FIFO_LEN):
        target_ff.run(int(cfg["sample_interval"]), np.zeros(target_count, dtype=DATA_TYPE))
        target_rows.append(target_ff.obs.copy()[2:8])
    target_sensors = np.asarray(target_rows, dtype=np.float64)
    del target_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(int(4 * ff.FIELD_SHAPE[0] / float(cfg["u0"])), np.zeros(n_obj, dtype=DATA_TYPE))
    ff.get_ddf(); ff.save_ddf(); ff.apply_ddf()
    command = np.zeros(n_obj, dtype=DATA_TYPE)
    command[-3:] = np.asarray(cfg["policy_init_action"], dtype=DATA_TYPE) * float(cfg["u0"])
    current = None
    for _ in range(FIFO_LEN):
        ff.run(int(cfg["sample_interval"]), command)
        current = ff.obs.copy()[slice(*cfg["obs_slice"])]
    assert current is not None
    return _LegacyEnvironment(ff, cfg, current, target_sensors)


def build_steady_environment(plan: ValidationPlan, device_id: int) -> RolloutEnvironment:
    FlowField, add_pinball, _, load_configs = _runtime_cfd()
    cfg = plan.cfg
    cuda_cfg, field_cfg = load_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))
    l0 = 20.0

    target_ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target_cy = (target_ff.FIELD_SHAPE[1] - 1) / 2
    for y in (target_cy + 2 * l0, target_cy, target_cy - 2 * l0):
        target_ff.add_sensor((float(cfg["sensor_x"]) * l0, y, 0.0), 5.0)
    target_count = target_ff.obs.size // 2
    target_ff.run(int(4 * target_ff.FIELD_SHAPE[0] / float(cfg["u0"])), np.zeros(target_count, dtype=DATA_TYPE))
    target_rows = []
    for _ in range(FIFO_LEN):
        target_ff.run(int(cfg["sample_interval"]), np.zeros(target_count, dtype=DATA_TYPE))
        target_rows.append(target_ff.obs.copy()[:6])
    target_sensors = np.asarray(target_rows, dtype=np.float64)
    del target_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)
    initial_alpha = np.asarray(cfg.get("constant_alpha", cfg["policy_init_action"]), dtype=np.float64)
    norm = add_pinball(
        ff, l0=l0, u0=cfg["u0"], sample_interval=cfg["sample_interval"], fifo_len=FIFO_LEN,
        data_type=DATA_TYPE, action_bias=tuple(initial_alpha), 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], n_objects_total=cfg["n_objects_env"],
    )
    current = np.asarray(norm["save_states"])[-1]
    return _LegacyEnvironment(ff, cfg, current, target_sensors)


def build_environment(plan: ValidationPlan, device_id: int) -> RolloutEnvironment:
    if plan.cfg["scene_id"] == "karman":
        return build_karman_environment(plan, device_id)
    if plan.cfg["scene_id"] == "steady":
        return build_steady_environment(plan, device_id)
    return build_illusion_environment(plan, device_id)


def _finite(name: str, values: np.ndarray) -> None:
    if not np.all(np.isfinite(values)):
        raise FloatingPointError(f"{name} contains non-finite values")


class RolloutFailure(RuntimeError):
    def __init__(self, message: str, *, step: int, telemetry: Mapping[str, np.ndarray]):
        super().__init__(message)
        self.step = step
        self.telemetry = telemetry


def _telemetry_arrays(
    environment: RolloutEnvironment,
    sensors: list[np.ndarray],
    forces: list[np.ndarray],
    actions_omega: list[np.ndarray],
    actions_alpha: list[np.ndarray],
    target_forces: list[np.ndarray],
    target_phases: list[int],
    termination_reason: str,
) -> dict[str, np.ndarray]:
    return {
        "sensors": np.asarray(sensors, dtype=np.float64).reshape(-1, 6),
        "forces": np.asarray(forces, dtype=np.float64).reshape(-1, 6),
        "actions_omega": np.asarray(actions_omega, dtype=np.float64).reshape(-1, 3),
        "actions_alpha": np.asarray(actions_alpha, dtype=np.float64).reshape(-1, 3),
        "target_forces": np.asarray(target_forces, dtype=np.float64).reshape(-1, 2),
        "target_phase": np.asarray(target_phases, dtype=np.int64),
        "target_sensors": np.asarray(environment.target_sensors, dtype=np.float64).reshape(-1, 6),
        "termination_reason": np.asarray(termination_reason),
    }


def run_rollout(plan: ValidationPlan, environment: RolloutEnvironment, policy: Policy) -> dict[str, Any]:
    sensors: list[np.ndarray] = []
    forces: list[np.ndarray] = []
    actions_omega: list[np.ndarray] = []
    actions_alpha: list[np.ndarray] = []
    target_forces: list[np.ndarray] = []
    target_phases: list[int] = []
    raw = np.asarray(environment.current_raw, dtype=np.float64).reshape(12)
    termination_reason = "completed"
    for step in range(plan.n_steps):
        try:
            omega, target_force, phase = policy.action(raw, step)
            omega = np.asarray(omega, dtype=np.float64).reshape(3)
            _finite("action omega", omega)
            next_raw = np.asarray(environment.step(omega), dtype=np.float64).reshape(12)
            _finite("raw observation", next_raw)
        except Exception as exc:
            reason = f"{type(exc).__name__}: {exc}"
            partial = _telemetry_arrays(
                environment, sensors, forces, actions_omega, actions_alpha,
                target_forces, target_phases, reason,
            )
            raise RolloutFailure(reason, step=step, telemetry=partial) from exc
        raw = next_raw
        policy.observe(omega)
        sensors.append(raw[:6].copy())
        forces.append(raw[6:].copy())
        actions_omega.append(omega.copy())
        actions_alpha.append(omega / float(plan.cfg["u0"]))
        if target_force is not None:
            target_forces.append(np.asarray(target_force, dtype=np.float64).reshape(2))
        if phase is not None:
            target_phases.append(int(phase))
    arrays = _telemetry_arrays(
        environment, sensors, forces, actions_omega, actions_alpha,
        target_forces, target_phases, termination_reason,
    )
    for name, values in arrays.items():
        if name not in {"termination_reason", "target_phase"}:
            _finite(name, values)
    return arrays


def _legacy_metric(target: np.ndarray, state: np.ndarray, conv_len: int) -> dict[str, Any]:
    return legacy_cycle_similarity(target, state, conv_len)


def compute_versioned_metrics(target: np.ndarray, sensors: np.ndarray, conv_len: int) -> dict[str, Any]:
    """Compute the historical reference-cycle versus final recorded-cycle metric."""
    exact = _legacy_metric(target, np.asarray(sensors, dtype=np.float64), conv_len)
    return {
        "metric_version": METRIC_VERSION,
        "adapter": "SR_analysis.utils.cfd_interface.legacy_dtw_v1_detail_adapter",
        "legacy_reference_cycle_vs_last_recorded_cycle": exact,
    }


def _action_ranges(telemetry: Mapping[str, np.ndarray]) -> dict[str, Any]:
    result = {}
    for key in ("actions_omega", "actions_alpha"):
        values = np.asarray(telemetry[key])
        result[key] = {
            "min": values.min(axis=0).tolist(),
            "max": values.max(axis=0).tolist(),
            "layout": list(ACTION_ORDER),
        }
    return result


def save_outputs(plan: ValidationPlan, telemetry: Mapping[str, np.ndarray], metrics: Mapping[str, Any]) -> dict[str, Any]:
    plan.validation_dir.mkdir(parents=True, exist_ok=True)
    plan.telemetry_dir.mkdir(parents=True, exist_ok=True)
    for path in (plan.validation_path, plan.telemetry_path):
        if path.exists():
            raise FileExistsError(f"Refusing to overwrite existing artifact: {path}")
    np.savez_compressed(
        plan.telemetry_path,
        **telemetry,
        schema_version=np.asarray(TELEMETRY_VERSION),
        scene=np.asarray(plan.scene),
        mode=np.asarray(plan.mode),
        sample_interval=np.asarray(plan.cfg["sample_interval"], dtype=np.int64),
    )
    formula_hashes = None
    formula_validation = None
    if plan.formula_pair is not None:
        formula_hashes = {
            "pair": plan.formula_pair.pair_hash,
            "front_source": plan.formula_pair.front["source_hash"],
            "rear_source": plan.formula_pair.rear["source_hash"],
            "front_deployment_expression": plan.formula_pair.front["deployment_expression_hash"],
            "rear_deployment_expression": plan.formula_pair.rear["deployment_expression_hash"],
        }
        formula_validation = plan.formula_pair.compatibility
    record = {
        "schema_version": VALIDATION_VERSION,
        "run_id": plan.run_id,
        "scene": plan.scene,
        "scene_id": plan.cfg["scene_id"],
        "mode": plan.mode,
        "artifact_id": f"{plan.run_id}-{plan.stem}",
        "formula_hash": plan.formula_pair.pair_hash if plan.formula_pair is not None else None,
        "metric_version": METRIC_VERSION,
        "status": "completed",
        "termination_reason": str(np.asarray(telemetry["termination_reason"]).item()),
        "n_steps": plan.n_steps,
        "sample_interval": int(plan.cfg["sample_interval"]),
        "si": int(plan.cfg["sample_interval"]),
        "dt_c": float(plan.cfg["sample_interval"]) * float(plan.cfg["u0"]) / 20.0,
        "alignment": CAUSAL_ALIGNMENT,
        "action_layout": list(ACTION_ORDER),
        "metrics": _jsonable(metrics),
        "action_ranges": _action_ranges(telemetry),
        "formula_hashes": formula_hashes,
        "formula_validation": formula_validation,
        "data_hashes": {
            "norm": _hash_optional(plan.norm_path),
            "harmonics": _hash_optional(plan.harmonics_path),
            "target_sensors": hash_json(np.asarray(telemetry["target_sensors"]).tolist()),
        },
        "model": {"path": str(plan.model_path) if plan.model_path else None, "sha256": _hash_optional(plan.model_path)},
        "config": _jsonable(plan.cfg),
        "config_hash": hash_json(_jsonable(plan.cfg)),
        "environment": {"python": sys.version, "platform": platform.platform(), "numpy": np.__version__},
        "git_sha": _git_sha(),
        "telemetry": {
            "path": str(plan.telemetry_path),
            "sha256": hash_file(plan.telemetry_path),
            "telemetry_hash": hash_file(plan.telemetry_path),
        },
        "provenance": {
            "git_sha": _git_sha(),
            "environment": {
                "conda_env": os.environ.get("CONDA_DEFAULT_ENV"),
                "gpu": _gpu_provenance(plan.cfg.get("device_id")),
            },
            "command": " ".join(sys.argv),
            "created_at": datetime.now(timezone.utc).isoformat(),
        },
        "created_utc": datetime.now(timezone.utc).isoformat(),
    }
    record["record_hash"] = hash_json(record)
    atomic_write_json(plan.validation_path, record)
    return record


def save_failure_outputs(
    plan: ValidationPlan,
    failure: RolloutFailure,
) -> dict[str, Any]:
    plan.validation_dir.mkdir(parents=True, exist_ok=True)
    plan.telemetry_dir.mkdir(parents=True, exist_ok=True)
    for path in (plan.validation_path, plan.telemetry_path):
        if path.exists():
            raise FileExistsError(f"Refusing to overwrite existing artifact: {path}")
    np.savez_compressed(
        plan.telemetry_path,
        **failure.telemetry,
        schema_version=np.asarray(TELEMETRY_VERSION),
        scene=np.asarray(plan.scene),
        mode=np.asarray(plan.mode),
        sample_interval=np.asarray(plan.cfg["sample_interval"], dtype=np.int64),
    )
    formula_hashes = None
    if plan.formula_pair is not None:
        formula_hashes = {
            "pair": plan.formula_pair.pair_hash,
            "front_source": plan.formula_pair.front["source_hash"],
            "rear_source": plan.formula_pair.rear["source_hash"],
            "front_deployment_expression": plan.formula_pair.front["deployment_expression_hash"],
            "rear_deployment_expression": plan.formula_pair.rear["deployment_expression_hash"],
        }
    created_at = datetime.now(timezone.utc).isoformat()
    record = {
        "schema_version": VALIDATION_VERSION,
        "run_id": plan.run_id,
        "scene": plan.scene,
        "scene_id": plan.cfg["scene_id"],
        "mode": plan.mode,
        "artifact_id": f"{plan.run_id}-{plan.stem}",
        "formula_hash": plan.formula_pair.pair_hash if plan.formula_pair is not None else None,
        "metric_version": METRIC_VERSION,
        "status": "failed",
        "termination_reason": str(failure),
        "exception_type": type(failure.__cause__).__name__ if failure.__cause__ else type(failure).__name__,
        "failed_step": int(failure.step),
        "completed_steps": int(np.asarray(failure.telemetry["sensors"]).shape[0]),
        "n_steps": plan.n_steps,
        "sample_interval": int(plan.cfg["sample_interval"]),
        "si": int(plan.cfg["sample_interval"]),
        "dt_c": float(plan.cfg["sample_interval"]) * float(plan.cfg["u0"]) / 20.0,
        "alignment": CAUSAL_ALIGNMENT,
        "action_layout": list(ACTION_ORDER),
        "metrics": None,
        "action_ranges": None,
        "formula_hashes": formula_hashes,
        "formula_validation": plan.formula_pair.compatibility if plan.formula_pair is not None else None,
        "config": _jsonable(plan.cfg),
        "config_hash": hash_json(_jsonable(plan.cfg)),
        "git_sha": _git_sha(),
        "telemetry": {
            "path": str(plan.telemetry_path),
            "sha256": hash_file(plan.telemetry_path),
            "telemetry_hash": hash_file(plan.telemetry_path),
        },
        "provenance": {
            "git_sha": _git_sha(),
            "environment": {
                "conda_env": os.environ.get("CONDA_DEFAULT_ENV"),
                "gpu": _gpu_provenance(plan.cfg.get("device_id")),
            },
            "command": " ".join(sys.argv),
            "created_at": created_at,
        },
        "created_utc": created_at,
    }
    record["record_hash"] = hash_json(record)
    atomic_write_json(plan.validation_path, record)
    return record


def execute_plan(
    plan: ValidationPlan,
    device_id: int,
    *,
    environment_factory: Callable[[ValidationPlan, int], RolloutEnvironment] = build_environment,
    model_loader: Callable[..., Any] | None = None,
) -> dict[str, Any]:
    cfg = dict(plan.cfg)
    cfg["device_id"] = device_id
    runtime_plan = ValidationPlan(**{**plan.__dict__, "cfg": cfg})
    environment = environment_factory(runtime_plan, device_id)
    try:
        policy = build_policy(runtime_plan, model_loader=model_loader)
        try:
            telemetry = run_rollout(runtime_plan, environment, policy)
        except RolloutFailure as failure:
            save_failure_outputs(runtime_plan, failure)
            raise
        metrics = compute_versioned_metrics(environment.target_sensors, telemetry["sensors"], int(cfg["conv_len"]))
        return save_outputs(runtime_plan, telemetry, metrics)
    finally:
        environment.close()


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


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Stage 3 legacy CFD validation (round1)")
    selection = parser.add_mutually_exclusive_group(required=True)
    selection.add_argument("--scene")
    selection.add_argument("--group")
    parser.add_argument("--mode", choices=("pysr", "ppo", "uncontrolled", "constant"), default="pysr")
    parser.add_argument("--constant-alpha", nargs=3, type=float, metavar=("FRONT", "UPPER", "LOWER"))
    parser.add_argument("--formula-front")
    parser.add_argument("--formula-rear", "--formula-top", dest="formula_rear")
    parser.add_argument("--formula-lower")
    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("--run-id", default=None)
    parser.add_argument("--output-root", default=None)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--prepare-only", 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")
    if args.mode == "pysr" and (not args.formula_front or not args.formula_rear):
        parser.error("--mode pysr requires --formula-front and --formula-rear/--formula-top")
    if args.mode == "constant" and args.constant_alpha is None:
        parser.error("--mode constant requires --constant-alpha FRONT UPPER LOWER")
    run_id = args.run_id or default_run_id()
    output_root = resolve_output_root(args.output_root)
    failures = 0
    try:
        pair = load_formula_pair(args.formula_front, args.formula_rear, args.formula_lower) if args.mode == "pysr" else None
        plans = [
            prepare_plan(
                scene=scene, mode=args.mode, n_steps=args.steps, run_id=run_id,
                output_root=output_root, formula_pair=pair,
            )
            for scene in _scene_names(args)
        ]
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    if args.dry_run or args.prepare_only:
        for plan in plans:
            print(
                f"{plan.scene}: scene_id={plan.cfg['scene_id']} mode={plan.mode} "
                f"SI={plan.cfg['sample_interval']} validation={plan.validation_path} telemetry={plan.telemetry_path}"
            )
        if args.prepare_only:
            for plan in plans:
                plan.validation_dir.mkdir(parents=True, exist_ok=True)
                plan.telemetry_dir.mkdir(parents=True, exist_ok=True)
        return 0

    plans = [
        ValidationPlan(**{**plan.__dict__, "cfg": {**plan.cfg, "model_device": args.model_device}})
        for plan in plans
    ]
    if args.mode == "constant":
        plans = [
            ValidationPlan(**{**plan.__dict__, "cfg": {**plan.cfg, "constant_alpha": tuple(args.constant_alpha)}})
            for plan in plans
        ]
    for plan in plans:
        try:
            record = execute_plan(plan, args.device)
            metric = record["metrics"]["legacy_reference_cycle_vs_last_recorded_cycle"]
            print(f"{plan.scene}: legacy_reference_cycle_vs_last_recorded_cycle={metric['similarity']:.6f}")
        except Exception as exc:
            failures += 1
            print(f"ERROR [{plan.scene}]: {exc}", file=sys.stderr)
            traceback.print_exc()
    return 1 if failures else 0


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