DynamisLab/src/drl_pinball/legacy_test/acquire.py

"""Legacy Karman Re100 role acquisition; solver imports stay lazy."""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
from collections import deque
from pathlib import Path
from typing import Any, Callable

import numpy as np

from drl_pinball.acquisition import (
    accumulate_mean_fields, assign_phase, assign_periodic_phase,
    cleanup_scratch, complete_cycle_field_indices, create_scratch,
    default_reproduction_mapping, dual_cycle_dtw, pooled_phase_bins,
    prepare_role_output, publish_role_output, publish_selected_fields,
    select_phase_fields, validate_output_storage, write_boundary_artifacts,
    write_phase_cycle_artifacts,
)
from .cases import get_case
from .metrics import erase_reward_terms, frozen_reference_comparison, reward_terms, sha256_file
from .core.dtw_metrics import gen_target_states_at
from .runtime import erase_policy_observation, load_policy_norm, policy_observation, reset_runtime, run_historical_interval

CASE_NAME, ROLE = "karman_re100", "controlled"  # compatibility defaults
PERIODIC_CASES = ("karman_re50", "karman_re100", "karman_re200", "karman_re400",
                  "illusion_075L", "illusion_1L", "illusion_15L")
SUPPORTED_CASES = PERIODIC_CASES + ("steady", "vortex_lamb", "vortex_taylor", "erase")
PERIODIC_ROLES = ("controlled", "target", "zero")
STEADY_ROLES = ("target", "constant", "zero")
ROLES = PERIODIC_ROLES
WARMUP_INTERVALS, COLLECT_BOUNDARIES = 480, 160
NX, U0, SI, CYCLE_LENGTH = 1280, 0.01, 800, 30  # compatibility defaults
MIN_FREE_BYTES = 4 * 1024**3


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    temporary = path.with_name(f".{path.name}.tmp")
    with temporary.open("w", encoding="utf-8") as stream:
        json.dump(payload, stream, indent=2)
        stream.write("\n")
        stream.flush()
        os.fsync(stream.fileno())
    os.replace(temporary, path)


def _identity(path: Path) -> dict[str, Any]:
    return {"path": str(path.resolve()), "sha256": sha256_file(str(path)), "bytes": path.stat().st_size}


def _canonical_json_value(value):
    """Convert builder contract values to deterministic JSON-native values."""
    if isinstance(value, dict):
        return {str(key): _canonical_json_value(current) for key, current in value.items()}
    if isinstance(value, (list, tuple)):
        return [_canonical_json_value(current) for current in value]
    if isinstance(value, np.ndarray):
        return _canonical_json_value(value.tolist())
    if isinstance(value, np.generic):
        return value.item()
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    raise TypeError(f"builder config contains non-JSON value {type(value).__name__}")


def _optional_file_provenance(path: Path | None):
    if path is None:
        return None
    path = Path(path)
    return {"path": str(path.resolve()), "exists": path.is_file(),
            "file_identity": _identity(path) if path.is_file() else None}


def _generated_array_identity(value) -> dict[str, Any]:
    """Identify a run-generated numeric array by dtype, shape, and C-order content."""
    array = np.asarray(value)
    if array.dtype.hasobject:
        raise TypeError("run-generated identity does not support object arrays")
    descriptor = {"dtype": array.dtype.str, "shape": list(array.shape), "order": "C"}
    header = json.dumps(descriptor, sort_keys=True, separators=(",", ":"),
                        allow_nan=False).encode("ascii")
    content = np.ascontiguousarray(array).tobytes(order="C")
    digest = hashlib.sha256(header + b"\0" + content).hexdigest()
    return {"identity_kind": "run-generated-array-content", **descriptor,
            "content_sha256": digest, "bytes": len(content)}


def _generated_config_identity(config) -> dict[str, Any]:
    """Identify a run-generated builder config by canonical JSON content."""
    values = _canonical_json_value(config)
    canonical = json.dumps(values, sort_keys=True, separators=(",", ":"),
                           ensure_ascii=True, allow_nan=False)
    return {"identity_kind": "run-generated-builder-config", "values": values,
            "canonical_json_sha256": hashlib.sha256(canonical.encode("ascii")).hexdigest()}


def _role_semantics(role: str) -> str:
    return {
        "controlled": "frozen-policy controlled trajectory",
        "target": "live builder-generated target trajectory; actions, rewards, and policy normalization unavailable",
        "zero": "physical-zero/uncontrolled trajectory using counter-bias normalized action; frozen norm used only for native reward evaluation",
        "constant": "established steady constant control at case.action_bias*U0 from normalized action zero",
    }[role]


def _controlled_reference(role, bundle, columns, conv_len):
    if role != "controlled":
        return None
    return frozen_reference_comparison(
        str(bundle["reference_path"]), columns["sensors"], columns["forces"],
        columns["action_normalized"], conv_len,
    )


def _case_roles(case):
    if case.name not in SUPPORTED_CASES:
        raise ValueError(f"unsupported Legacy acquisition case {case.name!r}")
    return STEADY_ROLES if case.scene == "steady" else PERIODIC_ROLES


def _resolve_bundle(
    repo_root: Path, case_name: str = CASE_NAME, role: str = "controlled",
) -> dict[str, Any]:
    """Resolve immutable CPU artifacts before solver construction."""
    case = get_case(case_name)
    _case_roles(case)
    frozen = repo_root / "src" / "SR_analysis" / "data" / case.scene / (case.reference_case or case.name)
    needs_model = case.model is not None and role == "controlled"
    needs_frozen_policy = (case.model is not None and role != "target" and
                           case.scene != "erase" and
                           not (case.scene == "vortex" and role == "zero"))
    norm_path = frozen / "norm.json" if needs_frozen_policy else None
    norm = load_policy_norm(str(norm_path)) if norm_path is not None else None
    model = None
    if needs_model:
        from .core.model_loader import model_path
        model = Path(model_path(case.model))
    return {"case": case, "model_path": model, "norm_path": norm_path, "norm": norm,
            "reference_path": frozen / "controlled.npz"}


