# How Epidemics Spread 10, Exercise 1
# Exercise 1: when should a closure be triggered? Isolation finds 10% of infectious people a day (assumed), a 50% contact cut lasts 60 days.
# The functions below are the simulator from the chapter.
import math
import random


def poisson(lam):
    """Poisson draw: Knuth's method for small means, a normal approximation for large ones."""
    if lam <= 0:
        return 0
    if lam > 30:
        return max(0, round(random.gauss(lam, math.sqrt(lam))))
    limit, k, p = math.exp(-lam), 0, 1.0
    while True:
        p *= random.random()
        if p <= limit:
            return k
        k += 1

def binom(n, p):
    """Binomial draw: exact for small n, normal approximation for large n."""
    if n <= 0 or p <= 0:
        return 0
    if p >= 1:
        return n
    if n < 40:
        return sum(1 for _ in range(n) if random.random() < p)
    return min(n, max(0, round(random.gauss(n * p, math.sqrt(n * p * (1 - p))))))

def simulate(R0=2.5, detect=0.0, trigger=None, cut=0.5, closure_days=60,
             n=100_000, seeds=10, latent=3.0, infectious=5.0, days=300):
    """Return (share ever infected, peak infectious, closure days used).
    detect: chance per day that an infectious person is found and isolated (0 = never).
    trigger: close when this many cases have been found (None = never close)."""
    s, e, i, q, r = n - seeds, 0, seeds, 0, 0
    beta = R0 / infectious
    found, closed_left, closed_used, peak = 0, 0, 0, seeds
    for day in range(days):
        factor = (1 - cut) if closed_left > 0 else 1.0
        if closed_left > 0:
            closed_left -= 1
            closed_used += 1
        new_e = min(s, poisson(beta * factor * s * i / n))
        new_i = binom(e, 1 / latent)
        new_found = binom(i, detect)
        new_r = binom(i - new_found, 1 / infectious)
        s -= new_e
        e += new_e - new_i
        i += new_i - new_found - new_r
        q += new_found
        r += new_r
        found += new_found
        peak = max(peak, i)
        if trigger is not None and closed_used == 0 and closed_left == 0 and found >= trigger:
            closed_left = closure_days
    return (n - s) / n, peak, closed_used

print("trigger (found cases)   median ever infected   mean closed days")
random.seed(41)
for trig in (30, 100, 300, 1000, 3000):
    runs = [simulate(detect=0.10, trigger=trig) for _ in range(100)]
    attack = sorted(a for a, _, _ in runs)
    print(f"{trig:21d} {attack[50]*100:20.1f}% {sum(c for _, _, c in runs)/100:17.1f}")
# Output when run:
# trigger (found cases)   median ever infected   mean closed days
#                    30                 73.2%              60.0
#                   100                 72.8%              60.0
#                   300                 71.4%              60.0
#                  1000                 67.4%              60.0
#                  3000                 57.7%              60.0

# Reading the output: later triggers did a little better here (73.2% ever infected at 30 found cases, 57.7% at 3,000), with the same 60 closed days each time. My reading, which I did not test, is that a closure placed nearer the peak of the epidemic stops more of the overshoot, as in How Epidemics Spread 9. Even the best row leaves more than half the population infected: a 50% cut for 60 days with slow isolation is not a plan that works. All values are assumed, and the median hides the spread between runs.