import pandas as pd
import statsmodels.formula.api as smf
from scipy import stats
d = pd.read_csv("density.csv")
d["genotype"] = pd.Categorical(d.genotype, ["WT", "KO"])
fmt = lambda p: "< 0.001" if p < 0.001 else f"= {p:.3f}"
for g in ["WT", "KO"]:
    s = d[d.genotype == g]
    r = stats.pearsonr(s.age_weeks, s.density)
    m = smf.ols("density ~ age_weeks", s).fit()
    lo, hi = m.conf_int().loc["age_weeks"]
    print(f"{g}: r = {r.statistic:.2f}, p {fmt(r.pvalue)}; slope = {m.params['age_weeks']:.2f} per week (95% CI {lo:.2f} to {hi:.2f}), "
          f"slope p {fmt(m.pvalues['age_weeks'])}, R^2 = {m.rsquared:.2f}")
m = smf.ols("density ~ age_weeks * genotype", d).fit()
term = "age_weeks:genotype[T.KO]"
lo, hi = m.conf_int().loc[term]
print(f"Difference in slopes (KO - WT) = {m.params[term]:.2f} per week (95% CI {lo:.2f} to {hi:.2f}), "
      f"t = {m.tvalues[term]:.2f}, df = {m.df_resid:.0f}, p {fmt(m.pvalues[term])}")
c = smf.ols("density ~ age_weeks + genotype", d).fit()
lo, hi = c.conf_int().loc["age_weeks"]
print(f"Common slope = {c.params['age_weeks']:.2f} per week (95% CI {lo:.2f} to {hi:.2f})")
for g in ["WT", "KO"]:
    s = d[d.genotype == g]
    print(f"{g} residual SD = {smf.ols('density ~ age_weeks', s).fit().resid.std(ddof=2):.2f}")
rob = smf.ols("density ~ age_weeks * genotype", d).fit(cov_type="HC3", use_t=True)
print(f"Robust (HC3) test of the difference in slopes: p {fmt(rob.pvalues[term])}")