# How Epidemics Spread 3, Exercise 3: fitting the boarding-school counts
# Exercise 3: fit beta and gamma to the 1978 boarding-school in-bed counts (Chapter 1 data).
# Assumptions: the number in bed equals I in the model, and day 1 starts with 3 infectious.
in_bed = [3, 8, 26, 76, 225, 298, 258, 233, 189, 128, 68, 29, 14, 4]
boys, dt = 763, 0.1

def model(beta, gamma):
    S, I = boys - 3.0, 3.0
    out = []
    for day in range(14):
        out.append(I)
        for _ in range(int(1 / dt)):
            new_inf = beta * S * I / boys * dt
            S -= new_inf
            I += new_inf - gamma * I * dt
    return out

best = None
for b in range(100, 301, 5):
    for g in range(10, 101, 5):
        beta, gamma = b / 100, g / 100
        error = sum((m - o) ** 2 for m, o in zip(model(beta, gamma), in_bed))
        if best is None or error < best[0]:
            best = (error, beta, gamma)
error, beta, gamma = best
print(f"best beta {beta:.2f}, gamma {gamma:.2f}, R0 = {beta / gamma:.2f}, mean infectious period {1 / gamma:.1f} days")
print("day  observed  model")
for day, (o, m) in enumerate(zip(in_bed, model(beta, gamma)), start=1):
    print(f"{day:3d} {o:9d} {m:7.0f}")

# The model's final size at this R0 (repeat the self-consistency guess until steady).
import math
z = 0.5
for _ in range(200):
    z = 1 - math.exp(-(beta / gamma) * z)
print(f"model share ever infected: {z:.0%}; recorded: 512 of 763 = {512 / 763:.0%}")

# Output when run:
# best beta 1.75, gamma 0.45, R0 = 3.89, mean infectious period 2.2 days
# day  observed  model
#   1         3       3
#   2         8      10
#   3        26      33
#   4        76      95
#   5       225     210
#   6       298     301
#   7       258     290
#   8       233     227
#   9       189     162
#  10       128     111
#  11        68      75
#  12        29      50
#  13        14      33
#  14         4      22
# model share ever infected: 98%; recorded: 512 of 763 = 67%