# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# Mito stress test in WT and KO cells: OCR (pmol/min) per well at four stages, on several independent plates.
import numpy as np
STAGES = ("basal", "oligomycin", "fccp", "rot_aa")
PER_CELL = {"basal": 2.0, "oligomycin": 0.9, "fccp": 4.0, "rot_aa": 0.3}   # pmol/min per 1,000 cells (WT)

def study(rng, plates=5, wells=10, ko_cells=0.75, ko_rate=1.0, sd_plate=0.20, sd_plate_geno=0.08, sd_well=0.10, sd_cells=0.15):
    """Rows (plate, genotype, well, cells, OCR by stage). KO wells hold fewer cells (`ko_cells` x WT) because KO cells
    grow more slowly; `ko_rate` scales respiration per cell. Plate effects and noise are multiplicative."""
    out = []
    for p in range(1, plates + 1):
        pe = rng.normal(0, sd_plate)
        for g in ("WT", "KO"):
            ge = rng.normal(0, sd_plate_geno)                     # this plate's genotype-specific deviation
            for w in range(wells):
                cells = 20_000 * (ko_cells if g == "KO" else 1) * np.exp(rng.normal(0, sd_cells))
                k = np.exp(pe + ge + rng.normal(0, sd_well)) * cells / 1000 * (ko_rate if g == "KO" else 1)
                out.append((p, g, w + 1, round(cells), *[PER_CELL[s] * k * np.exp(rng.normal(0, 0.05)) for s in STAGES]))   # 5% noise per reading
    return out