"""Bandit environments and an experiment runner for the Chapter 2 lab.

You don't need to change anything in this file. Read it once so you know
what the agents in `bandits.py` will be plugged into.

Every agent has the same tiny interface:

    agent.select_action() -> int        # which arm to pull now
    agent.update(action, reward)        # learn from what happened

and every bandit has:

    bandit.k                            # number of arms K
    bandit.pull(a) -> float             # sample a reward from arm a
    bandit.means                        # true mean rewards q_*(a), shape (K,)
    bandit.step()                       # advance time (only the switching bandit changes)
"""

from __future__ import annotations

from typing import Callable

import numpy as np


class BernoulliBandit:
    """K slot machines. Arm a pays 1 with probability means[a], else 0."""

    def __init__(self, means, rng: np.random.Generator):
        self.means = np.asarray(means, dtype=float)
        self.k = len(self.means)
        self.rng = rng

    def pull(self, a: int) -> float:
        return float(self.rng.random() < self.means[a])

    def step(self) -> None:
        pass


class GaussianBandit:
    """The Sutton & Barto 10-armed testbed: q_*(a) ~ N(0, 1), reward ~ N(q_*(a), 1)."""

    def __init__(self, means, rng: np.random.Generator):
        self.means = np.asarray(means, dtype=float)
        self.k = len(self.means)
        self.rng = rng

    def pull(self, a: int) -> float:
        return float(self.means[a] + self.rng.standard_normal())

    def step(self) -> None:
        pass


class SwitchingBernoulliBandit(BernoulliBandit):
    """A nonstationary Bernoulli bandit: after `switch_at` pulls the win
    probabilities are shuffled among the arms, so the best arm suddenly
    becomes some other arm. Everything learned before the switch is now stale."""

    def __init__(self, means, rng: np.random.Generator, switch_at: int = 1000):
        super().__init__(means, rng)
        self.switch_at = switch_at
        self.t = 0

    def step(self) -> None:
        self.t += 1
        if self.t == self.switch_at:
            self.means = self.rng.permutation(self.means)


def run_bandit(agent, bandit, steps: int) -> dict:
    """Let `agent` play `bandit` for `steps` pulls.

    Returns a dict of arrays, each of shape (steps,):
      reward   - the reward received at each step
      optimal  - 1.0 if the pulled arm was (one of) the best arm(s) at that step
      regret   - max_a q_*(a) - q_*(a_t): the expected reward lost at that step
    """
    rewards = np.zeros(steps)
    optimal = np.zeros(steps)
    regret = np.zeros(steps)
    for t in range(steps):
        a = agent.select_action()
        r = bandit.pull(a)
        agent.update(a, r)
        best = bandit.means.max()
        rewards[t] = r
        optimal[t] = float(bandit.means[a] == best)
        regret[t] = best - bandit.means[a]
        bandit.step()
    return {"reward": rewards, "optimal": optimal, "regret": regret}


def run_testbed(
    make_agent: Callable[[int, np.random.Generator], object],
    make_bandit: Callable[[np.random.Generator], object],
    runs: int = 200,
    steps: int = 1000,
    seed: int = 0,
) -> dict:
    """Average an agent's performance over many independent bandit problems.

    make_agent(k, rng)  -> a fresh agent for a K-armed problem
    make_bandit(rng)    -> a fresh bandit (its true means drawn with rng)

    Run i uses seed (seed, i) for the bandit and a separate stream for the agent,
    so two agents tested with the same seed face exactly the same problems.

    Returns a dict of arrays of shape (steps,):
      avg_reward   - mean reward at each step (over runs)
      pct_optimal  - fraction of runs that pulled the best arm at each step
      cum_regret   - mean cumulative regret, sum_{s<=t} (max q_* - q_*(a_s))
    """
    reward = np.zeros(steps)
    optimal = np.zeros(steps)
    regret = np.zeros(steps)
    for i in range(runs):
        env_rng = np.random.default_rng([seed, i, 0])
        agent_rng = np.random.default_rng([seed, i, 1])
        bandit = make_bandit(env_rng)
        agent = make_agent(bandit.k, agent_rng)
        out = run_bandit(agent, bandit, steps)
        reward += out["reward"]
        optimal += out["optimal"]
        regret += out["regret"]
    return {
        "avg_reward": reward / runs,
        "pct_optimal": optimal / runs,
        "cum_regret": np.cumsum(regret / runs),
    }


def bernoulli_testbed(k: int = 10) -> Callable[[np.random.Generator], BernoulliBandit]:
    """K Bernoulli arms with win probabilities drawn uniformly from [0, 1]."""
    return lambda rng: BernoulliBandit(rng.random(k), rng)


def gaussian_testbed(k: int = 10) -> Callable[[np.random.Generator], GaussianBandit]:
    """The Sutton & Barto testbed: q_*(a) ~ N(0, 1)."""
    return lambda rng: GaussianBandit(rng.standard_normal(k), rng)


def switching_testbed(k: int = 10, switch_at: int = 1000) -> Callable[[np.random.Generator], SwitchingBernoulliBandit]:
    """K Bernoulli arms (win probabilities uniform on [0, 1]) that get shuffled after `switch_at` pulls."""
    return lambda rng: SwitchingBernoulliBandit(rng.random(k), rng, switch_at)