"""Lai-Robbins asymptotic lower bound and the 'regret vs log T' picture.

For any *uniformly efficient* policy (regret o(T^a) for every a > 0 on every
instance of the model), Lai & Robbins (1985) show

    liminf_{T->inf}  E[Regret_T] / log T  >=  sum_{i: Delta_i > 0} Delta_i / KL(mu_i, mu*)

where KL is the divergence between the reward distributions of arm i and the
best arm.  The bound is *asymptotic*: at finite T a good policy's regret can
sit above or below C*log T (there is an additive constant), but its slope
against log T should approach C.  For Bernoulli arms KL is the Bernoulli KL;
for Gaussian arms with common variance sigma^2 it is Delta^2/(2 sigma^2), so
each term becomes 2 sigma^2 / Delta_i.
"""
from __future__ import annotations

import numpy as np

from .algorithms import kl_bernoulli


def lai_robbins_constant(means, kind: str = "bernoulli", sigma: float = 1.0) -> float:
    means = np.asarray(means, dtype=float)
    best = means.max()
    sub = means[means < best]
    gaps = best - sub
    if kind == "bernoulli":
        return float(np.sum(gaps / kl_bernoulli(sub, best)))
    if kind == "gaussian":
        return float(np.sum(2.0 * sigma**2 / gaps))
    raise ValueError(f"unknown kind {kind!r}")


def fitted_slope(regret_mean: np.ndarray, t_from: int) -> float:
    """Least-squares slope of mean regret against log t over t >= t_from."""
    t = np.arange(1, len(regret_mean) + 1)
    mask = t >= t_from
    x, y = np.log(t[mask]), regret_mean[mask]
    return float(np.polyfit(x, y, 1)[0])


def plot_regret_vs_logT(results: list[dict], means, path: str, kind: str = "bernoulli",
                        sigma: float = 1.0, title: str | None = None) -> float:
    """Plot mean regret (with 10-90% band) of several policies against log t,
    together with the line C log t.  Returns C."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    C = lai_robbins_constant(means, kind, sigma)
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
    for res in results:
        reg = res["regret"]
        T = reg.shape[1]
        t = np.arange(1, T + 1)
        m = reg.mean(axis=0)
        lo, hi = np.percentile(reg, [10, 90], axis=0)
        for ax in axes:
            line, = ax.plot(t, m, label=res["name"], lw=1.6)
            ax.fill_between(t, lo, hi, color=line.get_color(), alpha=0.12, lw=0)
    t = np.arange(1, results[0]["regret"].shape[1] + 1)
    for ax in axes:
        ax.plot(t, C * np.log(t), "k--", lw=1.4, label=f"Lai-Robbins  {C:.1f}·log t")
        ax.set_xlabel("t (log scale)" if ax is axes[0] else "t")
        ax.set_ylabel("cumulative pseudo-regret")
        ax.grid(alpha=0.3)
    axes[0].set_xscale("log")
    axes[0].set_title("Regret vs log t: slopes should approach C")
    axes[1].set_title("Same, linear time axis")
    axes[0].legend(fontsize=8, loc="upper left")
    if title:
        fig.suptitle(title)
    fig.tight_layout()
    fig.savefig(path, dpi=130)
    plt.close(fig)
    return C