"""Multiplayer bandits without communication (collision model).

M players face the same K Bernoulli arms (M <= K).  Each round every player
pulls one arm.  If two or more players pull the same arm they *collide*: all
of them get reward 0 and each of them observes that a collision happened.
Players never exchange messages; they only see their own rewards and
collision indicators.  The ideal is that the M players spread over the M best
arms, so the collective (pseudo-)regret per round is

    sum of the M largest means  -  sum_k mu_k * 1{exactly one player on arm k}.

Three strategies are implemented:

* ``selfish_ucb``     each player runs UCB1 on its own observations, treating
                      a collision as reward 0.  No coordination at all.
* ``musical_chairs``  Rosenski, Shamir & Szlak (2016): uniform random
                      exploration for T0 rounds (also used to estimate M from
                      the collision frequency), then each player picks random
                      arms among its estimated top-M until it gets a round
                      without collision, and sits there forever.
* ``sic_simplified``  a simplified collision-*signalling* protocol in the
                      spirit of SIC-MMAB (Boursier & Perchet, 2019); see the
                      class docstring for exactly what is simplified.
"""
from __future__ import annotations

import math

import numpy as np


def top_m_sum(means, M: int) -> float:
    return float(np.sort(np.asarray(means))[::-1][:M].sum())


def occupancy(arms: np.ndarray, K: int) -> np.ndarray:
    """arms: (R, M) chosen arms -> (R, K) number of players on each arm."""
    R = arms.shape[0]
    flat = (np.arange(R)[:, None] * K + arms).ravel()
    return np.bincount(flat, minlength=R * K).reshape(R, K)


def collisions_and_regret(arms: np.ndarray, means: np.ndarray, M: int):
    """Collision mask (R, M) and per-run pseudo-regret (R,) of one round."""
    K = len(means)
    occ = occupancy(arms, K)
    collided = np.take_along_axis(occ, arms, axis=1) > 1
    regret = top_m_sum(means, M) - (means[None, :] * (occ == 1)).sum(axis=1)
    return collided, regret


# --------------------------------------------------------------------------
# 1. Selfish UCB
# --------------------------------------------------------------------------


def selfish_ucb(means, M: int, T: int, n_runs: int = 50, seed: int = 0):
    """Each player runs UCB1 independently; a collision is just a 0 reward."""
    rng = np.random.default_rng(seed)
    means = np.asarray(means, float)
    K = len(means)
    counts = np.zeros((n_runs, M, K))
    sums = np.zeros((n_runs, M, K))
    inst = np.empty((n_runs, T))
    n_coll = np.zeros(n_runs)
    for t in range(T):
        with np.errstate(divide="ignore", invalid="ignore"):
            idx = sums / counts + np.sqrt(2 * np.log(t + 1) / counts)
        idx = np.where(counts == 0, np.inf, idx)
        # random tie-breaking: players start on different random arms
        idx = idx + 1e-9 * rng.random(idx.shape)
        arms = idx.argmax(axis=2)
        collided, inst[:, t] = collisions_and_regret(arms, means, M)
        rewards = (rng.random(arms.shape) < means[arms]) & ~collided
        np.put_along_axis(counts, arms[..., None],
                          np.take_along_axis(counts, arms[..., None], 2) + 1, 2)
        np.put_along_axis(sums, arms[..., None],
                          np.take_along_axis(sums, arms[..., None], 2) + rewards[..., None], 2)
        n_coll += collided.sum(axis=1)
    return {"name": "selfish UCB", "regret": np.cumsum(inst, axis=1), "collisions": n_coll}


# --------------------------------------------------------------------------
# 2. Musical chairs (Rosenski, Shamir & Szlak 2016)
# --------------------------------------------------------------------------


def estimate_num_players(n_collisions, T0: int, K: int):
    """Invert P(no collision) = (1 - 1/K)^(M-1) under uniform play."""
    frac_free = (T0 - np.asarray(n_collisions, float)) / T0
    with np.errstate(divide="ignore"):
        m_hat = 1 + np.round(np.log(np.maximum(frac_free, 1e-12)) / np.log(1 - 1 / K))
    return np.clip(m_hat, 1, K).astype(int)