def _default_runtime(case, bundle, device_id, role="controlled"):
    from .core.legacy_env_builder import (
        build_illusion, build_illusion_target, build_karman_cloak, build_karman_target,
        build_erase, build_steady_cloak, build_steady_target, build_vortex,
    )
    if case.scene == "vortex":
        data = build_vortex(device_id, case.re_code, vortex_type=case.vortex_type, action_scale=case.action_scale, target_only=(role == "target"))
    elif case.scene == "erase":
        data = build_erase(device_id, case.re_code, target_only=(role == "target"))
    elif case.scene == "karman":
        data = (build_karman_target(device_id, case.re_code, sample_interval=case.sample_interval)
                if role == "target" else
                build_karman_cloak(device_id, case.re_code, sample_interval=case.sample_interval))
    elif case.scene == "illusion":
        kwargs = dict(device_id=device_id, re_code=case.re_code,
                      target_diameter_L=case.target_radius_l, sample_interval=case.sample_interval)
        data = build_illusion_target(**kwargs) if role == "target" else build_illusion(**kwargs)
    elif case.scene == "steady":
        data = (build_steady_target(device_id, case.re_code, sample_interval=case.sample_interval)
                if role == "target" else build_steady_cloak(device_id, case.re_code))
    else:
        raise ValueError(f"unsupported Legacy acquisition scene {case.scene!r}")
    model = None
    if role == "controlled":
        from .core.model_loader import load_model
        model = load_model(case.model)
    return data, model, data["flow_field"]


def _run_interval(ff, steps: int, command: np.ndarray) -> None:
    context = getattr(ff, "context", None)
    if context is None:
        run_historical_interval(ff, steps, command)
        return
    context.push()
    try:
        run_historical_interval(ff, steps, command)
    finally:
        context.pop()


def _command(case, action, n_objects, dtype):
    action = np.asarray(action, dtype=np.float32).reshape(3)
    command = np.zeros(n_objects, dtype=dtype)
    command[-3:] = (action * case.action_scale + np.asarray(case.action_bias)) * U0
    return action, command


def _clock(ff, completed_intervals: int, sample_interval: int = SI) -> tuple[int, int, int | None]:
    """Return relative policy clocks plus the solver's absolute lattice clock."""
    lattice_step = completed_intervals * sample_interval
    absolute = None
    if hasattr(ff, "solver_clock_state"):
        state = ff.solver_clock_state()
        absolute = int(state["solver_absolute_lattice_clock"])
        if absolute < lattice_step:
            raise ValueError("solver lattice clock regressed below acquisition timeline")
    return lattice_step, completed_intervals, absolute


def _capture_boundary_field(ff, nx: int, ny: int) -> dict[str, np.ndarray]:
    """Capture public Legacy q/U0, convert to physical q/RHO_ref, and canonicalize."""
    velocity = ff.current_step_velocity_field()
    if not isinstance(velocity, tuple) or len(velocity) != 2:
        raise ValueError("public Legacy velocity accessor must return (ux, uy)")
    flags = np.asarray(ff.completed_flags_xy())
    if flags.dtype != np.uint8 or flags.shape != (nx, ny):
        raise ValueError("completed Legacy flags must be exact uint8 (nx,ny)")
    fluid_xy = (flags & np.uint8(0b00000001)) != 0
    u0 = float(ff.field_config.velocity)
    if not np.isfinite(u0) or u0 <= 0:
        raise ValueError("Legacy field_config.velocity U0 must be positive and finite")
    result = {}
    for name, field in zip(("ux", "uy"), velocity):
        value = np.asarray(field)
        if value.shape != (nx, ny) or value.dtype.kind != "f":
            raise ValueError("public Legacy velocity must be floating (nx,ny) q/U0")
        if not np.isfinite(value[fluid_xy]).all():
            raise ValueError("public Legacy fluid velocity contains non-finite values")
        if not np.array_equal(value[~fluid_xy], np.zeros(np.count_nonzero(~fluid_xy), value.dtype)):
            raise ValueError("public Legacy velocity must be exact zero on nonfluid cells")
        result[name] = np.asarray((value * u0).T, dtype=np.float32)
    return result


def _target_sensors(target_states):
    target = np.asarray(target_states)
    if target.ndim != 2:
        raise ValueError("target_states must be two-dimensional")
    if target.shape[1] == 6:
        return target
    if target.shape[1] == 8:
        return target[:, 2:8]
    raise ValueError("target_states must expose exactly six sensor channels")


def _observation_slices(case, role, boundary_obs):
    if role == "target":
        return ((boundary_obs[2:8], np.full(6, np.nan, np.float32))
                if case.scene in ("karman", "illusion") else
                (boundary_obs[:6], np.full(6, np.nan, np.float32)))
    raw = boundary_obs[:12] if case.scene in ("illusion", "steady") else boundary_obs[2:14]
    return raw[:6], raw[6:12]


