# Cosinor analysis (24-h period) by least squares, per genotype and compared between genotypes.
# Python 3.13, NumPy, pandas, SciPy.
import numpy as np, pandas as pd
from scipy import stats

d = pd.read_csv("rhythm.csv")                        # columns: genotype, zt, mouse, expr; one row per mouse
w = 2 * np.pi / 24
d["c"], d["s"] = np.cos(w * d.zt), np.sin(w * d.zt)

def ols(X, y):
    b, *_ = np.linalg.lstsq(X, y, rcond=None); res = y - X @ b; dfe = len(y) - X.shape[1]
    s2 = res @ res / dfe; V = s2 * np.linalg.inv(X.T @ X); return b, V, res @ res, dfe

def describe(b1, b2, V):                             # amplitude and peak time with delta-method SEs
    A = np.hypot(b1, b2); phi = np.arctan2(b2, b1)
    gA = np.array([b1 / A, b2 / A]); gP = np.array([-b2 / A ** 2, b1 / A ** 2])
    return A, np.sqrt(gA @ V @ gA), (phi / w) % 24, np.sqrt(gP @ V @ gP) / w

for g in ("WT", "KO"):
    x = d[d.genotype == g]; X = np.column_stack([np.ones(len(x)), x.c, x.s])
    b, V, sse, dfe = ols(X, x.expr.to_numpy())
    sst0 = ((x.expr - x.expr.mean()) ** 2).sum(); F = ((sst0 - sse) / 2) / (sse / dfe)
    A, sA, pk, sP = describe(b[1], b[2], V[1:, 1:]); tq = stats.t.ppf(0.975, dfe)
    print(f"{g}: mesor {b[0]:.2f}; amplitude {A:.3f} (95% CI {A - tq * sA:.3f} to {A + tq * sA:.3f}); "
          f"peak at ZT {pk:.1f} (95% CI {(pk - tq * sP) % 24:.1f} to {(pk + tq * sP) % 24:.1f}); zero-amplitude F(2, {dfe}) = {F:.1f}, p = {stats.f.sf(F, 2, dfe):.2g}")

k = (d.genotype == "KO").astype(float).to_numpy()
Xf = np.column_stack([np.ones(len(d)), d.c, d.s, k, k * d.c, k * d.s]); y = d.expr.to_numpy()
bf, Vf, ssef, dfef = ols(Xf, y)
br, Vr, sser, dfer = ols(np.column_stack([np.ones(len(d)), d.c, d.s, k]), y)
F = ((sser - ssef) / 2) / (ssef / dfef)
print(f"Rhythm differs between genotypes (KO x cos and KO x sin terms): F(2, {dfef}) = {F:.2f}, p = {stats.f.sf(F, 2, dfef):.2g}")
bw, bk = bf[1:3], bf[1:3] + bf[4:6]
Aw, Ak = np.hypot(*bw), np.hypot(*bk)
g = np.zeros(6); g[1:3] = -bw / Aw; g[1:3] += bk / Ak; g[4:6] = bk / Ak       # gradient of Ak - Aw
sd = np.sqrt(g @ Vf @ g); tq = stats.t.ppf(0.975, dfef)
print(f"Amplitude, KO - WT: {Ak - Aw:.3f} (95% CI {Ak - Aw - tq * sd:.3f} to {Ak - Aw + tq * sd:.3f})")
pw, pk_ = np.arctan2(bw[1], bw[0]), np.arctan2(bk[1], bk[0]); dp = (pk_ - pw + np.pi) % (2 * np.pi) - np.pi
gp = np.zeros(6); gp[1:3] = np.array([-bk[1], bk[0]]) / Ak ** 2 - np.array([-bw[1], bw[0]]) / Aw ** 2; gp[4:6] = np.array([-bk[1], bk[0]]) / Ak ** 2
sp = np.sqrt(gp @ Vf @ gp) / w
print(f"Peak time, KO - WT: {dp / w:.1f} h (95% CI {dp / w - tq * sp:.1f} to {dp / w + tq * sp:.1f})")
pt = [t for t in sorted(d.zt.unique())]
ps = [stats.ttest_ind(d.expr[(d.zt == t) & (d.genotype == "KO")], d.expr[(d.zt == t) & (d.genotype == "WT")], equal_var=True).pvalue for t in pt]
print(f"For comparison, separate t-tests at each of {len(pt)} time points: p < 0.05 at ZT {', '.join(str(t) for t, p in zip(pt, ps) if p < 0.05) or 'none'}")
print(f"With a Bonferroni correction for {len(pt)} tests (p < 0.05/{len(pt)}): ZT {', '.join(str(t) for t, p in zip(pt, ps) if p < 0.05 / len(pt)) or 'none'}")