# How Epidemics Spread 3, Exercise 1: same R0, different infectious periods
# Exercise 1: same R0 = 2.5, different infectious periods (assumed values).
population, r0, dt = 10_000, 2.5, 0.05

print("infectious period  beta   gamma  peak day  peak infectious  never infected")
for period in [2, 5, 10]:
    gamma = 1 / period
    beta = r0 * gamma
    S, I, R = population - 1.0, 1.0, 0.0
    t, peak_I, peak_t = 0.0, I, 0.0
    while I >= 0.01:
        new_inf = beta * S * I / population * dt
        new_rec = gamma * I * dt
        S -= new_inf
        I += new_inf - new_rec
        R += new_rec
        t += dt
        if I > peak_I:
            peak_I, peak_t = I, t
    print(f"{period:14d} d {beta:6.3f} {gamma:6.3f} {peak_t:9.1f} {peak_I:16.0f} {S:15.0f}")
print("Same R0, same final size, but a longer infectious period slows everything down.")

# Output when run:
# infectious period  beta   gamma  peak day  peak infectious  never infected
#              2 d  1.250  0.500      13.0             2354            1058
#              5 d  0.500  0.200      32.3             2343            1067
#             10 d  0.250  0.100      64.3             2339            1070
# Same R0, same final size, but a longer infectious period slows everything down.