def _collect_role(role, data, model, ff, policy_norm, scratch: Path,
                  capture_field: Callable = _capture_boundary_field, case_name: str = CASE_NAME):
    case = get_case(case_name)
    if role not in _case_roles(case):
        raise ValueError(f"unknown acquisition role {role!r} for {case.name}")
    config = data["config"]; si = int(case.sample_interval)
    expected_objects = 4 if role == "target" and case.scene in ("karman", "illusion") else (3 if role == "target" else (7 if case.scene == "karman" else 6))
    if (int(config["nx"]), int(config["sample_interval"]), int(config["n_obj_total"])) != (NX, si, expected_objects):
        raise ValueError(f"built Legacy {case.name} {role} lattice/control contract changed")
    if role == "target":
        fifo = obs = None
    else:
        saved = np.asarray(data["norm"]["save_states"], dtype=np.float32)
        obs_dim = 14 if case.scene == "illusion" else 12
        fifo, obs = reset_runtime(ff, saved, obs_dim)
        if not isinstance(fifo, deque) or len(fifo) != len(saved):
            raise ValueError("canonical reset FIFO contract failed")
    rows = []
    field_buffer = None
    total = WARMUP_INTERVALS + COLLECT_BOUNDARIES
    harmonics = data.get("target_harmonics")
    steady = case.scene == "steady"
    for interval in range(total):
        if role == "controlled":
            action, _ = model.predict(obs, deterministic=True)
            action, command = _command(case, action, expected_objects, ff.DATA_TYPE)
        elif role == "zero":
            action = -np.asarray(case.action_bias, np.float32) / np.float32(case.action_scale)
            action, command = _command(case, action, expected_objects, ff.DATA_TYPE)
            if not np.allclose(command[-3:], 0.0, rtol=0.0, atol=np.finfo(ff.DATA_TYPE).eps):
                raise ValueError("Legacy zero role counter-bias did not cancel the affine command")
            command[-3:] = np.zeros(3, dtype=ff.DATA_TYPE)
        elif role == "constant":
            action = np.zeros(3, np.float32)
            action, command = _command(case, action, expected_objects, ff.DATA_TYPE)
        else:
            action = np.full(3, np.nan, np.float32); command = np.zeros(expected_objects, ff.DATA_TYPE)
        _run_interval(ff, si, command)
        boundary_obs = np.asarray(ff.obs, np.float32).copy()
        sensors, forces = _observation_slices(case, role, boundary_obs)
        if role == "target":
            effective = np.full(3, np.nan, np.float32)
            terms = {name: np.nan for name in ("reward", "reward_cd", "reward_cl", "native_legacy_dtw")}
        else:
            raw = np.concatenate((sensors, forces)).astype(np.float32)
            fifo.append(raw)
            if steady:
                terms = {name: np.nan for name in ("reward", "reward_cd", "reward_cl", "native_legacy_dtw")}
            else:
                terms = reward_terms(case, data["target_states"], harmonics, np.asarray(fifo), policy_norm["force_norm_fact"], interval)
            if not steady:
                target_force = gen_target_states_at(interval + 1, harmonics)[:2] if harmonics is not None else None
                obs = policy_observation(raw, policy_norm, target_force=target_force)
            effective = np.asarray(ff.current_effective_action(), np.float32)[-3:].copy()
        if interval < WARMUP_INTERVALS:
            continue
        boundary = interval - WARMUP_INTERVALS
        lattice_step, control_clock, absolute = _clock(ff, interval + 1, si)
        native = float(terms["native_legacy_dtw"])
        row = {"physical_time": lattice_step * U0 / NX, "lattice_step": lattice_step,
               "control_index": control_clock, "sensors": sensors.copy(), "forces": forces.copy(),
               "action_normalized": action.copy(),
               "commanded_target_omega": (np.full(3, np.nan, np.float32) if role == "target" else command[-3:].astype(np.float32, copy=True)),
               "effective_smoothed_omega": effective, "reward_raw": float(terms["reward"]),
               "reward_cd": float(terms["reward_cd"]), "reward_cl": float(terms["reward_cl"]),
               "reward_sim": float(np.exp(-10 * abs(native - 1))) if np.isfinite(native) else np.nan,
               "native_reward_dtw": native}
        if absolute is not None: row["solver_absolute_lattice_step"] = absolute
        rows.append(row)
        if not steady or boundary == COLLECT_BOUNDARIES - 1:
            fields = capture_field(ff, int(config["nx"]), int(config["ny"]))
            if field_buffer is None:
                count = 1 if steady else COLLECT_BOUNDARIES
                field_buffer = _allocate_field_buffer(fields["ux"].shape, count)
            store_at = 0 if steady else boundary
            _store_field_candidate(field_buffer, store_at, fields)
    return rows, field_buffer


def _collect(data, model, ff, policy_norm, scratch: Path, capture_field: Callable = _capture_boundary_field):
    return _collect_role("controlled", data, model, ff, policy_norm, scratch, capture_field=capture_field)


def _allocate_field_buffer(shape_yx, count: int = COLLECT_BOUNDARIES):
    ny, nx = map(int, shape_yx)
    return {
        "ux": np.empty((int(count), ny, nx), dtype=np.float32),
        "uy": np.empty((int(count), ny, nx), dtype=np.float32),
    }


def _store_field_candidate(field_buffer, boundary: int, fields):
    field_buffer["ux"][boundary] = np.asarray(fields["ux"], dtype=np.float32)
    field_buffer["uy"][boundary] = np.asarray(fields["uy"], dtype=np.float32)


def _columns(rows):
    return {name: np.asarray([row[name] for row in rows]) for name in rows[0]}


def _require_periodic_crossings(phase, minimum: int = 4):
    count = len(phase["crossing_times"])
    if count < minimum:
        raise ValueError(f"periodic publication requires at least {minimum} center-uy rising crossings")
    return count


def _phase_values(columns, role="controlled"):
    excluded = {"physical_time", "lattice_step", "solver_absolute_lattice_step",
                "control_index", "phase", "cycle_id"}
    if role == "target":
        return {"sensors": columns["sensors"]}
    return {name: value for name, value in columns.items() if name not in excluded}


def _minimum_crossing_interval(physical_time, cycle_length):
    times = np.asarray(physical_time, dtype=np.float64)
    if len(times) < 2:
        raise ValueError("physical_time must contain at least two boundaries")
    dt = float(np.median(np.diff(times)))
    if not np.isfinite(dt) or dt <= 0:
        raise ValueError("physical_time spacing must be positive and finite")
    return 0.5 * float(cycle_length) * dt


PHASE_FIELD_KEYS = {
    "ux", "uy", "mean_ux", "mean_uy", "field_indices", "cycle_id",
    "target_phase", "actual_phase", "phase_error",
}


