# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# 16S read counts for 30 taxa in mice housed 4 per cage; each genotype in its own cages (a common design).
import numpy as np
NTAXA, CAGES, PER_CAGE, DEPTH = 30, 4, 4, 20_000
BASE = np.log(np.r_[0.20, np.geomspace(0.12, 0.002, NTAXA - 1)])   # taxon 1 is common (about 20% of reads)

def study(rng, bloom=4.0, sd_cage=0.38, sd_mouse=0.45):
    """Returns rows (mouse, genotype, cage, counts). Genotype changes only taxon 1, by `bloom`-fold in absolute terms;
    mice sharing a cage share a cage-specific shift in every taxon (cage effect). Cage and mouse variances (0.38², 0.45²)
    are in the ratio of the caging and inter-individual shares reported by Hildebrand et al. 2013 (31.7% and 45.5%)."""
    rows = []
    for gi, g in enumerate(("WT", "KO")):
        for c in range(CAGES):
            cage_eff = rng.normal(0, sd_cage, NTAXA)
            for m in range(PER_CAGE):
                logabs = BASE + cage_eff + rng.normal(0, sd_mouse, NTAXA)
                if g == "KO": logabs[0] += np.log(bloom)
                p = np.exp(logabs); p /= p.sum()
                rows.append((f"{g}{c * PER_CAGE + m + 1}", g, f"{g}-cage{c + 1}", rng.multinomial(DEPTH, p)))
    return rows