def musical_chairs(means, M: int, T: int, T0: int, n_runs: int = 50, seed: int = 0):
    """Musical Chairs: explore uniformly for T0 rounds, then fix on a chair.

    The paper picks T0 from the gap and a confidence level; here T0 is a
    parameter (too small -> wrong top-M or wrong M-hat; too large -> linear
    exploration cost).
    """
    rng = np.random.default_rng(seed)
    means = np.asarray(means, float)
    K = len(means)
    sums = np.zeros((n_runs, M, K))
    counts = np.zeros((n_runs, M, K))
    n_coll_player = np.zeros((n_runs, M))
    inst = np.empty((n_runs, T))
    n_coll = np.zeros(n_runs)
    fixed = np.zeros((n_runs, M), bool)
    seat = np.zeros((n_runs, M), int)
    ranked = m_hat = None
    for t in range(T):
        if t < T0:
            arms = rng.integers(K, size=(n_runs, M))
        else:
            if ranked is None:  # end of exploration: rank arms, estimate M
                est = sums / np.maximum(counts, 1)
                ranked = np.argsort(-est, axis=2)                 # (R, M, K)
                m_hat = estimate_num_players(n_coll_player, T0, K)  # (R, M)
            pick = np.floor(rng.random((n_runs, M)) * m_hat).astype(int)
            candidate = np.take_along_axis(ranked, pick[..., None], 2)[..., 0]
            arms = np.where(fixed, seat, candidate)
        collided, inst[:, t] = collisions_and_regret(arms, means, M)
        n_coll += collided.sum(axis=1)
        if t < T0:
            rew = (rng.random(arms.shape) < means[arms]) & ~collided
            ok = ~collided
            # only collision-free rounds give unbiased samples of mu_k
            r_idx, m_idx = np.nonzero(ok)
            np.add.at(counts, (r_idx, m_idx, arms[ok]), 1)
            np.add.at(sums, (r_idx, m_idx, arms[ok]), rew[ok])
            n_coll_player += collided
        else:
            newly = ~fixed & ~collided
            seat = np.where(newly, arms, seat)
            fixed |= newly
    return {"name": f"musical chairs (T0={T0})", "regret": np.cumsum(inst, axis=1),
            "collisions": n_coll}


# --------------------------------------------------------------------------
# 3. Simplified SIC-MMAB-style protocol
# --------------------------------------------------------------------------


