# How Epidemics Spread 7, Exercise 3
# Exercise 3: how many people must be vaccinated on a hub network, chosen at random or busiest first (T = 0.25 assumed).
import random

def random_network(n, mean_degree):
    """Random network: pick random pairs until there are n * mean_degree / 2 links."""
    nbrs = [set() for _ in range(n)]
    links = 0
    while links < n * mean_degree // 2:
        a, b = random.randrange(n), random.randrange(n)
        if a != b and b not in nbrs[a]:
            nbrs[a].add(b)
            nbrs[b].add(a)
            links += 1
    return nbrs

def clustered_network(n, mean_degree, rewire):
    """Ring of neighbours (everyone knows the people nearby); each link is rewired at random with chance `rewire`."""
    nbrs = [set() for _ in range(n)]
    for i in range(n):
        for step in range(1, mean_degree // 2 + 1):
            j = (i + step) % n
            if random.random() < rewire:
                j = random.randrange(n)
                while j == i or j in nbrs[i]:
                    j = random.randrange(n)
            nbrs[i].add(j)
            nbrs[j].add(i)
    return nbrs

def hub_network(n, m):
    """Growth with preferential attachment: each newcomer links to m people, favouring the already well connected."""
    nbrs = [set() for _ in range(n)]
    pool = []                      # a person appears here once per link they have
    targets = list(range(m))
    for new in range(m, n):
        for t in set(targets):
            nbrs[new].add(t)
            nbrs[t].add(new)
        pool.extend(targets)
        pool.extend([new] * m)
        targets = [random.choice(pool) for _ in range(m)]
    return nbrs

def describe(nbrs):
    n = len(nbrs)
    degrees = [len(x) for x in nbrs]
    mean = sum(degrees) / n
    excess = (sum(d * d for d in degrees) / n - mean) / mean   # average number of *other* contacts a contact has
    tri, trip = 0, 0
    for i in range(n):
        d = degrees[i]
        if d >= 2:
            trip += d * (d - 1) / 2
            tri += sum(len(nbrs[i] & nbrs[j]) for j in nbrs[i]) / 2
    return mean, max(degrees), excess, tri / trip


def outbreak(nbrs, T, blocked=frozenset()):
    """One outbreak from a random unblocked person; returns how many people were ever infected."""
    start = random.randrange(len(nbrs))
    while start in blocked:
        start = random.randrange(len(nbrs))
    reached, frontier = {start}, [start]
    while frontier:
        following = []
        for u in frontier:
            for v in nbrs[u]:
                if v not in reached and v not in blocked and random.random() < T:
                    reached.add(v)
                    following.append(v)
        frontier = following
    return len(reached)

def summarise(nbrs, T, runs, blocked=frozenset()):
    sizes = [outbreak(nbrs, T, blocked) for _ in range(runs)]
    big = [s for s in sizes if s > 20]
    fizzled = 1 - len(big) / runs
    share = sum(big) / len(big) / len(nbrs) if big else 0.0   # 0 if none took off
    return fizzled, share, sum(sizes) / runs / len(nbrs)


random.seed(23)
n, T, runs = 2000, 0.25, 300
g = hub_network(n, 3)
ranked = sorted(range(n), key=lambda i: len(g[i]), reverse=True)
print("vaccinated   random choice   busiest first")
for percent in [0, 1, 2, 3, 5, 8, 10]:
    count = n * percent // 100
    chosen_random = frozenset(random.sample(range(n), count))
    chosen_busiest = frozenset(ranked[:count])
    a = summarise(g, T, runs, chosen_random)[2]
    b = summarise(g, T, runs, chosen_busiest)[2]
    print(f"{percent:9d}% {a:15.1%} {b:15.1%}")

# Output when run:
# vaccinated   random choice   busiest first
#         0%           32.3%           32.5%
#         1%           30.8%           16.3%
#         2%           31.9%            7.6%
#         3%           27.9%            2.8%
#         5%           29.1%            0.5%
#         8%           23.5%            0.2%
#        10%           20.9%            0.1%

# Reading the output: choosing the busiest people works far better than choosing at random. With 5% vaccinated, busiest-first gives 0.5% infected on average against 29.1% for random choice; 3% busiest-first already gives 2.8%. Random choice still leaves 20.9% at 10%. This assumes you know who the busiest people are, which is often the hard part, and the network is assumed static. All numbers are for assumed T and sizes, on one network drawn with one seed.