# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# UMI counts for one cell type: 4 control and 4 treated mice, 250 cells each, 100 genes.
import numpy as np
MICE, CELLS, GENES, TRUE_DE = 4, 250, 100, 5

def study(rng, effect=1.0, sd_mouse=0.3, mice=MICE):
    """Rows (mouse, condition, counts matrix cells x genes). Genes 1-5 change by `effect` (log2) in treated mice; every gene
    has a mouse-specific shift on the log scale (biological variation between animals); cells are Poisson around their mouse's mean."""
    base = np.log(rng.gamma(2.0, 1.0, GENES) + 0.2)          # gene means per cell (UMIs), fixed per study
    out = []
    for cond in ("control", "treated"):
        for m in range(mice):
            mu = base + rng.normal(0, sd_mouse, GENES)
            if cond == "treated": mu[:TRUE_DE] += effect * np.log(2)
            lib = rng.gamma(10, 0.1, CELLS)[:, None]           # cell-to-cell capture (size factor, mean 1)
            out.append((f"{cond[0].upper()}{m + 1}", cond, rng.poisson(lib * np.exp(mu)[None, :])))
    return out