# How Epidemics Spread 10, Exercise 3
# Exercise 3: scoring plans. loss = infections per 1000 people + w x full-closure days (cut x mean closed days). w is how many infections you would trade for one full day of closure; all values are assumed.
# 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

plans = {
    "isolate only (25%/day)":            dict(detect=0.25),
    "isolate + closure at 30 found":     dict(detect=0.25, trigger=30),
    "isolate + closure at 200 found":    dict(detect=0.25, trigger=200),
    "slow find (10%/day) + closure at 30": dict(detect=0.10, trigger=30),
}
random.seed(43)
results = {}
for name, kw in plans.items():
    runs = [simulate(**kw) for _ in range(100)]
    attack = sorted(a for a, _, _ in runs)[50]
    full_days = sum(c for _, _, c in runs) / 100 * 0.5
    results[name] = (attack, full_days)
print("loss = infections per 1000 people + w x full-closure days; lowest loss wins")
print(f"{'plan':38s} {'w = 0.5':>9s} {'w = 2':>9s} {'w = 8':>9s}")
for name, (attack, full_days) in results.items():
    print(f"{name:38s} " + " ".join(f"{attack*1000 + w*full_days:9.1f}" for w in (0.5, 2, 8)))
best = [min(results, key=lambda k: results[k][0]*1000 + w*results[k][1]) for w in (0.5, 2, 8)]
print("best plan for each w:", best)
# Output when run:
# loss = infections per 1000 people + w x full-closure days; lowest loss wins
# plan                                     w = 0.5     w = 2     w = 8
# isolate only (25%/day)                     371.1     371.1     371.1
# isolate + closure at 30 found               15.8      59.9     236.3
# isolate + closure at 200 found              20.1      62.9     233.9
# slow find (10%/day) + closure at 30        746.4     791.4     971.4
# best plan for each w: ['isolate + closure at 30 found', 'isolate + closure at 30 found', 'isolate + closure at 200 found']

# Reading the output: isolation alone has a constant loss of 371.1 because it uses no closure days. For w = 0.5 and w = 2 the closure at 30 found cases wins. For w = 8 the closure at 200 found cases is ahead by 2.4 points (233.9 against 236.3), which is probably within the run-to-run scatter of 100 runs (not tested), so call it a tie. The slow-detection plan is worst in every column: a closure cannot make up for finding people slowly. Which plan is best depends on the weight w, which is a value judgement and not something the model can supply.