# How Epidemics Spread 8, Exercise 2
# Exercise 2: infections are flat (true R about 1) but testing doubles every 14 days. Everything is assumed.
import math
import random

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(41)
sens, spec, prevalence = 0.80, 0.98, 0.01        # assumed; the true prevalence never changes
q = sens * prevalence + (1 - spec) * (1 - prevalence)     # chance a random test is positive
tests, positives = [], []
for day in range(29):
    t = int(2000 * 2 ** (day / 14))                # assumed testing effort: doubles every 14 days
    tests.append(t)
    positives.append(sum(random.random() < q for _ in range(t)))

shape, scale = 4, 1.25
r_cases = growth_rate(positives, 0, 28)
share = [p / t for p, t in zip(positives, tests)]
print(f"day 0:  {tests[0]} tests, {positives[0]} positives, {share[0]:.2%} positive")
print(f"day 28: {tests[-1]} tests, {positives[-1]} positives, {share[-1]:.2%} positive")
print(f"growth rate of raw positives per day: {r_cases:.3f}  -> 'R' = {r_to_R(r_cases, shape, scale):.2f}")
r_share = growth_rate(share, 0, 28)
print(f"growth rate of share positive per day: {r_share:.3f}  -> 'R' = {r_to_R(r_share, shape, scale):.2f}")

# Output when run:
# day 0:  2000 tests, 55 positives, 2.75% positive
# day 28: 8000 tests, 210 positives, 2.62% positive
# growth rate of raw positives per day: 0.050  -> 'R' = 1.28
# growth rate of share positive per day: 0.001  -> 'R' = 1.00

# Reading the output: the true share infected never changes (assumed 1%), yet the number of positives grows from 55 to 210 simply because the number of tests grows from 2,000 to 8,000. Fitting the raw positives gives a growth rate of 0.050 per day and an apparent 'R' of 1.28, an epidemic that is not there. The share of tests that are positive stays near 2.7% (2.75% to 2.62%), with a growth rate of 0.001 and an 'R' of 1.00. Dividing by the number of tests is a simple fix when tests are a random sample of people. If tests go increasingly to people with symptoms, the share positive also moves, and no such simple fix exists. The testing numbers, sensitivity and specificity are assumed.