#!/usr/bin/env python3
"""Print a compact status summary from prime-rl's local metrics sink."""

from __future__ import annotations

import argparse
import json
import statistics
from collections import defaultdict
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("metrics", type=Path)
    parser.add_argument("--max-steps", type=int, required=True)
    args = parser.parse_args()

    by_step: dict[int, dict] = defaultdict(dict)
    for line in args.metrics.read_text().splitlines():
        row = json.loads(line)
        by_step[int(row["step"])].update(row)
    if not by_step:
        print("no completed steps")
        return

    steps = sorted(by_step)
    latest = by_step[steps[-1]]
    recent = [by_step[step] for step in steps[-10:]]
    # A compiled first step reports zero throughput and can be several minutes
    # slower; exclude it from the steady-state ETA.
    times = [
        row["time/step"]
        for row in recent
        if row.get("time/step", 0) > 0 and row.get("perf/throughput", 0) > 0
    ]
    throughputs = [row["perf/throughput"] for row in recent if row.get("perf/throughput", 0) > 0]
    losses = [row["loss/mean"] for row in recent if "loss/mean" in row]
    mean_step_seconds = statistics.mean(times)
    remaining_seconds = max(0, args.max_steps - steps[-1]) * mean_step_seconds
    result = {
        "step": steps[-1],
        "max_steps": args.max_steps,
        "tokens": latest.get("progress/num_tokens"),
        "samples": latest.get("progress/num_samples"),
        "lr": latest.get("optim/lr"),
        "loss": latest.get("loss/mean"),
        "recent_mean_loss": statistics.mean(losses),
        "recent_mean_throughput": statistics.mean(throughputs) if throughputs else 0,
        "recent_mean_step_seconds": mean_step_seconds,
        "estimated_remaining_hours": remaining_seconds / 3600,
        "peak_memory_gib": latest.get("perf/peak_memory"),
        "grad_norm": latest.get("optim/grad_norm"),
        "nan_count": latest.get("loss/nan_count"),
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()