Blame for scripts/h2_live_sim_console.py

#!/usr/bin/env python3
"""Interactive terminal remote for the live H2 MuJoCo session.

This is intentionally only a UI around ``h2_live_sim_control.py``.  It owns no
shared memory and implements no separate L1/R1 policy, so the one-shot CLI and
the interactive console always have the same behaviour.
"""

from __future__ import annotations

import argparse
from contextlib import redirect_stderr, redirect_stdout
import curses
from dataclasses import dataclass
import io
import math
from pathlib import Path
import sys
import time


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from runtime.planner_protocol import (  # noqa: E402
    BALL_SOURCE_MONOTONIC_INDEX,
    ROBOT_SOURCE_MONOTONIC_INDEX,
    source_age_s,
)


@dataclass(frozen=True)
class ConsoleSnapshot:
    run_dir: Path
    shm_name: str
    phase: str
    state_ready: bool
    planner_ready: bool
    state_age_s: float
    planner_age_s: float
    robot_age_s: float
    ball_age_s: float
    state_limit_s: float
    planner_limit_s: float
    robot_limit_s: float
    ball_limit_s: float
    start_requested: bool
    stop_requested: bool
    emergency_stop: bool
    l1_available: bool
    l1_reason: str
    policy_ticks: int
    active_ticks: int
    accepted_id: int
    reference_id: int
    root_z: float | None
    max_tilt: float | None
    max_dq: float | None
    max_torque_util: float | None
    stop_reason: str


def _number(value: object) -> float | None:
    try:
        result = float(value)
    except (TypeError, ValueError):
        return None
    return result if math.isfinite(result) else None


def collect_snapshot(run_dir: Path) -> ConsoleSnapshot:
    """Read one coherent-enough UI snapshot without changing simulator state."""
    # Import lazily so the plain dashboard renderer remains independently
    # testable on machines that do not have the deployed ROS Python stack.
    from scripts import h2_live_sim_control as control

    _, state = control._read_state(run_dir)
    runner = control._load_json(run_dir / "runner_status.json")
    plant = control._load_json(run_dir / "bridge_status.json")
    settings = control._settings_for_run(run_dir)
    now_ns = time.monotonic_ns()
    now_s = time.monotonic()

    def receipt_age(stamp: object) -> float:
        value = int(stamp)
        return float("inf") if value <= 0 else (now_ns - value) * 1.0e-9

    pelvis = plant.get("pelvis_pos")
    root_z = pelvis[2] if isinstance(pelvis, list) and len(pelvis) >= 3 else None
    try:
        control._require_frozen_and_fresh(run_dir, state)
        l1_available, l1_reason = True, "FROZEN_WAIT and inputs are fresh"
    except RuntimeError as error:
        l1_available, l1_reason = False, str(error).removeprefix("L1 rejected: ")

    return ConsoleSnapshot(
        run_dir=run_dir,
        shm_name=control._shm_name(run_dir),
        phase=str(runner.get("phase", "BOOT")),
        state_ready=bool(state.state_ready),
        planner_ready=bool(state.tt_ready),
        state_age_s=receipt_age(state.state_stamp_ns),
        planner_age_s=receipt_age(state.tt_stamp_ns),
        robot_age_s=source_age_s(
            now_s, state.tt_data[ROBOT_SOURCE_MONOTONIC_INDEX]
        ),
        ball_age_s=source_age_s(
            now_s, state.tt_data[BALL_SOURCE_MONOTONIC_INDEX]
        ),
        state_limit_s=float(settings.max_control_gap_s),
        planner_limit_s=0.10,
        robot_limit_s=float(settings.robot_mocap_timeout_s),
        ball_limit_s=float(settings.ball_mocap_timeout_s),
        start_requested=bool(state.start_requested),
        stop_requested=bool(state.stop_requested),
        emergency_stop=bool(state.emergency_stop),
        l1_available=l1_available,
        l1_reason=l1_reason,
        policy_ticks=int(runner.get("policy_ticks", 0)),
        active_ticks=int(runner.get("active_reference_policy_ticks", 0)),
        accepted_id=int(runner.get("accepted_command_id", 0)),
        reference_id=int(runner.get("reference_command_id", 0)),
        root_z=_number(root_z),
        max_tilt=_number(plant.get("max_tilt_rad")),
        max_dq=_number(plant.get("max_abs_joint_velocity_rad_s")),
        max_torque_util=_number(plant.get("max_torque_utilization")),
        stop_reason=str(runner.get("stop_reason") or ""),
    )


