#!/usr/bin/env python3

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path


def find_latest() -> Path | None:
    paths = sorted(Path("outputs").glob("*/reward_log.csv"), key=lambda p: p.stat().st_mtime, reverse=True)
    if paths:
        return paths[0]
    p = Path("docs/driftshield_proof_reward_log.csv")
    return p if p.exists() else None


def load(path: Path):
    with open(path, newline="") as f:
        r = list(csv.DictReader(f))
    if not r:
        return None
    return r


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("csv_path", nargs="?", type=Path)
    ap.add_argument("--out", type=Path, default=None)
    args = ap.parse_args()
    p = args.csv_path or find_latest()
    if not p or not p.exists():
        print("No CSV found. Provide path or add docs/driftshield_proof_reward_log.csv")
        return 1
    rows = load(p)
    if not rows:
        return 1
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        print("Install matplotlib: pip install matplotlib")
        return 1

    eps = [int(x["episode"]) for x in rows if x.get("episode", "").isdigit()]
    if not eps:
        eps = list(range(1, len(rows) + 1))
    tot = [float(x.get("total_reward", x.get("total", 0))) for x in rows]
    inv = [float(x.get("investigation", 0)) for x in rows]
    rout = [float(x.get("routing", 0)) for x in rows]
    pen = [float(x.get("penalty", x.get("penalty_sum", 0))) for x in rows]
    p95 = [1.0 if t >= 0.9 else 0.0 for t in tot]

    fig, axes = plt.subplots(2, 2, figsize=(11, 7))
    ax0, ax1, ax2, ax3 = axes[0, 0], axes[0, 1], axes[1, 0], axes[1, 1]
    n = list(range(1, len(tot) + 1)) if not eps or len(eps) != len(tot) else eps
    if len(n) != len(tot):
        n = list(range(1, len(tot) + 1))
    ax0.plot(n, tot, "-o", markersize=2, color="#1f77b4")
    ax0.set_ylabel("total_reward")
    ax0.set_title("DriftShield — total reward per episode")
    ax0.grid(alpha=0.3)
    ax1.plot(n, inv, label="investigation", alpha=0.8)
    ax1.plot(n, rout, label="routing", alpha=0.8)
    ax1.set_ylabel("subscore (0-1)")
    ax1.set_title("Base components (sampled)")
    ax1.legend(fontsize=8)
    ax1.grid(alpha=0.3)
    ax2.plot(n, pen, color="darkred", alpha=0.85)
    ax2.set_ylabel("penalty_sum")
    ax2.set_title("Penalty mass")
    ax2.grid(alpha=0.3)
    ax3.plot(n, p95, drawstyle="steps-post", color="#2ca02c")
    ax3.set_ylabel("1 if total>=0.9 else 0")
    ax3.set_xlabel("episode (or index)")
    ax3.set_title("Pass-rate proxy")
    ax3.set_ylim(-0.1, 1.1)
    ax3.grid(alpha=0.3)
    fig.suptitle(f"Source: {p.name}", fontsize=9)
    fig.tight_layout()
    out = args.out or p.with_name("reward_curve.png")
    fig.savefig(out, dpi=150)
    print(f"Wrote {out.resolve()}")
    return 0


if __name__ == "__main__":
    sys.exit(main())