"""Epidemic gossip: one rumor, synchronous rounds, three arrow directions.

Every node contacts `fanout` uniformly random peers every round. The three
modes differ only in which way the rumor travels along that contact:

    push      a caller that knows tells a callee that does not
    pull      a caller that does not know asks, and learns if the callee does
    pushpull  both directions on the same call

The contact schedule is drawn ONCE from a seeded `random.Random` and handed
to every mode, so three runs of the same seed disagree about direction and
never about who talked to whom. There is no wall clock anywhere in this file:
a round is an integer, and every number it produces is reproducible.

    schedule(n, fanout, rounds, seed)  -> the shared contact plan
    simulate(n, sched, mode, ...)      -> coverage curve, calls placed
    callers(mode, n, known, lazy)      -> who bothers to call this round
    first_round_at(curve, n, frac)     -> round coverage crossed a fraction
    band_costs(curves, n, fanout, ...) -> where the messages actually went
"""

import random

MODES = ("push", "pull", "pushpull")


def schedule(n, fanout, rounds, seed):
    """plan[r][i] = the `fanout` peers node i calls in round r.

    Peers are drawn with replacement from the other n-1 nodes: `j >= i` is
    bumped so a node never calls itself. Both of those are measured to be
    inert at n=1024 (see commentary section 5) — the schedule is generated
    up front purely so that push, pull and push-pull can replay it.
    """
    rng = random.Random(seed)
    plan = []
    for _ in range(rounds):
        row = []
        for i in range(n):
            peers = []
            for _ in range(fanout):
                j = rng.randrange(n - 1)
                if j >= i:
                    j += 1
                peers.append(j)
            row.append(tuple(peers))
        plan.append(tuple(row))
    return plan


def simulate(n, sched, mode="pushpull", max_rounds=None, start=0,
             lazy=False, live=False):
    """Run one epidemic over a fixed contact schedule.

    lazy   charge only the nodes with a reason to call. A push from a node
           with no news, or a pull by a node that already knows, cannot teach
           anyone anything, so this changes the bill and not the outcome.
    live   read the infection set live instead of from a start-of-round
           snapshot, letting a node infected earlier in this round spread
           within it. This is the load-bearing line; False is the honest
           model of a synchronous round.

    Returns a dict whose `curve[r]` is how many nodes know at the START of
    round r. A run that exhausts `max_rounds` comes back converged=False
    rather than looping until it finishes.
    """
    if mode not in MODES:
        raise ValueError(f"mode must be one of {MODES}, got {mode!r}")
    fanout = len(sched[0][0])
    cap = len(sched) if max_rounds is None else min(max_rounds, len(sched))
    inf = bytearray(n)
    inf[start] = 1
    known = 1
    curve = [known]
    calls = 0
    for r in range(cap):
        if known == n:
            break
        snap = inf if live else bytes(inf)   # <-- the load-bearing line
        calls += callers(mode, n, known, lazy) * fanout
        newly = set()
        for i, peers in enumerate(sched[r]):
            for j in peers:
                if mode != "pull" and snap[i] and not snap[j]:
                    newly.add(j)
                    inf[j] = 1
                if mode != "push" and snap[j] and not snap[i]:
                    newly.add(i)
                    inf[i] = 1
        known += len(newly)
        curve.append(known)
    return {"n": n, "mode": mode, "fanout": fanout, "curve": curve,
            "known": known, "converged": known == n, "calls": calls,
            "rounds": len(curve) - 1}


def callers(mode, n, known, lazy=False):
    """How many nodes actually place a call in a round that starts with
    `known` nodes informed.

    Naive accounting charges everybody every round. Lazy accounting is the
    steelman for push: a node with nothing to say stays silent, and under
    pull a node that already knows has nothing to ask for. Push-pull cannot
    be lazy — every node has one of the two reasons, always.
    """
    if not lazy or mode == "pushpull":
        return n
    return known if mode == "push" else n - known


def first_round_at(curve, n, frac):
    """First round index at which coverage reached `frac` of the cluster,
    or None if the run never got there."""
    target = frac * n
    for r, c in enumerate(curve):
        if c >= target:
            return r
    return None


def band_costs(curves, n, fanout, bands, mode="push", lazy=True):
    """Attribute every call placed to the coverage band its round STARTED in.

    A round that begins with c nodes informed places callers(...) * fanout
    calls and ends with curve[r+1] informed, so the spending and the nodes it
    bought are charged to the same band. That pairing is the entire cost
    argument: the bands cover equal fractions of the cluster at wildly
    unequal prices.
    """
    calls = {b: 0 for b in bands}
    gained = {b: 0 for b in bands}
    for curve in curves:
        for r in range(len(curve) - 1):
            frac = curve[r] / n
            for b in bands:
                if b[0] <= frac < b[1]:
                    calls[b] += callers(mode, n, curve[r], lazy) * fanout
                    gained[b] += curve[r + 1] - curve[r]
                    break
    return calls, gained