"""Figure: event timelines of the R1 interrupted runs (tool uses over elapsed time, kill marker, resumed session, outcome)."""
import os, sys, json
sys.path.insert(0, os.path.dirname(__file__)); from common import *
sys.path.insert(0, os.path.join(REPO, "src")); from pgar.state import StateStore
import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
def main():
    store = StateStore(os.path.join(REPO, "runs", "state.sqlite"))
    runs = store.query("SELECT r.run_id, r.task_id, r.treatment, r.manifest_json, e.scientific_success FROM runs r LEFT JOIN evaluations e ON e.run_id=r.run_id WHERE r.split='recovery' AND r.manifest_json LIKE '%interrupt%' ORDER BY r.task_id, r.treatment, r.run_id")
    runs = [r for r in runs if (json.loads(r["manifest_json"]).get("fault") or {}).get("fault_type") == "interrupt"]
    if not runs: print("no R1 runs"); return
    fig, ax = plt.subplots(figsize=(7.4, 0.4 * len(runs) + 1.4))
    for i, r in enumerate(runs):
        ev = store.query("SELECT seq, elapsed_s, event_type, tool_name, session_id FROM events WHERE run_id=? AND event_type='assistant' AND tool_name IS NOT NULL ORDER BY seq", (r["run_id"],))
        fault = store.query("SELECT detail_json FROM faults WHERE run_id=? AND fault_type='interrupt'", (r["run_id"],))
        kill_seq = None
        if fault:
            at = json.loads(fault[0]["detail_json"]).get("at_event", ""); kill_seq = int(at.split(":")[1]) if ":" in at else None
        sess = store.query("SELECT seq, kind FROM sessions WHERE run_id=? ORDER BY seq", (r["run_id"],))
        resume_seq = next((s["seq"] for s in sess if s["kind"] == "resume_after_fault"), None)
        # x axis = cumulative tool-use index (time axes restart per session; use event order)
        xs = list(range(1, len(ev) + 1)); colors = ["#4c72b0" if (resume_seq is None or e["seq"] <= resume_seq) else "#55a868" for e in ev]
        ax.scatter(xs, [i] * len(xs), c=colors, s=10, marker="|", linewidths=1.2)
        if kill_seq is not None:
            k = sum(1 for e in ev if e["seq"] <= kill_seq); ax.plot([k + 0.5, k + 0.5], [i - 0.35, i + 0.35], color="#b00020", lw=1.6)
        ax.text(len(xs) + 1.5, i, f"{'success' if r['scientific_success'] else 'fail'}", fontsize=7, va="center", color="#0a7a2f" if r["scientific_success"] else "#b00020")
    ax.set_yticks(range(len(runs))); ax.set_yticklabels([f"{r['run_id']} {r['task_id']} arm {r['treatment']}" for r in runs], fontsize=7); ax.invert_yaxis()
    ax.set_xlabel("Tool-use index within the run (blue: before the kill; red bar: controller kill after the 6th tool use; green: resumed session)", fontsize=7)
    ax.set_title("R1 interruption and resumption: tool-use timelines and outcomes", loc="left", fontsize=9)
    fig.tight_layout(); fig.savefig(os.path.join(REPO, "results", "figures", "fig_recovery.pdf")); plt.close(fig); print("fig_recovery ok", len(runs))
if __name__ == "__main__": main()