# How Epidemics Spread 6, Exercise 3
# Exercise 3: same R0, different dispersion. R0 = 1.2 and 2.5 and k = 0.16, 1 and infinity are assumed.
import math
import random

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

def offspring(r0, k):
    if k == float("inf"):
        return poisson(r0)
    return poisson(random.gammavariate(k, r0 / k))

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

random.seed(31)
runs = 3000
print("  R0      k   P(case infects nobody)   fizzle: formula   simulated")
for r0 in [1.2, 2.5]:
    for k in [0.16, 1.0, float("inf")]:
        if k == float("inf"):
            p0, gen = math.exp(-r0), (lambda q: math.exp(r0 * (q - 1)))
        else:
            p0, gen = (1 + r0 / k) ** (-k), (lambda q: (1 + r0 * (1 - q) / k) ** (-k))
        q = 0.0
        for _ in range(3000):
            q = gen(q)
        sim = 1 - sum(escapes(r0, k) for _ in range(runs)) / runs
        label = "inf" if k == float("inf") else f"{k:.2f}"
        print(f"{r0:5.1f} {label:>6} {p0:22.1%} {q:17.3f} {sim:11.3f}")

# Output when run:
#   R0      k   P(case infects nobody)   fizzle: formula   simulated
#   1.2   0.16                  71.0%             0.952       0.948
#   1.2   1.00                  45.5%             0.833       0.835
#   1.2    inf                  30.1%             0.686       0.699
#   2.5   0.16                  63.8%             0.795       0.792
#   2.5   1.00                  28.6%             0.400       0.398
#   2.5    inf                   8.2%             0.107       0.107

# Reading the output: at R0 = 1.2 the chance of fizzling is 0.952 with k = 0.16 against 0.686 with Poisson offspring;
# at R0 = 2.5 it is 0.795 against 0.107. Same average R0, very different chance that a single introduction goes nowhere.