def _finalize(role_dir: Path, scratch: Path, rows, field_buffer, data, bundle, storage, role="controlled"):
    case = bundle["case"]; columns = _columns(rows)
    phase_meta = {}
    if case.scene == "steady":
        write_boundary_artifacts(role_dir / "timeseries.npz", role_dir / "timeseries.csv", columns)
        publish_selected_fields(role_dir / "late_field.npz", {
            "ux": np.asarray(field_buffer["ux"][:1], dtype=np.float32),
            "uy": np.asarray(field_buffer["uy"][:1], dtype=np.float32),
            "field_indices": np.asarray([len(rows) - 1], dtype=np.int64),
        })
        phase_values = {}
    else:
        dtw = dual_cycle_dtw(_target_sensors(data["target_states"]), columns["sensors"],
                             columns["native_reward_dtw"], cycle_length=case.conv_len, lag_channel=3)
        columns.update(target_normalized_dtw=dtw["target_normalized_dtw"],
                       target_normalized_dtw_lag=dtw["target_normalized_dtw_lag"])
        min_gap = _minimum_crossing_interval(columns["physical_time"], case.conv_len)
        phase = assign_periodic_phase(
            columns["physical_time"], columns["sensors"], minimum_crossing_interval=min_gap,
        )
        _require_periodic_crossings(phase)
        columns.update(phase=phase["phase"], cycle_id=phase["cycle_id"])
        write_boundary_artifacts(role_dir / "timeseries.npz", role_dir / "timeseries.csv", columns)
        phase_values = _phase_values(columns, role)
        write_phase_cycle_artifacts(role_dir / "phase_cycle.npz", role_dir / "phase_cycle.csv", pooled_phase_bins(phase["phase"], phase_values, bins=32))
        selected = select_phase_fields(columns["physical_time"], phase["crossing_times"])
        mean_span = complete_cycle_field_indices(columns["physical_time"], phase["crossing_times"])
        mean = accumulate_mean_fields(field_buffer["ux"], field_buffer["uy"], mean_span["field_indices"])
        published = {
            "ux": np.asarray(field_buffer["ux"][selected["field_indices"]], dtype=np.float32),
            "uy": np.asarray(field_buffer["uy"][selected["field_indices"]], dtype=np.float32),
            "mean_ux": mean["mean_ux"],
            "mean_uy": mean["mean_uy"],
        }
        published.update(selected)
        field_path = publish_selected_fields(role_dir / "phase_fields.npz", published)
        with np.load(field_path, allow_pickle=False) as saved:
            if (saved["ux"].shape[0] != 8 or set(saved.files) != PHASE_FIELD_KEYS
                    or saved["mean_ux"].shape != saved["ux"].shape[1:]
                    or saved["mean_uy"].shape != saved["uy"].shape[1:]):
                raise ValueError("published phase fields failed exact-key validation")
        reference = _controlled_reference(role, bundle, columns, case.conv_len)
        _atomic_json(role_dir / "dtw_summary.json", {"native_pipeline": "unavailable for target role" if role == "target" else "original Legacy reward DTW",
            "target_normalized_dtw": dtw["metadata"]["definition"], "historical_window": case.conv_len,
            "lag_channel": 3, "target_scale": dtw["metadata"]["scale"].tolist(),
            "native_mean": (None if role == "target" else float(np.nanmean(columns["native_reward_dtw"]))),
            "target_normalized_dtw_finite_mean": float(np.nanmean(columns["target_normalized_dtw"])),
            "crossing_count": int(len(phase["crossing_times"])),
            "frozen_reference_comparison": reference})
        phase_meta = {
            "phase_contract": "smoothed rising crossings of sensors[:,3] with min-gap filter; complete half-open cycles",
            "phase_smoothing_kernel": phase["smoothing_kernel"],
            "minimum_crossing_interval": phase["minimum_crossing_interval"],
            "accepted_crossing_count": phase["accepted_crossing_count"],
            "rejected_crossing_count": phase["rejected_crossing_count"],
            "complete_cycle_count": phase["complete_cycle_count"],
            "mean_field_count": mean_span["mean_field_count"],
            "mean_first_crossing_time": mean_span["first_crossing_time"],
            "mean_last_crossing_time": mean_span["last_crossing_time"],
            "candidate_field_storage": "single-role in-memory FP32 ux/uy buffer; no boundary_*.npz scratch",
        }
    cleanup_scratch(scratch, root=role_dir / "scratch"); scratch.parent.rmdir()
    _atomic_json(role_dir / "metadata.json", {"schema": "drl-pinball-legacy-acquisition-v2", "case_id": case.name, "scene": case.scene, "role": role,
        "warmup_control_steps": WARMUP_INTERVALS, "collected_post_step_boundaries": COLLECT_BOUNDARIES,
        "si": case.sample_interval, "historical_window": case.conv_len, "phase_sensor": "fixed center sensor uy (sensors[:,3])" if case.scene != "steady" else None,
        "field_contract": {"quantity": "q/RHO_ref", "RHO_ref": 1, "source": "public current_step_velocity_field q/U0 multiplied by field_config.velocity U0"},
        "field_count": 1 if case.scene == "steady" else 8, "nonperiodic": case.scene == "steady",
        "role_semantics": _role_semantics(role),
        "timeline_contract": "V5 lattice_step is an absolute solver count; Legacy physical_time is derived only from relative Legacy lattice_step",
        "source_provenance": {
            "model_file": _identity(bundle["model_path"]) if role == "controlled" else None,
            "normalizer_file": (_identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None),
            "frozen_reference_file": (_optional_file_provenance(bundle["reference_path"])
                                      if role == "controlled" else None),
            "target_states": _generated_array_identity(data["target_states"]),
            "builder_config": _generated_config_identity(data["config"]),
        },
        "model": _identity(bundle["model_path"]) if role == "controlled" else None,
        "normalizer": (_identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None),
        "normalization_source": ("frozen_policy" if bundle["norm_path"] is not None else
                                 ("builder_recomputed" if role != "target" else "not_applicable")),
        "scratch_cleanup": {"complete": not scratch.exists(), "path": str(scratch)},
        "resolved_optane_path": str(storage["resolved_output_root"]), "phase_variables": list(phase_values),
        **phase_meta})


