#!/usr/bin/env python3
"""Reproduce all bandits-lab figures into out/.

  python experiments/run_experiments.py            # default (~1-2 min)
  python experiments/run_experiments.py --quick    # smoke run (~15 s)

Figures:
  out/single_player_regret.png   six policies on a 5-armed Bernoulli instance
  out/lai_robbins.png            KL-UCB / TS / UCB1 vs the Lai-Robbins line
  out/gaussian_ts.png            Gaussian TS vs UCB1 on Gaussian arms
  out/multiplayer.png            selfish UCB vs musical chairs vs SIC-style
"""
from __future__ import annotations

import argparse
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import matplotlib  # noqa: E402

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

from bandits.algorithms import (KLUCB, UCB1, BernoulliEnv, EpsilonGreedy,  # noqa: E402
                                ExploreThenCommit, GaussianEnv, ThompsonBeta,
                                ThompsonGaussian, simulate)
from bandits.lower_bound import fitted_slope, lai_robbins_constant, plot_regret_vs_logT  # noqa: E402
from bandits.multiplayer import musical_chairs, selfish_ucb, sic_simplified  # noqa: E402

OUT = os.path.join(os.path.dirname(__file__), "..", "out")


def single_player(T, runs):
    env = BernoulliEnv([0.9, 0.8, 0.7, 0.6, 0.5])
    policies = [EpsilonGreedy(0.1), EpsilonGreedy(decaying=True, c=5.0),
                ExploreThenCommit(m=50), UCB1(), KLUCB(), ThompsonBeta()]
    results = []
    print(f"[single-player] Bernoulli means {env.means.tolist()}, T={T}, runs={runs}")
    for i, pol in enumerate(policies):
        t0 = time.time()
        res = simulate(env, pol, T, runs, seed=100 + i)
        results.append(res)
        r = res["regret"][:, -1]
        print(f"  {res['name']:<18} regret@T = {r.mean():7.1f} ± {r.std() / np.sqrt(runs):5.1f}"
              f"   ({time.time() - t0:.1f}s)")
    fig, ax = plt.subplots(figsize=(7, 4.5))
    t = np.arange(1, T + 1)
    for res in results:
        ax.plot(t, res["regret"].mean(axis=0), label=res["name"])
    ax.set_xlabel("t")
    ax.set_ylabel("mean cumulative pseudo-regret")
    ax.set_title(f"Bernoulli {env.means.tolist()}, {runs} runs")
    ax.grid(alpha=0.3)
    ax.legend()
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "single_player_regret.png"), dpi=130)
    plt.close(fig)
    return env, results


def lower_bound(env, results, T):
    keep = [r for r in results if r["name"] in ("UCB1", "KL-UCB", "TS-Beta")]
    C = plot_regret_vs_logT(keep, env.means, os.path.join(OUT, "lai_robbins.png"),
                            title=f"Bernoulli {env.means.tolist()}")
    print(f"[lower bound] Lai-Robbins constant C = {C:.2f}  "
          f"(C log T at T={T}: {C * np.log(T):.1f})")
    for r in keep:
        s = fitted_slope(r["regret"].mean(axis=0), t_from=T // 10)
        print(f"  {r['name']:<8} slope of regret vs log t over [T/10, T] = {s:6.2f}"
              f"   (ratio to C: {s / C:.2f})")


def gaussian(T, runs):
    env = GaussianEnv([1.0, 0.5, 0.0], sigma=1.0)
    res = [simulate(env, UCB1(), T, runs, seed=7),
           simulate(env, ThompsonGaussian(sigma=1.0), T, runs, seed=8)]
    C = plot_regret_vs_logT(res, env.means, os.path.join(OUT, "gaussian_ts.png"),
                            kind="gaussian", sigma=1.0, title="Gaussian arms, sigma=1")
    print(f"[gaussian] means {env.means.tolist()}, C = {C:.1f}")
    for r in res:
        print(f"  {r['name']:<12} regret@T = {r['regret'][:, -1].mean():.1f}")


def multiplayer(T, runs, T0):
    means = np.array([0.9, 0.8, 0.7, 0.5, 0.4, 0.3])
    M = 3
    print(f"[multiplayer] K={len(means)}, M={M}, means {means.tolist()}, T={T}, runs={runs}")
    out = []
    for fn, kw in ((selfish_ucb, {}), (musical_chairs, {"T0": T0}),
                   (sic_simplified, {})):
        t0 = time.time()
        res = fn(means, M, T, n_runs=runs, seed=11, **kw)
        out.append(res)
        r = res["regret"][:, -1]
        final_step = res["regret"][:, -1] - res["regret"][:, -2]
        extra = f", final round on best {M} arms: {np.mean(final_step < 1e-9):.0%}"
        extra += f", collisions/run {res['collisions'].mean():.0f}" if "collisions" in res else ""
        extra += f", protocol failures {res['failures']}" if "failures" in res else ""
        print(f"  {res['name']:<26} collective regret@T = {r.mean():8.1f} "
              f"(median {np.median(r):.1f}){extra}   ({time.time() - t0:.1f}s)")
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
    t = np.arange(1, T + 1)
    for res in out:
        m = res["regret"].mean(axis=0)
        lo, hi = np.percentile(res["regret"], [10, 90], axis=0)
        for ax in axes:
            line, = ax.plot(t, m, label=res["name"])
            ax.fill_between(t, lo, hi, color=line.get_color(), alpha=0.12, lw=0)
    axes[1].set_xscale("log")
    axes[1].set_yscale("symlog", linthresh=10)
    for ax in axes:
        ax.set_xlabel("t")
        ax.set_ylabel("collective pseudo-regret")
        ax.grid(alpha=0.3)
    axes[0].legend()
    axes[0].set_title(f"M={M} players, K={len(means)} arms, no communication")
    axes[1].set_title("log-log view")
    axes[1].set_ylim(bottom=0)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "multiplayer.png"), dpi=130)
    plt.close(fig)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true", help="small T and few runs")
    args = ap.parse_args()
    os.makedirs(OUT, exist_ok=True)
    t0 = time.time()
    if args.quick:
        env, res = single_player(T=2000, runs=40)
        lower_bound(env, res, 2000)
        gaussian(T=2000, runs=40)
        multiplayer(T=5000, runs=10, T0=500)
    else:
        env, res = single_player(T=20000, runs=100)
        lower_bound(env, res, 20000)
        gaussian(T=20000, runs=100)
        multiplayer(T=50000, runs=30, T0=1000)
    print(f"done in {time.time() - t0:.0f}s; figures in {os.path.abspath(OUT)}")


if __name__ == "__main__":
    main()