# How Epidemics Spread 6, Exercise 1
# Exercise 1: stochastic SIR with R0 = 2 (assumed) for populations of 50, 200 and 1000.
import math
import random

gamma, r0, runs = 0.2, 2.0, 2000
beta = r0 * gamma

def outbreak(population):
    S, I, ever = population - 1, 1, 1
    while I > 0:
        infect_rate = beta * S * I / population
        recover_rate = gamma * I
        if random.random() < infect_rate / (infect_rate + recover_rate):
            S -= 1
            I += 1
            ever += 1
        else:
            I -= 1
    return ever

# Deterministic final share: repeat z = 1 - exp(-R0 * z) until it settles.
z = 0.5
for _ in range(200):
    z = 1 - math.exp(-r0 * z)

random.seed(11)
print(f"deterministic final share: {z:.1%}; 1/R0 = {1 / r0:.1%}")
print("population   fizzled (<=10 cases)   mean share infected in the rest")
for population in [50, 200, 1000]:
    sizes = [outbreak(population) for _ in range(runs)]
    big = [s for s in sizes if s > 10]
    small = len(sizes) - len(big)
    print(f"{population:10d} {small / runs:22.1%} {sum(big) / len(big) / population:30.1%}")

# Output when run:
# deterministic final share: 79.7%; 1/R0 = 50.0%
# population   fizzled (<=10 cases)   mean share infected in the rest
#         50                  50.7%                          75.0%
#        200                  48.2%                          76.7%
#       1000                  49.2%                          77.9%

# Reading the output: about half the introductions fizzle at every population size (48-51%), matching 1/R0 = 50%.
# The rest infect roughly 75-78% of the population, a little under the deterministic 79.7%. The reason (some of the 'rest'
# are unlucky chains that stop early) is my inference, not tested here.