def _validate_staged_role(role_dir: Path, case_name: str = CASE_NAME) -> None:
    steady = get_case(case_name).scene == "steady"
    expected = ({"timeseries.npz", "timeseries.csv", "late_field.npz", "metadata.json"} if steady else
                {"timeseries.npz", "timeseries.csv", "phase_cycle.npz", "phase_cycle.csv", "phase_fields.npz", "dtw_summary.json", "metadata.json"})
    actual = {path.name for path in role_dir.iterdir()}
    if actual != expected: raise ValueError(f"staged role files mismatch: expected {sorted(expected)}, got {sorted(actual)}")
    field_name = "late_field.npz" if steady else "phase_fields.npz"
    with np.load(role_dir / field_name, allow_pickle=False) as saved:
        count = 1 if steady else 8
        if saved["ux"].shape[0] != count or saved["uy"].shape[0] != count: raise ValueError("field count contract failed")
        if not steady:
            if (set(saved.files) != PHASE_FIELD_KEYS
                    or saved["mean_ux"].shape != saved["ux"].shape[1:]
                    or saved["mean_uy"].shape != saved["uy"].shape[1:]):
                raise ValueError("phase fields exact schema failed")
    with (role_dir / "metadata.json").open(encoding="utf-8") as stream: json.load(stream)


def _acquire_standard(role, *, case_name=CASE_NAME, output_root=None, overwrite=False, device_id=0,
                 repo_root=None, storage_validator=validate_output_storage,
                 runtime_factory=None, finalizer=_finalize, **unused):
    case = get_case(case_name)
    if role not in _case_roles(case): raise ValueError(f"unknown acquisition role {role!r} for {case_name}")
    repo = Path(repo_root) if repo_root is not None else Path(__file__).resolve().parents[3]
    storage = storage_validator(repo_mapping=default_reproduction_mapping(repo),
                                output_root=output_root, min_free_bytes=MIN_FREE_BYTES)
    if case_name == CASE_NAME and role == "controlled":
        bundle = _resolve_bundle(repo)
    else:
        bundle = _resolve_bundle(repo, case_name, role)
    prepared = prepare_role_output(storage, "legacy", case_name, role, overwrite=overwrite)
    scratch = create_scratch(prepared["scratch_root"])
    ff = None
    try:
        if runtime_factory is None:
            data, model, ff = _default_runtime(bundle["case"], bundle, device_id, role)
        else:
            data, model, ff = runtime_factory(bundle["case"], bundle, device_id, role)
        policy_norm = (None if role == "target" else
                       (data["norm"] if bundle["case"].scene == "steady" else bundle["norm"]))
        rows, fields = _collect_role(role, data, model, ff, policy_norm, scratch, case_name=case_name)
        finalizer(prepared["role_dir"], scratch, rows, fields, data, bundle, storage, role)
        _validate_staged_role(prepared["role_dir"], case_name)
        return publish_role_output(prepared)
    except Exception:
        if prepared["staging_dir"].exists():
            shutil.rmtree(prepared["staging_dir"])
        raise
    finally:
        if ff is not None and hasattr(ff, "close"):
            ff.close()



def _vortex_offset(case_name, value):
    if type(value) is not int:
        raise ValueError("vortex y offset must be an integer L0 multiple")
    allowed = {0} if case_name == "vortex_lamb" else {-2, -1, 0, 1, 2}
    if value not in allowed:
        raise ValueError(f"{case_name} vortex y offset must be one of {sorted(allowed)}")
    return value

def _vortex_scenario(case_name, offset):
    _vortex_offset(case_name, offset)
    label = {0:"y000", -2:"ym2L", -1:"ym1L", 1:"yp1L", 2:"yp2L"}[offset]
    return f"{case_name}_{label}"

def _vortex_event_channel(case_name):
    if case_name == "vortex_lamb":
        return 1, "upper_sensor_uy"
    if case_name == "vortex_taylor":
        return 3, "center_sensor_uy"
    raise ValueError(f"unknown Vortex case {case_name!r}")


def _select_vortex_event(sensors, case_name):
    sensors = np.asarray(sensors, np.float64)
    if sensors.ndim != 2 or sensors.shape[1:] != (6,) or len(sensors) != 150 or not np.all(np.isfinite(sensors)):
        raise ValueError("Vortex event selection requires finite sensors shape (150,6)")
    channel, channel_name = _vortex_event_channel(case_name)
    uy = sensors[:, channel]
    baseline = float(np.median(uy[:20]))
    score = np.abs(uy - baseline)
    lower, upper = 38, 87
    peak = lower + int(np.argmax(score[lower:upper + 1]))
    indices = peak + np.asarray([-10, -5, 0, 5, 10], np.int64)
    return {"baseline_uy": baseline, "event_score": score,
            "selected_event_score": score[indices], "peak_index": np.int64(peak),
            "field_indices": indices,
            "relative_offsets": np.asarray([-10, -5, 0, 5, 10], np.int64),
            "search_lower_boundary": lower, "search_upper_boundary": upper,
            "event_channel_index": channel, "event_channel": channel_name}


def _canonical_vortex_source(case_name):
    return "vortex_lamb_y000" if case_name == "vortex_lamb" else "vortex_taylor_y000"


def _load_vortex_event_source(storage, case_name):
    source = Path(storage["resolved_output_root"]) / "legacy" / _canonical_vortex_source(case_name) / "controlled" / "event_summary.json"
    if not source.is_file():
        raise FileNotFoundError(f"canonical Vortex event source is required before this run: {source}")
    try:
        with source.open(encoding="utf-8") as stream:
            summary = json.load(stream)
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(f"malformed canonical Vortex event source {source}") from exc
    channel, channel_name = _vortex_event_channel(case_name)
    expected_source = {"scenario": _canonical_vortex_source(case_name), "role": "controlled",
                       "channel_index": channel, "channel": channel_name}
    required = {"schema", "event_source", "common_peak_index", "field_indices", "relative_offsets"}
    if not isinstance(summary, dict) or not required.issubset(summary):
        raise ValueError(f"malformed canonical Vortex event source {source}")
    try:
        peak = int(summary["common_peak_index"])
        indices = np.asarray(summary["field_indices"], np.int64)
        offsets = np.asarray(summary["relative_offsets"], np.int64)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"malformed canonical Vortex event source {source}") from exc
    if (summary["schema"] != "drl-pinball-legacy-vortex-event-v2" or
            summary["event_source"] != expected_source or indices.shape != (5,) or
            not np.array_equal(offsets, [-10, -5, 0, 5, 10]) or
            not np.array_equal(indices, peak + offsets) or peak < 38 or peak > 87):
        raise ValueError(f"malformed canonical Vortex event source {source}")
    return {"source_path": source, "event_source": expected_source,
            "peak_index": peak, "field_indices": indices, "relative_offsets": offsets}

