"""Evaluate a few gridworld policies and plot their value functions.

    python run.py                  # uses your mdp.py
    RL_SOLUTIONS=1 python run.py   # uses solutions/mdp.py

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

import matplotlib

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

import envs

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

GAMMA, SLIP, STEP = 0.9, 0.2, -0.04


def draw_values(ax, V, pi, title):
    grid = np.full((envs.N_ROWS, envs.N_COLS), np.nan)
    for s in range(envs.N_STATES):
        if not envs.is_terminal(s):
            grid[envs.cell_of(s)] = V[s]
    ax.imshow(grid, cmap="RdYlGn", vmin=-1, vmax=1)
    for s in range(envs.N_STATES):
        r, c = envs.cell_of(s)
        k = envs.kind(s)
        if k == "#":
            ax.add_patch(plt.Rectangle((c - 0.5, r - 0.5), 1, 1, color="0.35"))
        elif k in ("G", "P"):
            ax.text(c, r, "+1" if k == "G" else "−1", ha="center", va="center", fontsize=12, weight="bold")
        else:
            arrow = "·" if np.allclose(pi[s], pi[s][0]) else envs.ARROWS[int(np.argmax(pi[s]))]
            ax.text(c, r - 0.18, arrow, ha="center", va="center", fontsize=13)
            ax.text(c, r + 0.2, f"{V[s]:+.2f}", ha="center", va="center", fontsize=9)
    ax.set_title(title, fontsize=10)
    ax.set_xticks([])
    ax.set_yticks([])


def main():
    P, R = mdp.build_gridworld(slip=SLIP, step_reward=STEP)
    policies = [
        ("uniform random", envs.uniform_policy(envs.N_STATES, envs.N_ACTIONS)),
        ("shortest path", envs.SHORTEST_PATH),
        ("safe detour", envs.SAFE_DETOUR),
    ]
    fig, axes = plt.subplots(1, 4, figsize=(16, 3.4))
    summary = []
    for ax, (name, pi) in zip(axes, policies):
        V = mdp.evaluate_exact(P, R, pi, GAMMA)
        draw_values(ax, V, pi, f"{name}: V(start) = {V[envs.START]:+.3f}")
        print(f"\n{name}  (gamma={GAMMA}, slip={SLIP}, step reward={STEP})")
        print(envs.format_grid(V, pi))
        summary.append(f"{name} {V[envs.START]:+.3f}")

    # Convergence of iterative evaluation: error after k sweeps vs gamma^k.
    ax = axes[3]
    pi = envs.uniform_policy(envs.N_STATES, envs.N_ACTIONS)
    V_star = mdp.evaluate_exact(P, R, pi, GAMMA)
    errs = []
    K = 120
    for k in range(1, K + 1):
        V_k, _ = mdp.evaluate_iterative(P, R, pi, GAMMA, tol=0.0, max_iters=k)
        errs.append(max(np.max(np.abs(V_k - V_star)), 1e-16))
    ks = np.arange(1, K + 1)
    ax.semilogy(ks, errs, label="max |V_k − V^π| (random policy)")
    ax.semilogy(ks, np.max(np.abs(V_star)) * GAMMA ** ks, "--", label="bound γ^k · max|V^π|")
    ax.set_xlabel("sweeps k")
    ax.set_title("iterative evaluation: error shrinks geometrically", fontsize=10)
    ax.legend(fontsize=8)

    fig.tight_layout()
    out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "values.png")
    fig.savefig(out, dpi=120)
    _, n = mdp.evaluate_iterative(P, R, pi, GAMMA, tol=1e-10)
    print(f"\nV(start): {', '.join(summary)} | iterative eval (random policy) needed {n} sweeps | saved {out}")


if __name__ == "__main__":
    main()