"""Solve the slippery 8x8 FrozenLake with value iteration and policy iteration, and plot.

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

Saves dp_frozenlake.png next to this file and prints a one-line summary.
"""
import importlib
import os

import matplotlib

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

from envs import FROZEN_LAKE_8x8, make_frozen_lake, plot_values  # noqa: E402

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

HERE = os.path.dirname(os.path.abspath(__file__))


def main():
    mdp = make_frozen_lake(FROZEN_LAKE_8x8, slippery=True)
    P, R = mdp.P, mdp.R
    start = mdp.cells.index((0, 0))

    fig, axes = plt.subplots(1, 3, figsize=(14, 4.3))

    # Panel 1: value iteration's max change per sweep, for three discount factors.
    ax = axes[0]
    for gamma, color in [(0.5, "C0"), (0.9, "C1"), (0.99, "C2")]:
        _, _, deltas = dp.value_iteration(P, R, gamma, tol=1e-8)
        k = np.arange(1, len(deltas) + 1)
        ax.semilogy(k, np.maximum(deltas, 1e-12), color=color, label=f"γ = {gamma}: {len(deltas)} sweeps")
        ax.semilogy(k, deltas[0] * gamma ** (k - 1), color=color, ls="--", lw=0.8)
    ax.set_xlabel("sweep k")
    ax.set_ylabel("max change  ‖V_k − V_{k−1}‖∞")
    ax.set_title("Value iteration contracts (dashed: δ₁·γ^(k−1))", fontsize=9)
    ax.set_ylim(1e-9, 2)
    ax.legend(fontsize=8)

    # Panel 2: V* and the optimal policy for gamma = 0.99.
    gamma = 0.99
    V_star, pi_vi, deltas = dp.value_iteration(P, R, gamma, tol=1e-10)
    plot_values(axes[1], mdp, V_star, pi_vi, title=f"V* and π* on slippery 8×8 FrozenLake (γ = {gamma})")

    # Panel 3: the greedy policy is optimal long before the values have converged.
    Q_star = dp.q_from_v(P, R, V_star, gamma)
    free = ~mdp.terminal
    V = np.zeros(mdp.n_states)
    errs, frac_opt = [], []
    for _ in range(len(deltas)):
        V = dp.q_from_v(P, R, V, gamma).max(axis=1)
        greedy = np.argmax(dp.q_from_v(P, R, V, gamma), axis=1)
        is_opt = Q_star[np.arange(mdp.n_states), greedy] >= Q_star.max(axis=1) - 1e-9
        frac_opt.append(is_opt[free].mean())
        errs.append(np.max(np.abs(V - V_star)))
    first_opt = next(i + 1 for i, f in enumerate(frac_opt) if all(x == 1.0 for x in frac_opt[i:]))
    ax = axes[2]
    # (The last iterate IS our V*, so its error is 0 by construction; leave it off the log plot.)
    ax.semilogy(np.arange(1, len(errs)), errs[:-1], color="C0", label="‖V_k − V*‖∞")
    ax.set_xlabel("sweep k")
    ax.set_ylabel("value error (log)", color="C0")
    ax2 = ax.twinx()
    ax2.plot(np.arange(1, len(frac_opt) + 1), frac_opt, color="C3", label="fraction of states greedy-optimal")
    ax2.set_ylabel("fraction of states where greedy(V_k) is optimal", color="C3", fontsize=8)
    ax2.set_ylim(0, 1.05)
    ax.axvline(first_opt, color="0.5", ls=":")
    ax.set_title(f"Greedy policy optimal from sweep {first_opt}; values need {len(deltas)}", fontsize=9)

    # Policy iteration for comparison.
    pi_pi, V_pi, history = dp.policy_iteration(P, R, gamma)
    eval_sweeps = sum(dp.policy_evaluation(P, R, h, gamma)[1] for h in history)
    assert np.allclose(V_pi, V_star, atol=1e-6), "policy iteration and value iteration disagree!"

    fig.tight_layout()
    out = os.path.join(HERE, "dp_frozenlake.png")
    fig.savefig(out, dpi=130)
    print(f"FrozenLake 8x8, γ={gamma}: V*(start)={V_star[start]:.3f}; value iteration {len(deltas)} sweeps "
          f"(greedy policy optimal from sweep {first_opt}); policy iteration {len(history)} evaluations "
          f"({eval_sweeps} evaluation sweeps). Saved {os.path.basename(out)}")


if __name__ == "__main__":
    main()