"""Train the tool-using agent three ways and plot what happens.

    python run.py                  # your agentic.py
    RL_SOLUTIONS=1 python run.py   # the reference solution

Variants (same seed, same budget):
    masked, trajectory-level   the correct recipe: only model tokens get gradient
    masked, turn-level         reward-to-go per turn instead of one number per trajectory
    UNmasked                   tool-output tokens are trained too (the bug)

Saves agentic_rl.png next to this file and prints a few episodes of the trained agent.
Takes about 1.5-3 minutes on a laptop CPU.
"""

import importlib
import os
import time

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from tinyagent import render  # noqa: E402

ag = importlib.import_module("solutions.agentic" if os.environ.get("RL_SOLUTIONS") else "agentic")

ITERS = 250
VARIANTS = [
    ("masked, trajectory-level", dict(mode="trajectory", mask_env=True), "C0"),
    ("masked, turn-level (reward-to-go)", dict(mode="turn", mask_env=True), "C2"),
    ("UNmasked (tool tokens trained)", dict(mode="trajectory", mask_env=False), "C3"),
]


def smooth(x, k=15):
    x = np.asarray(x, dtype=float)
    out = np.convolve(x, np.ones(k) / k, mode="valid")
    return np.arange(k - 1, len(x)), out


def main():
    t0 = time.time()
    results = {}
    for name, kw, _ in VARIANTS:
        policy, hist = ag.train(iters=ITERS, seed=0, **kw)
        results[name] = (policy, hist)
        last = hist[-20:]
        print(f"{name:36s}  solved: easy {np.mean([h['solved_easy'] for h in last]):.2f}, "
              f"hard {np.mean([h['solved_hard'] for h in last]):.2f}   tool calls: easy "
              f"{np.mean([h['calls_easy'] for h in last]):.2f}, hard {np.mean([h['calls_hard'] for h in last]):.2f}"
              f"   ({time.time() - t0:.0f}s)")

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
    ax = axes[0]
    for name, _, color in VARIANTS:
        hist = results[name][1]
        ax.plot(*smooth([h["solved_hard"] for h in hist]), color=color, label=name)
    ax.axhline(1 / 8, color="gray", ls=":", lw=1)
    ax.text(2, 1 / 8 + 0.02, "guessing without the tool (1/8)", color="gray", fontsize=8)
    ax.set(title="Hidden number: success rate", xlabel="iteration", ylabel="fraction solved", ylim=(0, 1.02))
    ax.legend(fontsize=8, loc="center right")

    ax = axes[1]
    hist = results[VARIANTS[0][0]][1]
    ax.plot(*smooth([h["calls_hard"] for h in hist]), color="C1", label="hidden number (tool helps)")
    ax.plot(*smooth([h["calls_easy"] for h in hist]), color="C0", ls="--", label="number in the prompt (tool useless)")
    ax.set(title="Tool calls per episode (masked, trajectory-level)", xlabel="iteration", ylabel="calls")
    ax.legend(fontsize=8)

    ax = axes[2]
    for name, _, color in VARIANTS:
        hist = results[name][1]
        ax.plot(*smooth([h["calls_hard"] for h in hist]), color=color, label=name)
    ax.set(title="Tool calls on hidden-number tasks", xlabel="iteration", ylabel="calls")
    ax.legend(fontsize=8)

    fig.tight_layout()
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "agentic_rl.png")
    fig.savefig(path, dpi=120)

    print("\nA few episodes of the trained agent (masked, trajectory-level):")
    policy = results[VARIANTS[0][0]][0]
    sampler, rng = policy.sampler(), np.random.default_rng(123)
    for secret, revealed in [(6, True), (2, False), (5, False), (0, False)]:
        print("  " + render(ag.rollout(sampler, ag.GuessEnv(secret, revealed), rng)))
    print(f"\nsaved {path}  ({time.time() - t0:.0f}s total)")


if __name__ == "__main__":
    main()