# How Epidemics Spread 10, Exercise 2
# Exercise 2: strength and length of a closure. Isolation finds 10% a day (assumed), the closure starts at 100 found cases. "Full-closure days" is cut x mean closed days, a crude cost.
# 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("cut   days   median ever infected   mean full-closure days (cut x days)")
random.seed(42)
for cut in (0.3, 0.5, 0.7):
    for length in (30, 60, 120):
        runs = [simulate(detect=0.10, trigger=100, cut=cut, closure_days=length) for _ in range(100)]
        attack = sorted(a for a, _, _ in runs)
        used = sum(c for _, _, c in runs) / 100
        print(f"{cut*100:3.0f}% {length:6d} {attack[50]*100:20.1f}% {used*cut:30.1f}")
# Output when run:
# cut   days   median ever infected   mean full-closure days (cut x days)
#  30%     30                 72.8%                            9.0
#  30%     60                 71.5%                           18.0
#  30%    120                 63.2%                           36.0
#  50%     30                 73.1%                           15.0
#  50%     60                 72.7%                           30.0
#  50%    120                 72.4%                           60.0
#  70%     30                 73.0%                           21.0
#  70%     60                  0.9%                           42.0
#  70%    120                  0.8%                           84.0

# Reading the output: for cuts of 30% and 50%, and for a 70% cut lasting only 30 days, the final share infected stays at 71-73%, except 30% for 120 days (63.2%). A 70% cut for 60 or 120 days brings it down to under 1%. The cost in full-closure days is 42 for the 60-day case and 84 for the 120-day case, so in these runs the extra 60 days of a 70% cut bought nothing. There is a threshold: the closure has to push the effective reproduction number well below 1 for long enough. The numbers are assumed and the model is a simple one.