def _erase_phase(times, center_uy, disturbance_force_y):
    try:
        phase = _assign_scalar_phase(times, center_uy)
        if len(phase["crossing_times"]) >= 4:
            return phase, "center_sensor_uy", False
    except ValueError as exc:
        if "rising crossings" not in str(exc): raise
    phase = _assign_scalar_phase(times, disturbance_force_y)
    if len(phase["crossing_times"]) < 4:
        raise ValueError("Erase publication requires at least 4 rising crossings in fixed disturbance-force fallback")
    return phase, "disturbance_force_y", True

def _assign_scalar_phase(times, probe):
    probe=np.asarray(probe,np.float64); times=np.asarray(times,np.float64)
    if probe.shape != times.shape or probe.ndim != 1 or not np.all(np.isfinite(probe)):
        raise ValueError("phase probe must be one-dimensional, aligned, and finite")
    synthetic=np.zeros((len(times),4),np.float64); synthetic[:,3]=probe
    return assign_phase(times,synthetic)

def _special_observation(case, role, raw):
    raw=np.asarray(raw,np.float32)
    if case.scene == "erase":
        if role == "target": return raw[:6], np.full(2,np.nan,np.float32), np.full(6,np.nan,np.float32), raw[:6]
        if raw.shape != (14,): raise ValueError("Erase observation must contain exactly 14 channels")
        return raw[:6], raw[6:8], raw[8:14], raw
    if raw.shape != (6 if role == "target" else 12,): raise ValueError("Vortex observation inventory changed")
    return raw[:6], np.full(2,np.nan,np.float32), (np.full(6,np.nan,np.float32) if role=="target" else raw[6:12]), raw

def _special_schedule(case, target):
    if case.scene == "erase":
        return (0 if target else 480), 160
    return 0, 150


def _special_timeline_metadata(case, target):
    warmup, collected = _special_schedule(case, target)
    return {"warmup_control_steps": warmup,
            "collected_post_step_boundaries": collected,
            "historical_episode_steps": 150 if case.scene == "vortex" else None}


def _collect_special(role,data,model,ff,norm,scratch,case,capture_field=_capture_boundary_field):
    target = role == "target"; expected = 3 if target else (7 if case.scene=="erase" else 6)
    if int(data["config"]["n_obj_total"]) != expected: raise ValueError("special scene object inventory changed")
    if target: fifo=None; obs=None
    else:
        saved=np.asarray(data["norm"]["save_states"],np.float32)
        fifo,obs=reset_runtime(ff,saved,14 if case.scene=="erase" else 12)
        if case.scene == "erase":
            if saved.ndim != 2 or saved.shape[1] != 14:
                raise ValueError("Erase reset FIFO must retain exactly 14 raw channels")
            obs=np.zeros(12,dtype=np.float32)
    warmup,collect=_special_schedule(case,target)
    rows=[]; fields=[]
    for interval in range(warmup+collect):
        if role=="controlled": action,_=model.predict(obs,deterministic=True); action,command=_command(case,action,expected,ff.DATA_TYPE)
        elif role=="zero":
            action=-np.asarray(case.action_bias,np.float32)/np.float32(case.action_scale); action,command=_command(case,action,expected,ff.DATA_TYPE); command[-3:]=0
        else: action=np.full(3,np.nan,np.float32); command=np.zeros(expected,ff.DATA_TYPE)
        _run_interval(ff,case.sample_interval,command); sensors,disturbance,pinball,raw=_special_observation(case,role,np.asarray(ff.obs,np.float32).copy())
        if target: terms={}; effective=np.full(3,np.nan,np.float32)
        else:
            fifo.append(raw.copy())
            if case.scene=="erase": terms=erase_reward_terms(np.asarray(fifo),np.mean(data["target_states"],axis=0),norm,case.conv_len); obs=erase_policy_observation(raw,norm)
            else: terms=reward_terms(case,data["target_states"],None,np.asarray(fifo),norm["force_norm_fact"],interval); obs=policy_observation(raw,norm)
            effective=np.asarray(ff.current_effective_action(),np.float32)[-3:].copy()
        if interval < warmup: continue
        lattice,clock,absolute=_clock(ff,interval+1,case.sample_interval)
        row={"physical_time":lattice*U0/NX,"lattice_step":lattice,"control_index":clock,"sensors":sensors.copy(),
             "disturbance_force":disturbance.copy(),"pinball_forces":pinball.copy(),"action_normalized":action.copy(),
             "commanded_target_omega":np.full(3,np.nan,np.float32) if target else command[-3:].astype(np.float32),
             "effective_smoothed_omega":effective,"reward_raw":float(terms.get("reward",np.nan)),
             "reward_u":float(terms.get("reward_u",np.nan)),"reward_v":float(terms.get("reward_v",np.nan)),
             "native_reward_dtw":float(terms.get("native_legacy_dtw",np.nan)),"native_lag":float(terms.get("native_lag",np.nan)),
             "native_component_similarity":np.asarray(terms.get("native_component_similarity",[np.nan,np.nan]))}
        if absolute is not None: row["solver_absolute_lattice_step"]=absolute
        rows.append(row); value=capture_field(ff,int(data["config"]["nx"]),int(data["config"]["ny"])); path=scratch/f"boundary_{len(rows)-1:04d}.npz"; np.savez(path,**value); fields.append(path)
    return rows,fields

def _publish_fields(path,field_paths,selection):
    result={"ux":[],"uy":[]}; indices=np.asarray(selection["field_indices"],np.int64)
    for index in indices:
        with np.load(field_paths[int(index)],allow_pickle=False) as item:
            result["ux"].append(item["ux"]); result["uy"].append(item["uy"])
    result={key:np.stack(value) for key,value in result.items()}; result.update(selection); publish_selected_fields(path,result)

