"""Plot all statistics from YC-Bench result JSON files.

Usage:
    uv run python scripts/plot_run.py results/yc_bench_result_medium_1_*.json
    uv run python scripts/plot_run.py results/some_result.json  # single run
"""
from __future__ import annotations

import json
import sys
from datetime import datetime
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

COLORS = ['#00d4aa', '#ff6b6b', '#4ecdc4', '#ffe66d', '#a29bfe', '#fd79a8', '#6c5ce7', '#00b894']
DOMAIN_COLORS = {
    "research": "#3498db",
    "inference": "#9b59b6",
    "data_environment": "#1abc9c",
    "training": "#e67e22",
}


def load(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def dt(iso: str) -> datetime:
    return datetime.fromisoformat(iso)


def short_name(data: dict, path: str) -> str:
    model = data.get("model", "")
    if "/" in model:
        return model.split("/")[-1]
    return Path(path).stem.split("_", 4)[-1]


# ---------------------------------------------------------------------------
# Individual plot functions
# ---------------------------------------------------------------------------

def plot_funds(ax, runs):
    for i, (path, data) in enumerate(runs):
        funds = data["time_series"]["funds"]
        if not funds:
            continue
        times = [dt(f["time"]) for f in funds]
        vals = [f["funds_cents"] / 100 for f in funds]
        ax.plot(times, vals, color=COLORS[i % len(COLORS)], linewidth=2, label=short_name(data, path))
    ax.axhline(y=200000, color='gray', linestyle='--', alpha=0.3)
    ax.set_ylabel("Funds ($)")
    ax.set_title("Funds Over Time")


def plot_tasks_cumulative(ax, runs):
    for i, (path, data) in enumerate(runs):
        tasks = data["time_series"].get("tasks", [])
        ok = sorted([t for t in tasks if t.get("success") is True and t.get("completed_at")], key=lambda t: t["completed_at"])
        fail = sorted([t for t in tasks if t.get("success") is False and t.get("completed_at")], key=lambda t: t["completed_at"])
        color = COLORS[i % len(COLORS)]
        name = short_name(data, path)
        if ok:
            ax.step([dt(t["completed_at"]) for t in ok], range(1, len(ok)+1), color=color, linewidth=2, label=f"{name} OK", where='post')
        if fail:
            ax.step([dt(t["completed_at"]) for t in fail], range(1, len(fail)+1), color=color, linewidth=1.5, linestyle='--', label=f"{name} fail", where='post', alpha=0.6)
    ax.set_ylabel("Cumulative Tasks")
    ax.set_title("Task Completions (OK vs Fail)")


def plot_prestige(ax, runs):
    # Only plot first run's prestige to avoid clutter
    if not runs:
        return
    path, data = runs[0]
    prestige = data["time_series"].get("prestige", [])
    if not prestige:
        return
    domains = sorted(set(p["domain"] for p in prestige))
    for domain in domains:
        pts = [p for p in prestige if p["domain"] == domain]
        times = [dt(p["time"]) for p in pts]
        levels = [p["level"] for p in pts]
        ax.plot(times, levels, color=DOMAIN_COLORS.get(domain, 'gray'), linewidth=1.5, label=domain)
    ax.set_ylabel("Prestige Level")
    ax.set_title(f"Prestige by Domain ({short_name(data, path)})")


def plot_trust(ax, runs):
    if not runs:
        return
    path, data = runs[0]
    trust = data["time_series"].get("client_trust", [])
    if not trust:
        return
    clients = sorted(set(t["client_name"] for t in trust))
    for client in clients:
        pts = [t for t in trust if t["client_name"] == client]
        times = [dt(t["time"]) for t in pts]
        levels = [t["trust_level"] for t in pts]
        is_rat = pts[0].get("loyalty", 0) < -0.3
        ax.plot(times, levels, linewidth=1.5, linestyle='--' if is_rat else '-', label=f"{client}{'*' if is_rat else ''}")
    ax.set_ylabel("Trust Level")
    ax.set_title(f"Client Trust ({short_name(data, path)}) (* = RAT)")


def plot_payroll(ax, runs):
    for i, (path, data) in enumerate(runs):
        ledger = data["time_series"].get("ledger", [])
        payrolls = [e for e in ledger if e["category"] == "monthly_payroll"]
        if not payrolls:
            continue
        # Group by month
        monthly = {}
        for p in payrolls:
            m = p["time"][:7]
            monthly[m] = monthly.get(m, 0) + abs(p["amount_cents"])
        months = sorted(monthly.keys())
        times = [datetime.strptime(m, "%Y-%m") for m in months]
        amounts = [monthly[m] / 100 for m in months]
        ax.plot(times, amounts, color=COLORS[i % len(COLORS)], linewidth=2, marker='o', markersize=3, label=short_name(data, path))
    ax.set_ylabel("Monthly Payroll ($)")
    ax.set_title("Payroll Growth")


def plot_assignments(ax, runs):
    for i, (path, data) in enumerate(runs):
        assignments = data["time_series"].get("assignments", [])
        completed = [a for a in assignments if a.get("completed_at")]
        if not completed:
            continue
        times = [dt(a["completed_at"]) for a in completed]
        counts = [a["num_assigned"] for a in completed]
        ax.scatter(times, counts, color=COLORS[i % len(COLORS)], alpha=0.5, s=15, label=short_name(data, path))
    ax.axhline(y=4, color='green', linestyle='--', alpha=0.3, label='efficient (4)')
    ax.set_ylabel("Employees Assigned")
    ax.set_title("Assignment Pattern Per Task")


def plot_tokens(ax, runs):
    for i, (path, data) in enumerate(runs):
        transcript = data.get("transcript", [])
        if not transcript or not transcript[0].get("prompt_tokens"):
            continue
        turns = [t["turn"] for t in transcript]
        prompt = [t.get("prompt_tokens", 0) for t in transcript]
        color = COLORS[i % len(COLORS)]
        ax.plot(turns, prompt, color=color, linewidth=1, alpha=0.7, label=f"{short_name(data, path)} prompt")
    ax.set_ylabel("Tokens")
    ax.set_title("Prompt Tokens Per Turn")
    ax.set_xlabel("Turn")


def plot_cost(ax, runs):
    for i, (path, data) in enumerate(runs):
        transcript = data.get("transcript", [])
        if not transcript:
            continue
        costs = [t.get("cost_usd", 0) for t in transcript]
        cumulative = []
        running = 0
        for c in costs:
            running += c
            cumulative.append(running)
        turns = [t["turn"] for t in transcript]
        ax.plot(turns, cumulative, color=COLORS[i % len(COLORS)], linewidth=2, label=short_name(data, path))
    ax.set_ylabel("Cumulative Cost ($)")
    ax.set_title("API Cost")
    ax.set_xlabel("Turn")


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

def print_summary(data, path):
    ts = data["time_series"]
    ledger = ts.get("ledger", [])
    cats = {}
    for e in ledger:
        cats[e["category"]] = cats.get(e["category"], 0) + e["amount_cents"]

    revenue = cats.get("task_reward", 0)
    payroll = abs(cats.get("monthly_payroll", 0))
    tasks = ts.get("tasks", [])
    ok = sum(1 for t in tasks if t.get("success") is True)
    fail = sum(1 for t in tasks if t.get("success") is False)
    gated = sum(1 for t in tasks if t.get("success") is True and t.get("required_trust", 0) > 0)

    assignments = ts.get("assignments", [])
    avg_emp = sum(a["num_assigned"] for a in assignments) / len(assignments) if assignments else 0

    employees = ts.get("employees", [])
    final_payroll = sum(e["salary_cents"] for e in employees) / 100 if employees else 0

    clients = ts.get("clients", [])
    rats = [c for c in clients if c.get("is_rat")]

    transcript = data.get("transcript", [])
    total_prompt = sum(t.get("prompt_tokens", 0) for t in transcript)
    total_completion = sum(t.get("completion_tokens", 0) for t in transcript)
    final_funds = (200000 * 100 + revenue - payroll) / 100

    print(f"\n{'='*60}")
    print(f"  {short_name(data, path)}")
    print(f"{'='*60}")
    print(f"  Model:     {data.get('model', '?')}")
    print(f"  Seed:      {data.get('seed', '?')}")
    print(f"  Terminal:  {data.get('terminal_reason', '?')} at turn {data.get('turns_completed', '?')}")
    print(f"  Final:     ${final_funds:,.0f}")
    print(f"  Revenue:   ${revenue/100:,.0f} | Payroll: ${payroll/100:,.0f}")
    print(f"  Tasks:     {ok} OK, {fail} fail ({gated} trust-gated)")
    print(f"  Avg emp:   {avg_emp:.1f} per task")
    print(f"  Payroll:   ${final_payroll:,.0f}/mo (final)")
    print(f"  RATs:      {len(rats)} — {', '.join(c['name'] for c in rats) if rats else 'none'}")
    print(f"  Scratchpad: {'yes' if ts.get('scratchpad') else 'no'}")
    total_tokens = total_prompt + total_completion
    print(f"  Tokens:    {total_prompt:,} prompt + {total_completion:,} completion = {total_tokens:,} total")
    print(f"  Cost:      ${data.get('total_cost_usd', 0):.2f}")
    started = data.get('started_at', '')
    ended = data.get('ended_at', '')
    if started and ended:
        try:
            t0 = datetime.fromisoformat(started)
            t1 = datetime.fromisoformat(ended)
            duration = t1 - t0
            mins = duration.total_seconds() / 60
            print(f"  Time:      {started[:19]} → {ended[:19]} ({mins:.1f} min)")
        except Exception:
            print(f"  Time:      {started[:19]} → {ended[:19]}")
    else:
        print(f"  Time:      N/A")

    config = ts.get("config", {})
    if config:
        print(f"  Config:    salary_bump={config.get('salary_bump_pct')}, "
              f"trust_build={config.get('trust_build_rate')}, "
              f"rat_fraction={config.get('loyalty_rat_fraction')}, "
              f"fail_penalty={config.get('penalty_fail_funds_pct')}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    if len(sys.argv) < 2:
        print("Usage: uv run python scripts/plot_run.py results/*.json")
        sys.exit(1)

    paths = sys.argv[1:]
    runs = [(p, load(p)) for p in paths]

    for path, data in runs:
        print_summary(data, path)

    fig, axes = plt.subplots(4, 2, figsize=(16, 20))
    fig.suptitle(f"YC-Bench — {len(runs)} run(s)", fontsize=14, fontweight='bold')

    plot_funds(axes[0, 0], runs)
    plot_tasks_cumulative(axes[0, 1], runs)
    plot_prestige(axes[1, 0], runs)
    plot_trust(axes[1, 1], runs)
    plot_payroll(axes[2, 0], runs)
    plot_assignments(axes[2, 1], runs)
    plot_tokens(axes[3, 0], runs)
    plot_cost(axes[3, 1], runs)

    for ax in axes.flat:
        ax.legend(fontsize=7, loc='best')
        ax.grid(True, alpha=0.2)
        ax.tick_params(labelsize=8)
        if ax.get_xlabel() != "Turn":
            try:
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%b'))
                ax.xaxis.set_major_locator(mdates.MonthLocator())
            except Exception:
                pass

    plt.tight_layout()

    Path("plots").mkdir(exist_ok=True)
    out = "plots/run_analysis.png"
    plt.savefig(out, dpi=150, bbox_inches='tight')
    print(f"\nPlot saved to {out}")


if __name__ == "__main__":
    main()