# How Epidemics Spread 5, Exercise 1: testing how the latent period changes timing but not the final size
# Exercise 1: SEIR with R0 = 3 (assumed) and latent periods of 2, 4 and 8 days.
population, gamma, r0, dt = 10_000, 0.2, 3.0, 0.05
beta = r0 * gamma

def seir(latent_days):
    sigma = 1 / latent_days
    S, E, I, R = population - 1.0, 0.0, 1.0, 0.0
    peak_I, peak_t, t = I, 0.0, 0.0
    while E + I >= 0.01 and t < 2000:
        new_exposed = beta * S * I / population * dt
        new_infectious = sigma * E * dt
        new_rec = gamma * I * dt
        S -= new_exposed
        E += new_exposed - new_infectious
        I += new_infectious - new_rec
        R += new_rec
        t += dt
        if I > peak_I:
            peak_I, peak_t = I, t
    return peak_t, peak_I, R / population

print("latent days   peak day   peak infectious   share ever infected")
for latent in [2, 4, 8]:
    day, peak, share = seir(latent)
    print(f"{latent:11d} {day:10.0f} {peak:17.0f} {share:21.1%}")

# Output when run:
# latent days   peak day   peak infectious   share ever infected
#           2         45              2117                 94.1%
#           4         59              1635                 94.1%
#           8         84              1133                 94.1%