"""Watch a few episodes, then compare the bridge and the long way round as the floor gets slippery.

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

Saves returns_vs_slip.png next to this file.
"""
import importlib
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

gw = importlib.import_module("solutions.gridworld" if os.environ.get("RL_SOLUTIONS") else "gridworld")

GAMMA = 0.9
GLYPH = "↑→↓←"


def show_episode(env, policy, name, seed):
    """Print one episode as a list of (state, action, reward)."""
    rng = np.random.default_rng(seed)
    s = env.reset()
    rewards, lines = [], []
    for t in range(100):
        a = policy(s, rng)
        s2, r, done = env.step(a)
        rewards.append(r)
        lines.append(f"{s}{GLYPH[a]}{'' if r == 0 else f'({r:+g})'}")
        s = s2
        if done:
            break
    G = gw.discounted_return(rewards, GAMMA)
    path = " ".join(lines)
    print(f"  {name:>7}: {len(rewards):3d} steps, G_0 = {G:+.3f}   {path[:110]}{' ...' if len(path) > 110 else ''}")


def main():
    print(f"One episode each (slip = 0.2, gamma = {GAMMA}). Each entry is state + chosen action (+ reward if nonzero);")
    print("a repeated state means you bumped into a wall or edge (on purpose, or because the floor slipped):")
    for name, pol, seed in [("random", gw.random_policy, 0), ("bridge", gw.bridge_policy, 1), ("safe", gw.safe_policy, 2)]:
        show_episode(gw.GridWorld(slip=0.2, seed=seed), pol, name, seed)

    print("\nAverage return over 2000 episodes (mean ± standard error):")
    for slip in (0.0, 0.2):
        row = []
        for name, pol in [("random", gw.random_policy), ("bridge", gw.bridge_policy), ("safe", gw.safe_policy)]:
            m, se = gw.estimate_return(gw.GridWorld(slip=slip, seed=10), pol, GAMMA, n_episodes=2000, seed=11)
            row.append(f"{name} {m:+.3f} ± {se:.3f}")
        print(f"  slip {slip:.1f}:  " + "   ".join(row))

    slips = np.linspace(0.0, 0.4, 17)
    results = {}
    for name, pol in [("bridge", gw.bridge_policy), ("safe", gw.safe_policy)]:
        ms, ses = [], []
        for i, p in enumerate(slips):
            m, se = gw.estimate_return(gw.GridWorld(slip=p, seed=100 + i), pol, GAMMA, n_episodes=2000, seed=200 + i)
            ms.append(m)
            ses.append(se)
        results[name] = (np.array(ms), np.array(ses))

    diff = results["bridge"][0] - results["safe"][0]
    cross = None
    for i in range(len(slips) - 1):
        if diff[i] > 0 >= diff[i + 1]:
            cross = slips[i] + (slips[i + 1] - slips[i]) * diff[i] / (diff[i] - diff[i + 1])
            break

    fig, ax = plt.subplots(figsize=(7, 4.2))
    for name, color in [("bridge", "tab:red"), ("safe", "tab:blue")]:
        m, se = results[name]
        ax.errorbar(slips, m, yerr=2 * se, marker="o", ms=4, capsize=2, color=color,
                    label={"bridge": "bridge policy (short, risky)", "safe": "safe policy (long way round)"}[name])
    if cross is not None:
        ax.axvline(cross, color="gray", ls="--", lw=1)
        ax.text(cross + 0.005, ax.get_ylim()[1] * 0.9, f"crossover ≈ {cross:.3f}", color="gray")
    ax.axhline(0, color="black", lw=0.5)
    ax.set_xlabel("slip probability")
    ax.set_ylabel(f"average return G_0  (γ = {GAMMA})")
    ax.set_title("Which route is better depends on the environment")
    ax.legend()
    fig.tight_layout()
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "returns_vs_slip.png")
    fig.savefig(out, dpi=120)

    if cross is None:
        msg = "no crossover found in slip 0-0.4 (the bridge never wins clearly, or always wins)"
    else:
        msg = f"the bridge wins below slip ≈ {cross:.3f}, the long way round above it"
        if GAMMA == 0.9:
            msg += " (exact answer: 0.068)"
    print(f"\ngamma = {GAMMA}: {msg}. Plot: {out}")


if __name__ == "__main__":
    main()