# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# Fear extinction: 20 tone presentations in one session; freezing scored in each of the 30 seconds of every tone.
# On the logit scale, freezing starts near 70% and declines across tones; mice differ in their overall level and,
# a little, in how fast they extinguish; trials scatter around each mouse's curve.
import numpy as np
TRIALS, SECONDS, MICE = 20, 30, 10

def mouse(rng, decline=0.12, start=0.85, sd_mouse=0.8, sd_slope=0.04, sd_trial=0.5):
    b0 = start + rng.normal(0, sd_mouse); b1 = decline + rng.normal(0, sd_slope)
    lp = b0 - b1 * np.arange(TRIALS) + rng.normal(0, sd_trial, TRIALS)
    return rng.binomial(SECONDS, 1 / (1 + np.exp(-lp)))              # seconds freezing out of 30, per tone

def study(rng, ko_decline=0.12):
    return [(g, m + 1, mouse(rng, decline=dec)) for g, dec in (("wild-type", 0.12), ("knockout", ko_decline)) for m in range(MICE)]