#!/usr/bin/env python3
"""Analyze sustained Dropbear learning trends from TensorBoard events."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
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"

TAGS = {
    "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 scalar_series(accumulator: EventAccumulator, tag: str) -> tuple[np.ndarray, np.ndarray]:
    events = accumulator.Scalars(tag)
    steps = np.asarray([event.step for event in events], dtype=np.float64)
    values = np.asarray([event.value for event in events], dtype=np.float64)
    return steps, values


def summarize(values: np.ndarray, steps: np.ndarray, window: int) -> tuple[float, float, float]:
    count = min(window, values.size)
    recent_values = values[-count:]
    recent_steps = steps[-count:]
    mean = float(np.mean(recent_values))
    std = float(np.std(recent_values))
    if count < 2 or np.ptp(recent_steps) == 0:
        slope_per_100 = 0.0
    else:
        slope_per_100 = float(np.polyfit(recent_steps, recent_values, 1)[0] * 100.0)
    return mean, std, slope_per_100


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("run", help="Run directory or its basename")
    parser.add_argument("--window", type=int, default=25, help="Recent iterations to summarize")
    args = parser.parse_args()

    run_dir = Path(args.run)
    if not run_dir.is_dir():
        run_dir = LOG_ROOT / args.run
    event_files = sorted(run_dir.glob("events.out*"), key=lambda path: path.stat().st_mtime)
    if not event_files:
        raise SystemExit(f"No TensorBoard event file in {run_dir}")

    accumulator = EventAccumulator(str(event_files[-1]), size_guidance={"scalars": 0})
    accumulator.Reload()
    scalar_tags = set(accumulator.Tags().get("scalars", []))

    print(f"run: {run_dir.name}")
    print(f"window: last {args.window} iterations")
    print("metric       latest      mean       std    slope/100")
    print("-----------  --------  --------  --------  ---------")
    history = 0
    summaries: dict[str, tuple[float, float, float, float]] = {}
    for label, tag in TAGS.items():
        if tag not in scalar_tags:
            continue
        steps, values = scalar_series(accumulator, tag)
        history = max(history, values.size)
        mean, std, slope = summarize(values, steps, args.window)
        latest = float(values[-1])
        summaries[label] = (latest, mean, std, slope)
        print(f"{label:11}  {latest:8.3f}  {mean:8.3f}  {std:8.3f}  {slope:9.3f}")

    print(f"\nhistory: {history} logged iterations")
    if history < max(100, args.window * 4):
        print("assessment: EARLY — insufficient history for a convergence claim")
        return

    command = summaries.get("cmd_lvl", (0.0, 0.0, 0.0, 0.0))[1]
    tracking = summaries.get("lin_track", (0.0, 0.0, 0.0, 0.0))[1]
    episode_length = summaries.get("ep_len", (0.0, 0.0, 0.0, 0.0))[1]
    bad_orientation = summaries.get("bad_orient", (1.0, 1.0, 0.0, 0.0))[1]
    reward_slope = summaries.get("reward", (0.0, 0.0, 0.0, 1.0))[3]

    gates = {
        "full command curriculum": command >= 1.0,
        "linear tracking >= 0.65": tracking >= 0.65,
        "episode length >= 700": episode_length >= 700.0,
        "bad orientation <= 0.30": bad_orientation <= 0.30,
        "reward plateau |slope/100| <= 1": abs(reward_slope) <= 1.0,
    }
    for name, passed in gates.items():
        print(f"{'PASS' if passed else 'WAIT'}: {name}")
    print(f"assessment: {'CONVERGED CANDIDATE' if all(gates.values()) else 'LEARNING'}")


if __name__ == "__main__":
    main()