class SICSimplified:
    """Collision-signalling protocol in the spirit of SIC-MMAB.

    Phases (all lengths are known to every player in advance, so players stay
    synchronised without talking):

    1. Orthogonalisation (T_c rounds): musical chairs on all K arms.  A player
       that gets a collision-free round sits on that arm ("seat").  T_c is
       chosen so every player is seated with prob. >= 1 - 1/T (union bound
       with per-round success prob. >= (K-M+1)/K * (1-1/K)^(M-1)).
    2. Roll call (K(K-1)/2 rounds): for each pair of seats a < b, the player
       on seat a (if any) moves to seat b while everyone else stays.  A
       collision happens iff both seats are occupied, and both players see it.
       Afterwards every player knows the set of occupied seats, hence M and
       its own rank j in {0..M-1}.
    3. Exploration, in phases p = 0, 1, 2, ...: with C the current candidate
       set, player j plays C[(j + s) mod |C|] for s = 0..|C| 2^p - 1.  Ranks
       are distinct and |C| >= M, so this sequential hopping is
       collision-free and every arm in C gets 2^p pulls from every player.
    4. Gathering: each follower sends its phase success count for every
       candidate arm to the rank-0 player (leader), in binary, one bit per
       round: bit 1 = move onto the leader's arm and collide, bit 0 = stay.
       Counts are integers in [0, 2^p], so this is exact and costs
       (M-1) |C| ceil(log2(2^p + 1)) rounds.
    5. Decision + broadcast: the leader pools everyone's samples, rejects arm
       k once at least M candidates have LCB above UCB_k (Hoeffding radius
       sqrt(conf_scale * log T / N)), and sends one bit per candidate arm to
       each follower (leader collides on the follower's arm for a 1).
    6. When |C| = M, player j commits to C[j] forever.

    What is simplified compared with SIC-MMAB:
      * statistics are gathered at a single leader (SIC-MMAB has everybody
        broadcast quantised means to everybody) and only the leader decides;
      * arms are not 'accepted' early -- everyone keeps exploring the
        candidate set until exactly M arms remain (SIC-MMAB lets a player
        exploit an accepted arm immediately);
      * the rank/M estimation uses O(K^2) pairwise probes instead of SIC's
        O(K) sequential hopping;
      * exact integer counts are sent instead of quantised empirical means.
    The scheme is still collision-free outside the deliberate signalling
    rounds, and its regret grows like log T (exploration + communication).
    If the orthogonalisation phase fails (prob. <= 1/T) the run is charged the
    maximal per-round regret for the rest of the horizon, so failures are
    not hidden.
    """

    def __init__(self, conf_scale: float = 0.5):
        self.conf_scale = conf_scale

    def run_one(self, means, M: int, T: int, rng: np.random.Generator):
        means = np.asarray(means, float)
        K = len(means)
        best_sum = top_m_sum(means, M)
        segs: list[np.ndarray] = []
        used = 0

        def emit(regrets):
            nonlocal used
            segs.append(np.asarray(regrets, float))
            used += len(segs[-1])

        def round_regret(arms):
            occ = np.bincount(np.asarray(arms), minlength=K)
            return best_sum - float(means[occ == 1].sum()), occ

        # ---- 1. orthogonalisation ------------------------------------
        p0 = (K - M + 1) / K * (1 - 1 / K) ** (M - 1)
        Tc = math.ceil(math.log(M * T) / p0)
        seated = np.zeros(M, bool)
        seat = np.full(M, -1)
        reg = []
        for _ in range(Tc):
            arms = np.where(seated, seat, rng.integers(K, size=M))
            r, occ = round_regret(arms)
            reg.append(r)
            ok = (occ[arms] == 1) & ~seated
            seat[ok] = arms[ok]
            seated |= ok
        emit(reg)
        if not seated.all():
            emit(np.full(max(T - used, 0), best_sum))
            return np.cumsum(np.concatenate(segs))[:T], True

        # ---- 2. roll call (pairwise probes) --------------------------
        occupied_seen = [set() for _ in range(M)]  # each player's knowledge
        owner = {int(s): i for i, s in enumerate(seat)}
        reg = []
        for a in range(K):
            for b in range(a + 1, K):
                arms = seat.copy()
                if a in owner:
                    arms[owner[a]] = b
                r, occ = round_regret(arms)
                reg.append(r)
                if occ[b] > 1:  # both a and b occupied; both players see it
                    occupied_seen[owner[a]].add(b)
                    occupied_seen[owner[b]].add(a)
        emit(reg)
        # every player now knows all occupied seats (its own + those probed)
        ranks = np.empty(M, int)
        for i in range(M):
            occupied = sorted(occupied_seen[i] | {int(seat[i])})
            assert len(occupied) == M  # each player has learned M exactly
            ranks[i] = occupied.index(int(seat[i]))
        # ranks are a permutation of 0..M-1; from here on we index players by
        # rank (player with rank j behaves identically whatever its seat was)
        assert sorted(ranks) == list(range(M))

        # ---- 3-6. explore / gather / decide / broadcast / commit ------
        C = list(range(K))
        S = np.zeros(K)   # pooled successes (as known to the leader)
        N = np.zeros(K)   # pooled pulls
        p = 0
        while used < T:
            c = len(C)
            if c == M:  # commit: rank j sits on C[j]
                r, _ = round_regret(np.array(C))
                emit(np.full(T - used, r))
                break
            reps = 2 ** p
            cycle = [round_regret(np.array([C[(j + s) % c] for j in range(M)]))[0]
                     for s in range(c)]
            emit(np.tile(cycle, reps))
            Carr = np.array(C)
            # every player pulled every candidate arm `reps` times, no collisions
            succ = rng.binomial(reps, means[Carr], size=(M, c))   # row j = rank j
            if M > 1:
                # gathering: follower j -> leader, n_bits per arm, one bit/round
                n_bits = math.ceil(math.log2(reps + 1))
                reg = []
                received = np.zeros(c, dtype=np.int64)
                for follower in range(1, M):
                    for a in range(c):
                        value = 0
                        for b in range(n_bits):
                            bit = (int(succ[follower, a]) >> b) & 1
                            arms = np.arange(M)          # rank j waits on arm j
                            if bit:
                                arms[follower] = 0       # collide on leader's arm
                            r, occ = round_regret(arms)
                            reg.append(r)
                            value |= int(occ[0] > 1) << b  # leader decodes collision
                        received[a] += value
                emit(reg)
                assert np.array_equal(received, succ[1:].sum(axis=0))
                S[Carr] += succ[0] + received
            else:
                S[Carr] += succ[0]
            N[Carr] += M * reps
            # leader's elimination decision on pooled data
            mu_hat = S[Carr] / N[Carr]
            rad = np.sqrt(self.conf_scale * math.log(T) / N[Carr])
            ucb, lcb = mu_hat + rad, mu_hat - rad
            n_better = (lcb[None, :] > ucb[:, None]).sum(axis=1)
            reject = n_better >= M  # at least M arms always survive this rule
            # broadcast: one bit per candidate arm to each follower
            if M > 1:
                reg = []
                for follower in range(1, M):
                    for bit in reject:
                        arms = np.arange(M)
                        if bit:
                            arms[0] = follower  # leader collides on purpose
                        reg.append(round_regret(arms)[0])
                emit(reg)
            C = [k for k, rj in zip(C, reject) if not rj]
            p += 1
        return np.cumsum(np.concatenate(segs))[:T], False

    def run(self, means, M: int, T: int, n_runs: int = 50, seed: int = 0):
        rng = np.random.default_rng(seed)
        curves, fails = [], 0
        for _ in range(n_runs):
            curve, failed = self.run_one(means, M, T, rng)
            curves.append(curve)
            fails += failed
        return {"name": "SIC-style (simplified)", "regret": np.array(curves),
                "failures": fails}


def sic_simplified(means, M: int, T: int, n_runs: int = 50, seed: int = 0):
    return SICSimplified().run(means, M, T, n_runs, seed)