"""Reproduce the chapter's two experiments with YOUR code and save td_learning.png.

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

Left panel: RMS error of TD(0) vs. Monte Carlo on the 5-state random walk,
averaged over 100 runs, for a few step sizes.
Right panel: return per episode while learning on the cliff, SARSA vs.
Q-learning (epsilon = 0.1, alpha = 0.5), averaged over 20 runs.
Takes about 10-20 seconds.
"""

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 envs import CliffWalking, prediction_curve, render_path, train  # noqa: E402

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

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


def smooth(x, k=10):
    return np.convolve(x, np.ones(k) / k, mode="valid")


def main():
    t0 = time.time()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))

    # --- Prediction: random walk ----------------------------------------
    final = {}
    for alpha, ls in [(0.05, "-"), (0.1, "--"), (0.15, ":")]:
        err = prediction_curve(td.td0_update, alpha)
        ax1.plot(err, color="C0", ls=ls, label=f"TD(0), α = {alpha}")
        final[f"TD α={alpha}"] = err[-1]
    for alpha, ls in [(0.01, "-"), (0.02, "--"), (0.04, ":")]:
        err = prediction_curve(td.mc_update, alpha)
        ax1.plot(err, color="C1", ls=ls, label=f"MC, α = {alpha}")
        final[f"MC α={alpha}"] = err[-1]
    ax1.set(xlabel="episodes", ylabel="RMS error (avg over 100 runs)",
            title="Random walk: TD(0) vs. Monte Carlo", ylim=(0, 0.25))
    ax1.legend(fontsize=8)

    # --- Control: cliff walking ------------------------------------------
    env = CliffWalking()
    n_runs, n_eps = 20, 500
    curves = {}
    last_Q = {}
    for name, fn, color in [("SARSA", td.sarsa_episode, "C2"), ("Q-learning", td.q_learning_episode, "C3")]:
        rets = np.zeros(n_eps)
        for seed in range(n_runs):
            Q, r = train(fn, env, n_eps, alpha=0.5, epsilon=0.1, seed=seed)
            rets += r / n_runs
        last_Q[name] = Q
        curves[name] = rets
        ax2.plot(np.arange(len(smooth(rets))) + 10, smooth(rets), color=color, label=name)
    ax2.axhline(-13, color="gray", lw=0.8, ls=":", label="optimal path (−13)")
    ax2.set(xlabel="episode", ylabel="return per episode (smoothed)",
            title="Cliff walking, ε = 0.1, α = 0.5", ylim=(-100, 0))
    ax2.legend(fontsize=8)

    fig.tight_layout()
    out = os.path.join(HERE, "td_learning.png")
    fig.savefig(out, dpi=120)

    for name, Q in last_Q.items():
        print(f"\n{name} greedy path (last run):\n{render_path(Q)}")
    print()
    print("RMS error after 100 episodes: " + ", ".join(f"{k}: {v:.3f}" for k, v in final.items()))
    print(f"Cliff, mean return over the last 100 episodes: SARSA {curves['SARSA'][-100:].mean():.1f}, "
          f"Q-learning {curves['Q-learning'][-100:].mean():.1f}  |  saved {os.path.basename(out)} "
          f"in {time.time() - t0:.1f}s")


if __name__ == "__main__":
    main()