#!/usr/bin/env python3
"""Render a compact release report from one RSL-RL TensorBoard event file."""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from tensorboard.backend.event_processing.event_accumulator import (
    EventAccumulator,
)


TAGS = {
    "reward": "Train/mean_reward",
    "episode_length": "Train/mean_episode_length",
    "velocity_error": "Metrics/base_velocity/error_vel_xy",
    "bad_orientation": "Episode_Termination/bad_orientation",
    "timeout": "Episode_Termination/time_out",
    "value_loss": "Loss/value",
    "surrogate_loss": "Loss/surrogate",
}


def load_series(
    accumulator: EventAccumulator,
    tag: str,
) -> tuple[np.ndarray, np.ndarray]:
    events = accumulator.Scalars(tag)
    steps = np.asarray([event.step for event in events])
    values = np.asarray([event.value for event in events])
    return steps, values


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("event_file", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--title", default="Dropbear Stage 46")
    args = parser.parse_args()

    accumulator = EventAccumulator(
        str(args.event_file),
        size_guidance={"scalars": 0},
    )
    accumulator.Reload()
    available = set(accumulator.Tags()["scalars"])
    missing = [tag for tag in TAGS.values() if tag not in available]
    if missing:
        raise RuntimeError(f"Missing TensorBoard tags: {missing}")
    data = {
        name: load_series(accumulator, tag)
        for name, tag in TAGS.items()
    }

    plt.style.use("dark_background")
    fig, axes = plt.subplots(2, 2, figsize=(13, 8), sharex=True)
    fig.patch.set_facecolor("#07111f")
    for axis in axes.flat:
        axis.set_facecolor("#0c1727")
        axis.grid(color="#334155", alpha=0.28)
        axis.spines[["top", "right"]].set_visible(False)

    axes[0, 0].plot(*data["reward"], color="#5eead4", linewidth=2)
    axes[0, 0].set_title("Mean episode reward")
    axes[0, 0].set_ylabel("reward")

    axes[0, 1].plot(
        *data["episode_length"],
        color="#60a5fa",
        linewidth=2,
        label="episode length",
    )
    axes[0, 1].axhline(
        1000,
        color="#a7f3d0",
        linestyle=":",
        linewidth=1,
        label="full horizon",
    )
    axes[0, 1].set_title("Survival recovery")
    axes[0, 1].set_ylabel("steps")
    axes[0, 1].legend(frameon=False, fontsize=8)

    axes[1, 0].plot(
        *data["velocity_error"],
        color="#fbbf24",
        linewidth=2,
    )
    axes[1, 0].set_title("Planar velocity tracking error")
    axes[1, 0].set_ylabel("m/s")
    axes[1, 0].set_xlabel("checkpoint iteration")

    axes[1, 1].plot(
        *data["bad_orientation"],
        color="#fb7185",
        linewidth=2,
        label="bad orientation",
    )
    axes[1, 1].plot(
        *data["timeout"],
        color="#4ade80",
        linewidth=2,
        label="time-out",
    )
    axes[1, 1].set_title("Episode completion")
    axes[1, 1].set_ylabel("fraction")
    axes[1, 1].set_xlabel("checkpoint iteration")
    axes[1, 1].legend(frameon=False, fontsize=8)

    start = int(data["reward"][0][0])
    end = int(data["reward"][0][-1])
    fig.suptitle(
        f"{args.title} · iterations {start}–{end}",
        x=0.065,
        ha="left",
        fontsize=17,
        fontweight="bold",
    )
    fig.text(
        0.985,
        0.015,
        "RSL-RL PPO continuation · 512 parallel environments",
        ha="right",
        color="#94a3b8",
        fontsize=8,
    )
    fig.tight_layout(rect=(0.04, 0.04, 0.99, 0.93))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.output, dpi=180, facecolor=fig.get_facecolor())


if __name__ == "__main__":
    main()