# How Epidemics Spread 5, Exercise 2: waning immunity: simulated endemic level against the formula
# Exercise 2: SIRS with births. Does the simulated endemic level match the formula?
# (1 - 1/R0) / (1 + gamma / (mu + omega)); all values assumed.
population = 1_000_000
gamma, r0, dt = 0.2, 2.5, 0.25
mu = 1 / (70 * 365)
beta = r0 * (gamma + mu)

def endemic_level(omega, years=60):
    S, I, R = 0.9 * population, 100.0, 0.1 * population
    for _ in range(int(years * 365 / dt)):
        infect = beta * S * I / population
        S += (mu * population - mu * S - infect + omega * R) * dt
        I += (infect - (gamma + mu) * I) * dt
        R += (gamma * I - (mu + omega) * R) * dt
    return I / population

print("immunity lasts (years)   simulated per 100,000   formula per 100,000")
for years_of_immunity in [0.5, 1, 3]:
    omega = 1 / (years_of_immunity * 365)
    simulated = endemic_level(omega)
    formula = (1 - 1 / r0) / (1 + gamma / (mu + omega))
    print(f"{years_of_immunity:22.1f} {simulated * 100_000:23.1f} {formula * 100_000:21.1f}")

# Output when run:
# immunity lasts (years)   simulated per 100,000   formula per 100,000
#                    0.5                  1612.1                1611.1
#                    1.0                   822.5                 822.2
#                    3.0                   284.4                 284.4