# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# A treatment is given to pregnant dams; each dam's litter has 4-10 pups, and every pup is weighed at weaning
# (postnatal day 21). Pup weight varies between litters (SD 0.8 g) and between pups of a litter (SD 0.6 g); pups in
# larger litters are lighter (0.25 g per extra pup) and males are 0.4 g heavier. `effect` is the treatment's change in
# mean weight (g).
import numpy as np

def study(rng, dams=6, effect=0.0, sd_litter=0.8, sd_pup=0.6):
    rows = []
    for g, e in (("control", 0.0), ("treated", effect)):
        for k in range(dams):
            size = int(rng.integers(4, 11)); u = rng.normal(0, sd_litter)
            for p in range(size):
                male = rng.random() < 0.5
                w = 10.0 + e + u - 0.25 * (size - 7) + 0.4 * male + rng.normal(0, sd_pup)
                rows.append((g, f"{g[0].upper()}{k + 1}", size, p + 1, "male" if male else "female", round(w, 2)))
    return rows