def _age(value: float, limit: float, ready: bool = True) -> tuple[str, bool]:
    if not ready or not math.isfinite(value):
        return "  --  ", False
    return f"{value * 1000:6.1f} ms", -0.01 <= value <= limit


def _metric(value: float | None, suffix: str, digits: int = 3) -> str:
    return "--" if value is None else f"{value:.{digits}f}{suffix}"


def dashboard_lines(
    snapshot: ConsoleSnapshot, message: str = "", *, compact: bool = False
) -> list[str]:
    """Return the terminal dashboard as plain lines (also used by ``--once``)."""
    state_age, state_ok = _age(
        snapshot.state_age_s, snapshot.state_limit_s, snapshot.state_ready
    )
    planner_age, planner_ok = _age(
        snapshot.planner_age_s, snapshot.planner_limit_s, snapshot.planner_ready
    )
    robot_age, robot_ok = _age(snapshot.robot_age_s, snapshot.robot_limit_s)
    ball_age, ball_ok = _age(snapshot.ball_age_s, snapshot.ball_limit_s)
    health = "READY" if all((state_ok, planner_ok, robot_ok, ball_ok)) else "WAIT"
    l1 = "READY - press L" if snapshot.l1_available else snapshot.l1_reason
    flags = (
        f"start={int(snapshot.start_requested)}  stop={int(snapshot.stop_requested)}  "
        f"fault={int(snapshot.emergency_stop)}"
    )
    if compact:
        return [
            f"H2 SIM REMOTE | {snapshot.phase} | {health}",
            f"L: L1 / POLICY | {l1}",
            "R: R1 / STOP (latch)",
            "Q: QUIT after R1 | Esc/Ctrl-C: close UI only",
            f"state {state_age} {'OK' if state_ok else 'STALE'} | "
            f"planner {planner_age} {'OK' if planner_ok else 'STALE'}",
            f"robot {robot_age} {'OK' if robot_ok else 'STALE'} | "
            f"ball {ball_age} {'OK' if ball_ok else 'STALE'}",
            f"policy={snapshot.policy_ticks} active={snapshot.active_ticks} "
            f"id={snapshot.accepted_id}/{snapshot.reference_id}",
            f"z={_metric(snapshot.root_z, 'm')} tilt={_metric(snapshot.max_tilt, '')} "
            f"dq={_metric(snapshot.max_dq, '')} torque={_metric(snapshot.max_torque_util, 'x')}",
            flags,
            f"> {message}",
        ]
    lines = [
        "+---------------- H2 MUJOCO SIMULATED REMOTE ----------------+",
        f"  SIMULATION ONLY / no LowCmd        phase: {snapshot.phase:<16} {health}",
        "+---------------------- LIVE INPUTS -------------------------+",
        f"  MuJoCo state   {state_age}  {'OK' if state_ok else 'STALE'}",
        f"  Planner ROS    {planner_age}  {'OK' if planner_ok else 'STALE'}",
        f"  Robot mocap    {robot_age}  {'OK' if robot_ok else 'STALE'}",
        f"  Ball mocap     {ball_age}  {'OK' if ball_ok else 'STALE'}",
        "+---------------- POLICY / PLANT -----------------------------+",
        f"  ticks policy={snapshot.policy_ticks:<8} active={snapshot.active_ticks:<8} "
        f"command accepted/reference={snapshot.accepted_id}/{snapshot.reference_id}",
        f"  pelvis z={_metric(snapshot.root_z, ' m')}  "
        f"tilt={_metric(snapshot.max_tilt, ' rad')}  "
        f"max dq={_metric(snapshot.max_dq, ' rad/s')}  "
        f"torque={_metric(snapshot.max_torque_util, ' x')}",
        f"  {flags}",
        "+---------------- SIMULATED REMOTE ---------------------------+",
        f"  [L] L1 / enter policy : {l1}",
        "  [R] R1 / stop policy  : latch simulated stop",
        "  [Q] quit simulation   : available after R1",
        "  [Esc or Ctrl-C] close this console only; simulation keeps running",
    ]
    if snapshot.stop_reason:
        lines.append(f"  stop reason: {snapshot.stop_reason}")
    lines.extend(("+-------------------------------------------------------------+", f"  {message}"))
    return lines


