# The same mixed model in Python: statsmodels MixedLM (REML). Python 3.13, statsmodels 0.15.0, SciPy 1.18.1.
# statsmodels gives the same estimates as lme4, but its p-values are Wald z-tests without small-sample degrees of
# freedom, so with 16 mice they are smaller than the Satterthwaite or Kenward-Roger values from R.
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from scipy import stats

d = pd.read_csv("cells.csv")                         # columns: group, mouse, cell, y; one row per cell
d["drug"] = (d.group == "drug").astype(int)
fit = smf.mixedlm("y ~ drug", d, groups=d["mouse"]).fit(reml=True)
print(f"Data: {len(d)} cells from {d.mouse.nunique()} mice ({d[d.drug == 1].mouse.nunique()} per group)")
b, se = fit.params["drug"], fit.bse["drug"]
print(f"Fixed effect, drug - vehicle: {b:.3f} (SE {se:.3f}; normal-approximation (Wald z) 95% CI {b - 1.96 * se:.3f} to {b + 1.96 * se:.3f})")
print(f"statsmodels Wald z-test: z = {b / se:.2f}, p = {2 * stats.norm.sf(abs(b / se)):.4f}")
sm, sr = np.sqrt(fit.cov_re.iloc[0, 0]), np.sqrt(fit.scale)
print(f"Random effects: between-mouse SD {sm:.3f}; residual (between-cell) SD {sr:.3f}; ICC {sm ** 2 / (sm ** 2 + sr ** 2):.3f}")
vf = np.var(fit.model.exog @ fit.fe_params.values, ddof=1)
print(f"R-squared (Nakagawa & Schielzeth): marginal {vf / (vf + sm ** 2 + sr ** 2):.3f}, conditional {(vf + sm ** 2) / (vf + sm ** 2 + sr ** 2):.3f}")
m = d.groupby(["mouse", "drug"]).y.mean().reset_index()
t = stats.ttest_ind(m.y[m.drug == 1], m.y[m.drug == 0])
print(f"Check: t-test on the 16 mouse means: difference {m.y[m.drug == 1].mean() - m.y[m.drug == 0].mean():.3f}, t({len(m) - 2}) = {t.statistic:.2f}, p = {t.pvalue:.4f}")