def _special_normalization_source(case, role):
    if role == "target":
        return "not_applicable"
    if case.scene == "erase":
        return "builder_recomputed_historical"
    if case.scene == "vortex" and role == "zero":
        return "builder_recomputed_zero_runtime"
    return "frozen_policy"


def _special_normalizer_file(bundle):
    return _identity(bundle["norm_path"]) if bundle["norm_path"] is not None else None


def _finalize_special(role_dir,scratch,rows,field_paths,data,bundle,storage,role,event_source=None):
    case=bundle["case"]; columns=_columns(rows); target=role=="target"; event_peak_index=None; event_diagnostics=None
    if case.scene=="vortex":
        channel, channel_name = _vortex_event_channel(case.name)
        own_uy = np.asarray(columns["sensors"][:, channel], np.float64)
        own_baseline = float(np.median(own_uy[:20]))
        own_score = np.abs(own_uy - own_baseline)
        columns["event_score"] = own_score
        if event_source is None:
            local = _select_vortex_event(columns["sensors"], case.name)
            event_source = {"event_source": {"scenario": _canonical_vortex_source(case.name),
                            "role": "controlled", "channel_index": channel, "channel": channel_name},
                            "peak_index": int(local["peak_index"]),
                            "field_indices": local["field_indices"],
                            "relative_offsets": local["relative_offsets"]}
        indices = np.asarray(event_source["field_indices"], np.int64)
        offsets = np.asarray(event_source["relative_offsets"], np.int64)
        event_selection = {"field_indices": indices, "relative_offsets": offsets,
                           "selected_event_score": own_score[indices]}
        write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns)
        _publish_fields(role_dir/"event_fields.npz",field_paths,event_selection)
        event_peak_index = int(event_source["peak_index"])
        event_diagnostics = {"search_lower_boundary": 38, "search_upper_boundary": 87,
                             "diagnostic_channel_index": channel, "diagnostic_channel": channel_name}
        _atomic_json(role_dir/"event_summary.json", {"schema":"drl-pinball-legacy-vortex-event-v2",
            "definition":"canonical controlled fixed-channel absolute deviation from median first 20; earliest argmax in inclusive boundaries 38..87",
            "event_source":event_source["event_source"], "common_peak_index":event_peak_index,
            "field_indices":indices.tolist(), "relative_offsets":offsets.tolist(),
            "selected_event_score":own_score[indices].tolist(), "diagnostic_baseline_uy":own_baseline,
            "diagnostic_peak_score":float(own_score[event_peak_index]), **event_diagnostics})
        field_count=5; nonperiodic=True; probe=None
    elif target:
        write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns); _publish_fields(role_dir/"late_field.npz",field_paths,{"field_indices":np.asarray([len(rows)-1],np.int64)}); field_count=1; nonperiodic=True; probe=None
        err=columns["sensors"]-np.mean(data["target_states"],axis=0); _atomic_json(role_dir/"erase_summary.json",{"native_metrics":"unavailable for clean target","clean_sensor_rmse":float(np.sqrt(np.mean(err*err))),"phase_probe":None,"nonperiodic_clean_reference":True})
    else:
        phase,probe,fallback=_erase_phase(columns["physical_time"],columns["sensors"][:,3],columns["disturbance_force"][:,1])
        columns.update(phase=phase["phase"],cycle_id=phase["cycle_id"]); write_boundary_artifacts(role_dir/"timeseries.npz",role_dir/"timeseries.csv",columns)
        values=_phase_values(columns,role); write_phase_cycle_artifacts(role_dir/"phase_cycle.npz",role_dir/"phase_cycle.csv",pooled_phase_bins(phase["phase"],values,bins=32)); _publish_fields(role_dir/"phase_fields.npz",field_paths,select_phase_fields(columns["physical_time"],phase["crossing_times"])); field_count=8; nonperiodic=False
        err=columns["sensors"]-np.mean(data["target_states"],axis=0); _atomic_json(role_dir/"erase_summary.json",{"native_reward_mean":float(np.mean(columns["reward_raw"])),"reward_u_mean":float(np.mean(columns["reward_u"])),"reward_v_mean":float(np.mean(columns["reward_v"])),"native_reward_dtw_mean":float(np.mean(columns["native_reward_dtw"])),"clean_sensor_rmse":float(np.sqrt(np.mean(err*err))),"phase_probe":probe,"phase_fallback_used":fallback,"force_balance_diagnostic":"native enhanced formula is scale invariant; normalized diagnostic equals native by definition"})
    cleanup_scratch(scratch,root=role_dir/"scratch"); scratch.parent.rmdir()
    timeline=_special_timeline_metadata(case,target)
    _atomic_json(role_dir/"metadata.json",{"schema":"drl-pinball-legacy-nonstandard-acquisition-v1","case_id":case.name,"scene":case.scene,"role":role,**timeline,"field_count":field_count,"nonperiodic":nonperiodic,"phase_probe":probe,"event_peak_index":event_peak_index,"event_source":(event_source["event_source"] if event_source else None),"event_diagnostics":event_diagnostics,"normalization_source":_special_normalization_source(case,role),"role_semantics":(("physical-zero/uncontrolled trajectory using counter-bias normalized action; builder-recomputed historical norm" if role=="zero" else "frozen-policy controlled trajectory using builder-recomputed historical norm") if case.scene=="erase" and not target else ("physical-zero/uncontrolled trajectory using counter-bias normalized action; builder-recomputed zero-runtime norm" if case.scene=="vortex" and role=="zero" else _role_semantics(role))),"field_contract":{"quantity":"q/RHO_ref","RHO_ref":1},"source_provenance":{"model_file":_identity(bundle["model_path"]) if role=="controlled" else None,"normalizer_file":_special_normalizer_file(bundle),"target_states":_generated_array_identity(data["target_states"]),"builder_config":_generated_config_identity(data["config"])},"resolved_optane_path":str(storage["resolved_output_root"])})