def dispatch_key(key: int, run_dir: Path) -> tuple[bool, str]:
    """Run one existing control command; return whether it was handled."""
    command = {
        ord("l"): "l1",
        ord("L"): "l1",
        ord("r"): "r1",
        ord("R"): "r1",
        ord("q"): "quit",
        ord("Q"): "quit",
    }.get(key)
    if command is None:
        return False, ""
    output = io.StringIO()
    with redirect_stdout(output), redirect_stderr(output):
        result = _control_main([command, "--run-dir", str(run_dir)])
    text = " ".join(output.getvalue().strip().splitlines())
    return True, (text or f"{command}: exit {result}")


def _control_main(argv: list[str]) -> int:
    from scripts.h2_live_sim_control import main

    return main(argv)


def _draw(screen, lines: list[str]) -> None:
    screen.erase()
    height, width = screen.getmaxyx()
    for row, line in enumerate(lines[: max(0, height - 1)]):
        style = curses.A_NORMAL
        if "phase:" in line or "SIMULATED REMOTE" in line:
            style = curses.A_BOLD
        if " STALE" in line or "fault=1" in line or line.startswith("  [FAIL]"):
            style |= curses.color_pair(2)
        elif " OK" in line or "READY" in line:
            style |= curses.color_pair(1)
        try:
            screen.addnstr(row, 0, line, max(0, width - 1), style)
        except curses.error:
            pass
    screen.refresh()


def _interactive(screen, run_dir: Path, refresh_s: float) -> None:
    try:
        curses.curs_set(0)
    except curses.error:
        pass
    curses.use_default_colors()
    if curses.has_colors():
        curses.init_pair(1, curses.COLOR_GREEN, -1)
        curses.init_pair(2, curses.COLOR_RED, -1)
    screen.keypad(True)
    screen.timeout(max(100, int(refresh_s * 1000)))
    message = "Press L for simulated L1, R for R1, Q to quit after R1."
    while True:
        try:
            snapshot = collect_snapshot(run_dir)
            height, width = screen.getmaxyx()
            lines = dashboard_lines(
                snapshot, message, compact=height < 19 or width < 70
            )
        except (OSError, RuntimeError, ValueError) as error:
            lines = [
                "H2 MUJOCO SIMULATED REMOTE",
                "",
                f"Session state unavailable: {error}",
                "",
                "Esc or Ctrl-C closes this console only.",
            ]
        _draw(screen, lines)
        key = screen.getch()
        if key in (27, 3):
            return
        handled, result = dispatch_key(key, run_dir)
        if handled:
            message = result
        elif key not in (-1, curses.KEY_RESIZE):
            message = "Keys: L=L1, R=R1, Q=quit after R1, Esc=close console only"


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-dir", default=None)
    parser.add_argument("--refresh-s", type=float, default=0.25)
    parser.add_argument(
        "--once", action="store_true", help="print one dashboard without curses"
    )
    args = parser.parse_args(argv)
    if not 0.10 <= args.refresh_s <= 5.0:
        parser.error("--refresh-s must be between 0.10 and 5.0")
    try:
        from scripts.h2_live_sim_control import _resolve_run_dir

        run_dir = _resolve_run_dir(args.run_dir)
        if args.once:
            print("\n".join(dashboard_lines(collect_snapshot(run_dir))))
            return 0
        if not (sys.stdin.isatty() and sys.stdout.isatty()):
            raise RuntimeError("interactive console requires a TTY; use --once for text")
        curses.wrapper(_interactive, run_dir, args.refresh_s)
        return 0
    except KeyboardInterrupt:
        # Ctrl-C is deliberately a UI-only operation.  It creates no control file.
        return 0
    except (OSError, RuntimeError, ValueError) as error:
        print(f"[FAIL] {error}", file=sys.stderr)
        return 1


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