#!/usr/bin/env python3
"""Print a compact comparison of active Dropbear convergence lanes."""

from __future__ import annotations

import argparse
from pathlib import Path

from tensorboard.backend.event_processing.event_accumulator import EventAccumulator


WORKSPACE_ROOT = Path(__file__).resolve().parents[1]
LOG_ROOT = WORKSPACE_ROOT / "logs" / "rsl_rl" / "dropbear_velocity"

METRICS = (
    ("reward", "Train/mean_reward"),
    ("ep_len", "Train/mean_episode_length"),
    ("lin_track", "Episode_Reward/track_lin_vel_xy"),
    ("vel_err", "Metrics/base_velocity/error_vel_xy"),
    ("gait", "Episode_Reward/gait"),
    ("cmd_lvl", "Curriculum/lin_vel_cmd_levels"),
    ("act_std", "Policy/mean_std"),
    ("bad_orient", "Episode_Termination/bad_orientation"),
    ("fps", "Perf/total_fps"),
)


def latest_values(event_file: Path) -> tuple[int, dict[str, float]]:
    accumulator = EventAccumulator(str(event_file), size_guidance={"scalars": 0})
    accumulator.Reload()
    values: dict[str, float] = {}
    step = -1
    scalar_tags = set(accumulator.Tags().get("scalars", []))
    for label, tag in METRICS:
        if tag not in scalar_tags:
            continue
        events = accumulator.Scalars(tag)
        if events:
            values[label] = events[-1].value
            step = max(step, events[-1].step)
    return step, values


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--pattern",
        default="*converge_*",
        help="Run-directory glob below logs/rsl_rl/dropbear_velocity",
    )
    args = parser.parse_args()

    rows = []
    for run_dir in sorted(LOG_ROOT.glob(args.pattern)):
        event_files = sorted(run_dir.glob("events.out*"), key=lambda path: path.stat().st_mtime)
        if not event_files:
            rows.append((run_dir.name, -1, {}))
            continue
        step, values = latest_values(event_files[-1])
        rows.append((run_dir.name, step, values))

    if not rows:
        raise SystemExit(f"No runs matched {LOG_ROOT / args.pattern}")

    columns = ("run", "iter", *(label for label, _ in METRICS))
    widths = {
        column: max(
            len(column),
            max(
                (
                    len(run)
                    if column == "run"
                    else len(str(step))
                    if column == "iter"
                    else len(f"{values.get(column, float('nan')):.3f}")
                )
                for run, step, values in rows
            ),
        )
        for column in columns
    }
    print("  ".join(column.ljust(widths[column]) for column in columns))
    print("  ".join("-" * widths[column] for column in columns))
    for run, step, values in rows:
        cells = [run.ljust(widths["run"]), str(step).rjust(widths["iter"])]
        for label, _ in METRICS:
            value = values.get(label)
            cells.append(("—" if value is None else f"{value:.3f}").rjust(widths[label]))
        print("  ".join(cells))

    print(
        "\nHealthy direction: reward/episode length/lin_track/cmd_lvl rise; "
        "bad_orient falls; act_std remains non-zero."
    )


if __name__ == "__main__":
    main()