# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# C. elegans lifespan: in each of 3 independent trials, control and treated worms are kept on 3 plates of 35 each and
# scored daily. Lifespans follow a Gompertz distribution (median about 18 days). Trials differ (both conditions shift
# together), the treatment effect varies a little between trials, and plates differ; about 12% of worms are censored
# (crawled off, bagged or lost) at a random day before death.
import numpy as np
TRIALS, PLATES, WORMS = 3, 3, 35

def study(rng, ratio=1.0, b=0.25, median=18.0, sd_trial=0.10, sd_int=0.04, sd_plate=0.08, p_cens=0.12):
    a = b * np.log(2) / (np.exp(b * median) - 1)                       # Gompertz rate giving the chosen median
    rows = []
    for t in range(TRIALS):
        e_t = rng.normal(0, sd_trial)
        for cond, r in (("control", 1.0), ("treated", ratio)):
            e_tc = e_t + rng.normal(0, sd_int)
            for p in range(PLATES):
                scale = r * np.exp(e_tc + rng.normal(0, sd_plate))             # multiplies every lifespan on the plate
                life = scale * np.log(1 - b / a * np.log(rng.uniform(size=WORMS))) / b
                cens = rng.uniform(size=WORMS) < p_cens; tc = life * rng.uniform(size=WORMS)
                day = np.ceil(np.where(cens, tc, life)).astype(int)
                rows += [(t + 1, cond, f"{cond[0]}{t + 1}{p + 1}", int(dd), int(not c)) for dd, c in zip(day, cens)]
    return rows