# How Epidemics Spread 6, Exercise 2
# Exercise 2: branching process with Poisson offspring; chance that 1, 3 or 5 introductions take off.
import math
import random

def poisson(lam):
    limit, k, p = math.exp(-lam), 0, 1.0
    while True:
        p *= random.random()
        if p <= limit:
            return k
        k += 1

def escapes(r0, cap=500):
    active, total = 1, 1
    while active > 0 and total < cap:
        active = sum(poisson(r0) for _ in range(active))
        total += active
    return active > 0

random.seed(21)
runs = 4000
print("R0    q (one seed)   takes off: 1 seed   3 seeds   5 seeds      simulated: 1 seed   3 seeds")
for r0 in [1.1, 1.5, 2.0, 3.0]:
    q = 0.0
    for _ in range(3000):
        q = math.exp(r0 * (q - 1))
    sim1 = sum(escapes(r0) for _ in range(runs)) / runs
    sim3 = sum(any(escapes(r0) for _ in range(3)) for _ in range(runs)) / runs
    print(f"{r0:3.1f} {q:14.3f} {1 - q:19.3f} {1 - q**3:9.3f} {1 - q**5:9.3f} {sim1:20.3f} {sim3:9.3f}")

# Output when run:
# R0    q (one seed)   takes off: 1 seed   3 seeds   5 seeds      simulated: 1 seed   3 seeds
# 1.1          0.824               0.176     0.441     0.620                0.172     0.446
# 1.5          0.417               0.583     0.927     0.987                0.589     0.926
# 2.0          0.203               0.797     0.992     1.000                0.805     0.994
# 3.0          0.060               0.940     1.000     1.000                0.933     1.000

# Reading the output: the formula column and the simulated column agree to within run-to-run noise (4,000 runs each).
# At R0 = 1.1 even five separate introductions take off only about 62% of the time; at R0 = 3, one is nearly enough.