# How Epidemics Spread 9, Exercise 1
# Exercise 1: when to close and for how long? 70% contact cut (assumed), R0 = 2.5, infectious period 5 days (assumed).
def run_sir(R0, c, start, stop, days=500, n=1_000_000, i0=10, infectious_days=5.0, dt=0.25):
    """Step a simple S-I-R model forward; contacts are cut by c from day `start` up to day `stop`."""
    gamma = 1.0 / infectious_days
    beta = R0 * gamma
    s, i = n - i0, float(i0)
    peak = i
    t = 0.0
    while t < days:
        factor = (1 - c) if start <= t < stop else 1.0
        new_inf = beta * factor * s * i / n * dt
        s -= new_inf
        i += new_inf - gamma * i * dt
        t += dt
        peak = max(peak, i)
    return peak / n, (n - s) / n

print("start day  length   peak infected   ever infected")
for start in (10, 30, 50):
    for length in (30, 60, 120):
        pk, fin = run_sir(2.5, 0.7, start, start + length)
        print(f"{start:9d} {length:7d} {pk*100:13.2f}% {fin*100:13.1f}%")
# Output when run:
# start day  length   peak infected   ever infected
#        10      30         23.68%          89.6%
#        10      60         23.67%          89.6%
#        10     120         23.67%          89.6%
#        30      30         14.46%          85.3%
#        30      60         13.48%          84.7%
#        30     120         13.38%          84.6%
#        50      30         23.72%          85.2%
#        50      60         23.72%          85.2%
#        50     120         23.72%          85.2%

# Reading the output: the same closure gives three different results depending on when it starts. Started on day 10 it ends (on day 40, 70 or 130) before the epidemic has been seen to grow much, and the peak and the total are the same as with no control (23.7 per cent and 89.6 per cent in the chapter's no-control run). The table does not print the day of the peak, so that the closure only moved the epidemic later is my inference, untested here. Started on day 30, while cases are climbing, it cuts the peak to about 13 or 14 per cent and the total to about 85 per cent. Lengthening it from 60 to 120 days changed the peak by 0.1 point; my explanation is that by about day 90 the epidemic is nearly over, so the extra 60 days have little left to hold back (untested). The 14.46 per cent for 30 days is probably a second, smaller wave after reopening (also an inference). Started on day 50, after the peak that happens around day 41 without control, it cannot lower the peak (23.72 per cent) and trims the total only to 85.2 per cent. All numbers are for assumed values.