"""Stochastic multi-armed bandit environments and index policies.

Every policy works on a *batch* of independent runs at once: its state is a
set of arrays of shape (n_runs, K), and

    arms = policy.select()            # (n_runs,) integer arm per run
    policy.update(arms, rewards)      # (n_runs,) observed rewards

advances all runs by one round.  This keeps the code readable (one round = a
handful of numpy operations) while making 100+ Monte Carlo runs cheap.
"""
from __future__ import annotations

import numpy as np

# --------------------------------------------------------------------------
# Environments
# --------------------------------------------------------------------------


class BernoulliEnv:
    """K arms; arm k pays 1 with probability means[k], else 0."""

    kind = "bernoulli"

    def __init__(self, means):
        self.means = np.asarray(means, dtype=float)
        if np.any((self.means < 0) | (self.means > 1)):
            raise ValueError("Bernoulli means must lie in [0, 1]")
        self.K = len(self.means)

    @property
    def best_mean(self) -> float:
        return float(self.means.max())

    @property
    def gaps(self) -> np.ndarray:
        return self.best_mean - self.means

    def pull(self, arms: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        return (rng.random(arms.shape) < self.means[arms]).astype(float)


class GaussianEnv:
    """K arms; arm k pays N(means[k], sigma^2)."""

    kind = "gaussian"

    def __init__(self, means, sigma: float = 1.0):
        self.means = np.asarray(means, dtype=float)
        self.sigma = float(sigma)
        self.K = len(self.means)

    @property
    def best_mean(self) -> float:
        return float(self.means.max())

    @property
    def gaps(self) -> np.ndarray:
        return self.best_mean - self.means

    def pull(self, arms: np.ndarray, rng: np.random.Generator) -> np.ndarray:
        return self.means[arms] + self.sigma * rng.standard_normal(arms.shape)


# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------


def argmax_random_ties(values: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    """Row-wise argmax that breaks ties uniformly at random."""
    is_max = values >= values.max(axis=1, keepdims=True)
    noise = np.where(is_max, rng.random(values.shape), -1.0)
    return noise.argmax(axis=1)


def kl_bernoulli(p, q, eps: float = 1e-15):
    """KL(Ber(p) || Ber(q)), elementwise, with clipping to avoid log(0)."""
    p = np.clip(p, eps, 1 - eps)
    q = np.clip(q, eps, 1 - eps)
    return p * np.log(p / q) + (1 - p) * np.log((1 - p) / (1 - q))


def klucb_index(mu, n, level, iters: int = 50):
    """Vectorized KL-UCB upper index by bisection.

    Returns max{ q in [mu, 1] : n * kl(mu, q) <= level } elementwise.
    kl(mu, .) is increasing on [mu, 1], so bisection is exact up to 2^-iters.
    """
    mu = np.asarray(mu, dtype=float)
    target = np.asarray(level, dtype=float) / np.maximum(n, 1e-12)
    lo = mu.copy()
    hi = np.ones_like(mu)
    for _ in range(iters):
        mid = 0.5 * (lo + hi)
        ok = kl_bernoulli(mu, mid) <= target
        lo = np.where(ok, mid, lo)
        hi = np.where(ok, hi, mid)
    return lo  # lo is always feasible


def klucb_index_newton(mu: float, n: float, level: float,
                       tol: float = 1e-12, max_iter: int = 100) -> float:
    """Scalar KL-UCB index by Newton's method (reference implementation).

    f(q) = kl(mu, q) - level/n is convex and increasing on [mu, 1); starting to
    the right of the root, Newton iterates decrease monotonically to it.
    """
    target = level / n
    q = 1.0 - 1e-12
    if kl_bernoulli(mu, q) <= target:  # constraint never binds before 1
        return 1.0
    for _ in range(max_iter):
        f = float(kl_bernoulli(mu, q)) - target
        fprime = (q - mu) / (q * (1 - q))
        step = f / fprime
        q_new = max(q - step, mu)
        if abs(q_new - q) < tol:
            return q_new
        q = q_new
    return q


# --------------------------------------------------------------------------
# Policies
# --------------------------------------------------------------------------


class Policy:
    """Base class: keeps per-run, per-arm pull counts and reward sums."""

    name = "policy"

    def reset(self, n_runs: int, K: int, rng: np.random.Generator) -> None:
        self.R, self.K, self.rng = n_runs, K, rng
        self.counts = np.zeros((n_runs, K))
        self.sums = np.zeros((n_runs, K))
        self.t = 0  # rounds played so far
        self._rows = np.arange(n_runs)

    @property
    def means(self) -> np.ndarray:
        return self.sums / np.maximum(self.counts, 1)

    def select(self) -> np.ndarray:  # pragma: no cover - abstract
        raise NotImplementedError

    def update(self, arms: np.ndarray, rewards: np.ndarray) -> None:
        self.counts[self._rows, arms] += 1
        self.sums[self._rows, arms] += rewards
        self.t += 1

    def _init_round_robin(self):
        """Pull every arm once first (common to index policies)."""
        if self.t < self.K:
            return np.full(self.R, self.t)
        return None


class EpsilonGreedy(Policy):
    """With prob. eps explore uniformly, otherwise play the empirical best.

    If ``decaying`` is set, eps_t = min(1, c*K/t) (Auer et al., 2002 use
    eps_t = min(1, cK/(d^2 t)) with d a lower bound on the gap).
    """

    def __init__(self, epsilon: float = 0.1, decaying: bool = False, c: float = 5.0):
        self.epsilon, self.decaying, self.c = epsilon, decaying, c
        self.name = f"eps-greedy(c={c})" if decaying else f"eps-greedy({epsilon})"

    def select(self):
        rr = self._init_round_robin()
        if rr is not None:
            return rr
        eps = min(1.0, self.c * self.K / self.t) if self.decaying else self.epsilon
        greedy = argmax_random_ties(self.means, self.rng)
        explore = self.rng.random(self.R) < eps
        return np.where(explore, self.rng.integers(self.K, size=self.R), greedy)


class ExploreThenCommit(Policy):
    """Round-robin m pulls per arm, then commit to the empirical best forever."""

    def __init__(self, m: int = 50):
        self.m = m
        self.name = f"ETC(m={m})"

    def reset(self, n_runs, K, rng):
        super().reset(n_runs, K, rng)
        self.committed = None

    def select(self):
        if self.t < self.m * self.K:
            return np.full(self.R, self.t % self.K)
        if self.committed is None:
            self.committed = argmax_random_ties(self.means, self.rng)
        return self.committed


class UCB1(Policy):
    """UCB1 (Auer, Cesa-Bianchi & Fischer, 2002): mean + sqrt(2 log t / n)."""

    name = "UCB1"

    def select(self):
        rr = self._init_round_robin()
        if rr is not None:
            return rr
        bonus = np.sqrt(2.0 * np.log(self.t + 1) / self.counts)
        return argmax_random_ties(self.means + bonus, self.rng)


class KLUCB(Policy):
    """KL-UCB for Bernoulli rewards (Garivier & Cappé, 2011).

    Index = max{q : N_k * kl(mu_k, q) <= log t + c log log t}.  c = 0 is the
    usual practical choice; the paper's analysis uses c = 3.
    """

    def __init__(self, c: float = 0.0):
        self.c = c
        self.name = "KL-UCB"

    def select(self):
        rr = self._init_round_robin()
        if rr is not None:
            return rr
        t = self.t + 1
        level = np.log(t) + (self.c * np.log(np.log(t)) if t > np.e else 0.0)
        idx = klucb_index(self.means, self.counts, level, iters=30)
        return argmax_random_ties(idx, self.rng)


class ThompsonBeta(Policy):
    """Thompson sampling with Beta(a0, b0) priors on Bernoulli arms.

    Rewards in [0,1] that are not binary are first binarized with a coin flip
    (Agrawal & Goyal's trick), which keeps the posterior conjugate.
    """

    name = "TS-Beta"

    def __init__(self, a0: float = 1.0, b0: float = 1.0):
        self.a0, self.b0 = a0, b0

    def reset(self, n_runs, K, rng):
        super().reset(n_runs, K, rng)
        self.alpha = np.full((n_runs, K), self.a0)
        self.beta = np.full((n_runs, K), self.b0)

    def select(self):
        return argmax_random_ties(self.rng.beta(self.alpha, self.beta), self.rng)

    def update(self, arms, rewards):
        binary = np.isin(rewards, (0.0, 1.0))
        r = np.where(binary, rewards, (self.rng.random(rewards.shape) < rewards))
        self.alpha[self._rows, arms] += r
        self.beta[self._rows, arms] += 1 - r
        super().update(arms, rewards)


class ThompsonGaussian(Policy):
    """Thompson sampling for Gaussian rewards with known noise sigma.

    Conjugate prior N(mu0, tau0^2) per arm; posterior after n pulls with sum s:
    precision = 1/tau0^2 + n/sigma^2, mean = (mu0/tau0^2 + s/sigma^2)/precision.
    """

    name = "TS-Gaussian"

    def __init__(self, sigma: float = 1.0, mu0: float = 0.0, tau0: float = 1e3):
        self.sigma, self.mu0, self.tau0 = sigma, mu0, tau0

    def posterior(self):
        prec = 1.0 / self.tau0**2 + self.counts / self.sigma**2
        mean = (self.mu0 / self.tau0**2 + self.sums / self.sigma**2) / prec
        return mean, 1.0 / np.sqrt(prec)

    def select(self):
        mean, sd = self.posterior()
        return argmax_random_ties(mean + sd * self.rng.standard_normal(mean.shape), self.rng)


# --------------------------------------------------------------------------
# Simulation
# --------------------------------------------------------------------------


def simulate(env, policy: Policy, T: int, n_runs: int = 100, seed: int = 0):
    """Run ``policy`` for T rounds on ``n_runs`` independent copies of ``env``.

    Returns a dict with
      regret : (n_runs, T) float32 cumulative *pseudo*-regret sum_t Delta_{A_t}
      counts : (n_runs, K) final pull counts
    Pseudo-regret (expected gap of the chosen arm) has much less variance than
    realized regret and is what the Lai-Robbins bound is about.
    """
    rng = np.random.default_rng(seed)
    policy.reset(n_runs, env.K, rng)
    gaps = env.gaps
    inst = np.empty((n_runs, T), dtype=np.float32)
    for t in range(T):
        arms = policy.select()
        rewards = env.pull(arms, rng)
        policy.update(arms, rewards)
        inst[:, t] = gaps[arms]
    return {"regret": np.cumsum(inst, axis=1, dtype=np.float64).astype(np.float32),
            "counts": policy.counts.copy(), "name": policy.name}