# How Epidemics Spread 2, Exercise 3: the limited-growth model for several R0 values
# Same limited-growth model as the chapter, for several R0 values.
# Assumed: population 10,000, generation time 4 days, one starting case.
def run(r0, population=10_000, T=4, days=200):
    susceptible = population - 1.0
    new = [1.0] + [0.0] * (T - 1)
    for day in range(T, days):
        cases = min(susceptible, r0 * (susceptible / population) * new[day - T])
        susceptible -= cases
        new.append(cases)
    peak = max(range(len(new)), key=lambda d: new[d])
    return peak, new[peak], population - susceptible

print("   R0  peak day  peak new  ever infected")
for r0 in [0.9, 1.0, 1.5, 2.0, 3.0, 5.0]:
    peak, size, total = run(r0)
    print(f"{r0:5.1f}  {peak:8d}  {size:8.1f}  {total:7.0f} ({total / 10_000:.0%})")

# Output when run:
#    R0  peak day  peak new  ever infected
#   0.9         0       1.0       10 (0%)
#   1.0         0       1.0       48 (0%)
#   1.5        76     723.9     6185 (62%)
#   2.0        48    1877.2     8688 (87%)
#   3.0        32    3946.5    10000 (100%)
#   5.0        24    6402.3    10000 (100%)