# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# Heterozygous x heterozygous mouse crosses: each litter has 5-10 pups at birth; each pup is wild-type, heterozygous
# or knockout with probabilities 1/4, 1/2, 1/4 (Mendel). Pups survive to weaning (genotyping at postnatal day 21)
# with probability 0.95, multiplied by `viability` for knockouts (1 = no lethality, 0.6 = 40% of knockouts lost).
import numpy as np
GENOTYPES = np.array(["wild-type", "heterozygous", "knockout"])

def study(rng, litters, viability=1.0, survive=0.95):
    rows = []
    for li in range(litters):
        for p in range(rng.integers(5, 11)):
            g = rng.choice(3, p=[0.25, 0.5, 0.25])
            if rng.random() < survive * (viability if g == 2 else 1.0):
                rows.append((f"L{li + 1}", p + 1, "female" if rng.random() < 0.5 else "male", GENOTYPES[g]))
    return rows

def counts(rng, litters, viability=1.0, survive=0.95):
    """Genotype counts at weaning (wild-type, heterozygous, knockout) for the same model, drawn without pup rows."""
    born = rng.multinomial(rng.integers(5, 11, size=litters).sum(), [0.25, 0.5, 0.25])
    return rng.binomial(born, survive * np.array([1.0, 1.0, viability]))