# How Epidemics Spread 8, Exercise 1
# Exercise 1: how the fitting window and the assumed serial interval change the estimate of R (true R = 2.0 assumed).
import math
import random
import statistics

def poisson(lam):
    """Poisson random number (Knuth's method; fine for small lam)."""
    limit, k, p = math.exp(-lam), 0, 1.0
    while True:
        p *= random.random()
        if p <= limit:
            return k
        k += 1

def simulate(R, days, shape, scale, start=20):
    """New cases per day. Each case has Poisson(R) children; each child is infected a whole number
    of days later, drawn from a gamma 'serial interval' (shape, scale) and rounded, minimum 1 day."""
    daily = [0] * (days + 1)
    daily[0] = start
    for day in range(days + 1):
        for _ in range(daily[day]):
            for _ in range(poisson(R)):
                lag = max(1, round(random.gammavariate(shape, scale)))
                if day + lag <= days:
                    daily[day + lag] += 1
    return daily

def growth_rate(daily, first, last):
    """Slope of log(cases) against day, by least squares over days first..last (None if a day has no cases)."""
    if min(daily[first:last + 1]) == 0:
        return None
    xs = list(range(first, last + 1))
    ys = [math.log(daily[d]) for d in xs]
    mx, my = sum(xs) / len(xs), sum(ys) / len(ys)
    return sum((x - mx) * (y - my) for x, y in zip(xs, ys)) / sum((x - mx) ** 2 for x in xs)

def r_to_R(r, shape, scale):
    """Growth rate to R for a gamma serial interval: R = (1 + r * scale) ** shape."""
    return (1 + r * scale) ** shape

random.seed(31)
shape, scale = 4, 1.25
runs = [simulate(2.0, 30, shape, scale) for _ in range(100)]
print("window (days)   mean R estimate   sd   runs used")
for first, last in [(0, 10), (5, 15), (10, 20), (15, 25), (20, 30)]:
    rates = [r for r in (growth_rate(d, first, last) for d in runs) if r is not None]
    est = [r_to_R(r, shape, scale) for r in rates]
    print(f"{first:3d} to {last:2d} {statistics.mean(est):15.2f} {statistics.stdev(est):9.2f} {len(est):8d}")
print()
print("assumed serial-interval sd (mean fixed at 5 days), window 10 to 20:")
rates = [r for r in (growth_rate(d, 10, 20) for d in runs) if r is not None]
for sd in [1.0, 2.5, 4.0]:
    k = (5 / sd) ** 2            # gamma shape from mean and sd
    est = [r_to_R(r, k, 5 / k) for r in rates]
    print(f"assumed sd {sd:.1f} days: mean R estimate {statistics.mean(est):5.2f}")

# Output when run:
# window (days)   mean R estimate   sd   runs used
#   0 to 10            1.95      0.25       75
#   5 to 15            2.05      0.18      100
#  10 to 20            2.02      0.11      100
#  15 to 25            2.00      0.07      100
#  20 to 30            2.00      0.05      100
#
# assumed serial-interval sd (mean fixed at 5 days), window 10 to 20:
# assumed sd 1.0 days: mean R estimate  2.13
# assumed sd 2.5 days: mean R estimate  2.02
# assumed sd 4.0 days: mean R estimate  1.87

# Reading the output: the estimate sits close to the true 2.0 in every window from day 5 onwards, but the spread (sd) shrinks as the window moves later and the counts get larger (0.18 for days 5 to 15, 0.05 for days 20 to 30). The earliest window, days 0 to 10, is the noisiest (sd 0.25) and only 75 of 100 runs could be fitted, because the others had a day with no cases and log(0) is undefined; dropping the slow runs may bias that row (my inference, untested). In this simulation R never changes, so a later window is simply better. In a real epidemic, later windows mix in the effect of controls and of people running out to infect, so a later window does not estimate the same R. In the second table the assumed spread of the serial interval matters less than the assumed mean (in the chapter, a mean of 3 to 7 days moved the estimate from 1.56 to 2.62, while an sd 1.0 to of 1.0 to 4.0 here moved it only from 2.13 to 1.87). All numbers are for assumed values.