# How Epidemics Spread 4, Exercise 3: uneven coverage hides inside a good average
# Exercise 3: average coverage above the threshold, but uneven (all values assumed).
# Two equal communities; each person has a share 'within' of contacts in their own community.
n, gamma, r0, dt = 5_000, 0.2, 5.0, 0.05
beta = r0 * gamma
coverage = {"A": 0.98, "B": 0.70}       # average 84%, threshold 1 - 1/5 = 80%
print(f"average coverage {(coverage['A'] + coverage['B']) / 2:.0%}, threshold {1 - 1 / r0:.0%}, perfect vaccine")

def simulate(within, seed_in):
    S = {g: n * (1 - c) for g, c in coverage.items()}
    I = {"A": 0.0, "B": 0.0}
    S[seed_in] -= 1.0
    I[seed_in] = 1.0
    start = dict(S)
    while sum(I.values()) >= 0.01:
        force = {}
        for g, other in (("A", "B"), ("B", "A")):
            force[g] = beta * (within * I[g] + (1 - within) * I[other]) / n
        new = {g: S[g] * force[g] * dt for g in S}
        for g in S:
            S[g] -= new[g]
            I[g] += new[g] - gamma * I[g] * dt
    return {g: start[g] - S[g] for g in S}

print("within-community  seeded in   infected in A   infected in B   share of all")
for within in [0.5, 0.8, 0.95]:
    for seed in ["B", "A"]:
        out = simulate(within, seed)
        total = (out["A"] + out["B"]) / (2 * n)
        print(f"{within:16.0%} {seed:>9} {out['A']:15.0f} {out['B']:15.0f} {total:13.1%}")

# Output when run:
# average coverage 84%, threshold 80%, perfect vaccine
# within-community  seeded in   infected in A   infected in B   share of all
#              50%         B               0               4          0.0%
#              50%         A               0               4          0.0%
#              80%         B              10             485          4.9%
#              80%         A              10             483          4.9%
#              95%         B               4             797          8.0%
#              95%         A               4             796          8.0%