def _validate_special(role_dir,case,role):
    if case.scene=="vortex": expected={"timeseries.npz","timeseries.csv","event_fields.npz","event_summary.json","metadata.json"}; field="event_fields.npz"; count=5
    elif role=="target": expected={"timeseries.npz","timeseries.csv","late_field.npz","erase_summary.json","metadata.json"}; field="late_field.npz"; count=1
    else: expected={"timeseries.npz","timeseries.csv","phase_cycle.npz","phase_cycle.csv","phase_fields.npz","erase_summary.json","metadata.json"}; field="phase_fields.npz"; count=8
    actual={x.name for x in role_dir.iterdir()}
    if actual!=expected: raise ValueError(f"staged role files mismatch: expected {sorted(expected)}, got {sorted(actual)}")
    with np.load(role_dir/field,allow_pickle=False) as saved:
        expected_keys = ({"ux", "uy", "field_indices", "relative_offsets", "selected_event_score"}
                         if case.scene == "vortex" else None)
        if expected_keys is not None and set(saved.files) != expected_keys:
            raise ValueError(f"Vortex event field keys mismatch: expected {sorted(expected_keys)}, got {sorted(saved.files)}")
        if any(np.asarray(saved[name]).ndim < 1 or np.asarray(saved[name]).shape[0] != count
               for name in saved.files):
            raise ValueError("special field arrays must be non-scalar with exact field count")
        if case.scene == "vortex":
            field_indices = np.asarray(saved["field_indices"], np.int64)
            relative_offsets = np.asarray(saved["relative_offsets"], np.int64)
    if case.scene == "vortex":
        try:
            with (role_dir/"event_summary.json").open(encoding="utf-8") as stream:
                summary = json.load(stream)
            channel, channel_name = _vortex_event_channel(case.name)
            expected_source = {"scenario": _canonical_vortex_source(case.name), "role": "controlled",
                               "channel_index": channel, "channel": channel_name}
            peak = int(summary["common_peak_index"])
            if (summary["schema"] != "drl-pinball-legacy-vortex-event-v2" or
                    summary["event_source"] != expected_source or
                    not np.array_equal(field_indices, peak + relative_offsets) or
                    not np.array_equal(field_indices, np.asarray(summary["field_indices"], np.int64)) or
                    not np.array_equal(relative_offsets, [-10, -5, 0, 5, 10])):
                raise ValueError
        except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
            raise ValueError("Vortex event summary/source contract failed") from exc

def _acquire_special(role,case_name,output_root,overwrite,device_id,repo_root,storage_validator,runtime_factory,vortex_y_offset_l0):
    case=get_case(case_name); repo=Path(repo_root) if repo_root else Path(__file__).resolve().parents[3]; offset=_vortex_offset(case_name,vortex_y_offset_l0) if case.scene=="vortex" else 0
    storage=storage_validator(repo_mapping=default_reproduction_mapping(repo),output_root=output_root,min_free_bytes=MIN_FREE_BYTES); bundle=_resolve_bundle(repo,case_name,role)
    scenario=_vortex_scenario(case_name,offset) if case.scene=="vortex" else case_name
    source_run = case.scene == "vortex" and role == "controlled" and offset == 0
    event_source = None if not case.scene == "vortex" or source_run else _load_vortex_event_source(storage, case_name)
    prepared=prepare_role_output(storage,"legacy",scenario,role,overwrite=overwrite); scratch=create_scratch(prepared["scratch_root"]); ff=None
    try:
        if runtime_factory is not None: data,model,ff=runtime_factory(case,bundle,device_id,role)
        else:
            from .core.legacy_env_builder import build_erase,build_vortex
            data=build_vortex(device_id,case.re_code,vortex_type=case.vortex_type,action_scale=case.action_scale,y_offset_l0=offset,target_only=role=="target",role=role) if case.scene=="vortex" else build_erase(device_id,case.re_code,target_only=role=="target")
            model=None
            if role=="controlled":
                from .core.model_loader import load_model
                model=load_model(case.model)
            ff=data["flow_field"]
        norm=None if role=="target" else (data["norm"] if case.scene=="erase" or (case.scene=="vortex" and role=="zero") else bundle["norm"]); rows,fields=_collect_special(role,data,model,ff,norm,scratch,case); _finalize_special(prepared["role_dir"],scratch,rows,fields,data,bundle,storage,role,event_source); _validate_special(prepared["role_dir"],case,role); return publish_role_output(prepared)
    except Exception:
        if prepared["staging_dir"].exists(): shutil.rmtree(prepared["staging_dir"])
        raise
    finally:
        if ff is not None and hasattr(ff,"close"): ff.close()

def acquire_role(role, *, case_name=CASE_NAME, output_root=None, overwrite=False, device_id=0, repo_root=None, storage_validator=validate_output_storage, runtime_factory=None, finalizer=_finalize, vortex_y_offset_l0=0):
    case=get_case(case_name)
    if role not in _case_roles(case): raise ValueError(f"unknown acquisition role {role!r} for {case_name}")
    if case.scene in ("vortex","erase"):
        return _acquire_special(role,case_name,output_root,overwrite,device_id,repo_root,storage_validator,runtime_factory,vortex_y_offset_l0)
    return _acquire_standard(role,case_name=case_name,output_root=output_root,overwrite=overwrite,device_id=device_id,repo_root=repo_root,storage_validator=storage_validator,runtime_factory=runtime_factory,finalizer=finalizer)

def acquire_controlled(**kwargs):
    return acquire_role("controlled", **kwargs)


def main() -> int:
    parser = argparse.ArgumentParser(description="Acquire Legacy full-matrix evidence")
    parser.add_argument("case", choices=SUPPORTED_CASES)
    parser.add_argument("--role", choices=ROLES + ("constant",), required=True)
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--vortex-y-offset-l0", type=int, default=0)
    parser.add_argument("--output-root", type=Path); parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    path = acquire_role(args.role, case_name=args.case, output_root=args.output_root, overwrite=args.overwrite, device_id=args.device_id, vortex_y_offset_l0=args.vortex_y_offset_l0)
    print(path); return 0


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