# Data-generating model shared by make_data.py and simulation.py. Python 3.13, NumPy.
# Each mouse has an anxiety level (higher = less anxious) and an activity level. A 5-min elevated plus maze gives arm
# entries and time in the open arms; a 10-min open field gives distance, time and entries in the centre, and rearing.
# Count measures (open-arm entries, centre entries) depend on activity as well as anxiety; percentages much less so.
import numpy as np
MICE = 10
COLS = ["pct_open_time", "pct_open_entries", "open_entries", "closed_entries", "pct_centre_time", "centre_entries", "distance_m", "rearing"]
logit_inv = lambda x: 1 / (1 + np.exp(-x))

def mouse(rng, anx=0.0, act=1.0):
    a = rng.normal(anx, 1.0); L = act * np.exp(rng.normal(0, 0.25))
    tot = rng.poisson(20 * L); p_open = logit_inv(-1.0 + 0.6 * a + rng.normal(0, 0.5)); op = rng.binomial(tot, p_open)
    pct_time = 0.0 if op == 0 else 100 * logit_inv(-1.6 + 0.7 * a + rng.normal(0, 0.5))      # no open entries, no open time
    pct_entries = 100 * op / tot if tot > 0 else 0.0
    dist = 30 * L * np.exp(rng.normal(0, 0.15)); pc = logit_inv(-2.2 + 0.5 * a + rng.normal(0, 0.5))
    return [round(pct_time, 1), round(pct_entries, 1), op, tot - op, round(100 * pc, 1), rng.poisson(3 * dist * pc), round(dist, 1), rng.poisson(40 * L)]

def study(rng, anx=0.0, act=1.0):
    return [(g, f"{g[0].upper()}{m + 1}", *mouse(rng, a, c)) for g, a, c in (("wild-type", 0.0, 1.0), ("knockout", anx, act)) for m in range(MICE)]