"""Race the Chapter 2 agents on the 10-armed Bernoulli testbed and plot the results.

    python run.py            # stationary testbed  -> bandits.png      (about 20 s)
    python run.py --switch   # arms reshuffled at step 1000 -> bandits_switch.png  (2000 steps)

Uses your bandits.py, or the reference solution with RL_SOLUTIONS=1.
"""

import argparse
import importlib
import os
import time

import matplotlib

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

from envs import bernoulli_testbed, switching_testbed, run_testbed

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


def agents_stationary():
    B = bandits
    return {
        "greedy": lambda k, rng: B.EpsilonGreedy(k, 0.0, rng),
        "ε-greedy (ε = 0.1)": lambda k, rng: B.EpsilonGreedy(k, 0.1, rng),
        "ε-greedy (ε = 0.01)": lambda k, rng: B.EpsilonGreedy(k, 0.01, rng),
        "optimistic greedy (Q₁ = 1, α = 0.1)": lambda k, rng: B.EpsilonGreedy(k, 0.0, rng, initial_value=1.0, step_size=0.1),
        "UCB1 (c = √2)": lambda k, rng: B.UCB1(k, rng),
        "Thompson (Beta)": lambda k, rng: B.ThompsonBernoulli(k, rng),
    }


def agents_switch():
    B = bandits
    return {
        "ε-greedy (ε = 0.1), sample average": lambda k, rng: B.EpsilonGreedy(k, 0.1, rng),
        "ε-greedy (ε = 0.1), constant α = 0.1": lambda k, rng: B.EpsilonGreedy(k, 0.1, rng, step_size=0.1),
        "UCB1 (c = √2)": lambda k, rng: B.UCB1(k, rng),
        "Thompson (Beta)": lambda k, rng: B.ThompsonBernoulli(k, rng),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--switch", action="store_true", help="nonstationary: shuffle the arms halfway through")
    parser.add_argument("--runs", type=int, default=200)
    parser.add_argument("--steps", type=int, default=None, help="default 1000 (2000 with --switch)")
    args = parser.parse_args()

    if args.switch:
        args.steps = args.steps or 2000
        agents, make_bandit, out_name = agents_switch(), switching_testbed(10, args.steps // 2), "bandits_switch.png"
        title = f"10-armed Bernoulli bandit, arms shuffled at step {args.steps // 2}"
    else:
        args.steps = args.steps or 1000
        agents, make_bandit, out_name = agents_stationary(), bernoulli_testbed(10), "bandits.png"
        title = "10-armed Bernoulli testbed"

    t0 = time.time()
    results = {}
    for name, make_agent in agents.items():
        results[name] = run_testbed(make_agent, make_bandit, runs=args.runs, steps=args.steps, seed=0)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
    smooth = lambda x, w=20: np.convolve(x, np.ones(w) / w, mode="valid")
    for name, out in results.items():
        axes[0].plot(smooth(out["avg_reward"]), label=name)
        axes[1].plot(100 * smooth(out["pct_optimal"]), label=name)
        axes[2].plot(out["cum_regret"], label=name)
    axes[0].set(title="Average reward", xlabel="step")
    axes[1].set(title="% optimal action", xlabel="step", ylim=(0, 100))
    axes[2].set(title="Cumulative regret", xlabel="step")
    axes[2].legend(fontsize=8)
    fig.suptitle(f"{title}, {args.runs} runs")
    fig.tight_layout()
    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), out_name)
    fig.savefig(path, dpi=120)

    summary = ", ".join(f"{n}: {r['cum_regret'][-1]:.0f}" for n, r in results.items())
    print(f"regret after {args.steps} steps -> {summary}. Saved {out_name} ({time.time() - t0:.0f}s)")


if __name__ == "__main__":
    main()