# How Epidemics Spread 5, Exercise 3: vaccinating newborns in a population with births and deaths
# Exercise 3: switch on vaccination of a share p of newborns (perfect vaccine, assumed), R0 = 2.5.
# Start from the unvaccinated endemic state, then compare the level after 100 years
# with the prediction mu * (R0 * (1 - p) - 1) / beta (zero or less means the disease dies out).
population = 1_000_000
gamma, r0, dt = 0.2, 2.5, 0.25
mu = 1 / (70 * 365)
beta = r0 * (gamma + mu)

def run(p, years=100):
    I = mu * (r0 - 1) / beta * population              # unvaccinated endemic level
    S = population / r0
    R = population - S - I
    peak_later = 0.0
    for step in range(int(years * 365 / dt)):
        infect = beta * S * I / population
        S += (mu * (1 - p) * population - mu * S - infect) * dt
        I += (infect - (gamma + mu) * I) * dt
        R += (gamma * I + mu * p * population - mu * R) * dt
        peak_later = max(peak_later, I)
    return I, peak_later

start = mu * (r0 - 1) / beta * 100_000
print(f"unvaccinated endemic level at the start: {start:.2f} per 100,000")
print("newborn coverage   predicted level   level after 100 years   highest in the 100 years")
for p in [0.0, 0.3, 0.5, 0.6, 0.7]:
    predicted = max(0.0, mu * (r0 * (1 - p) - 1) / beta) * 100_000
    final, high = run(p)
    print(f"{p:16.0%} {predicted:17.2f} {final / population * 100_000:23.2f} {high / population * 100_000:24.2f}")

# Output when run:
# unvaccinated endemic level at the start: 11.74 per 100,000
# newborn coverage   predicted level   level after 100 years   highest in the 100 years
#               0%             11.74                   11.74                    11.74
#              30%              5.87                    7.21                    11.74
#              50%              1.96                    2.03                    11.74
#              60%              0.00                    0.00                    11.74
#              70%              0.00                    0.00                    11.74