import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from scipy import stats
d = pd.read_csv("hearts.csv")
d["group"] = pd.Categorical(d.group, ["control", "treated"])
d["ratio"] = d.heart_mg / d.body_weight_g            # mg per g
fmt = lambda p: "< 0.001" if p < 0.001 else f"= {p:.3f}"
for col, label in [("body_weight_g", "Body weight (g)"), ("heart_mg", "Heart weight (mg)"), ("ratio", "Heart/body weight (mg/g)")]:
    t, c = d.loc[d.group == "treated", col], d.loc[d.group == "control", col]
    r = stats.ttest_ind(t, c, equal_var=False)
    print(f"{label}: control {c.mean():.2f}, treated {t.mean():.2f}, Welch t = {r.statistic:.2f}, df = {r.df:.1f}, p {fmt(r.pvalue)}")
# ANCOVA on the log scale: log heart ~ group + log body weight (allometric, slope estimated from the data)
d["log_heart"], d["log_bw"] = np.log(d.heart_mg), np.log(d.body_weight_g)
slopes = smf.ols("log_heart ~ group * log_bw", d).fit()
print(f"Slopes differ between groups? interaction p {fmt(slopes.pvalues['group[T.treated]:log_bw'])}")
m = smf.ols("log_heart ~ group + log_bw", d).fit()
b, (lo, hi) = m.params["group[T.treated]"], m.conf_int().loc["group[T.treated]"]
print(f"ANCOVA (log scale): common slope = {m.params['log_bw']:.2f} (95% CI {m.conf_int().loc['log_bw', 0]:.2f} to {m.conf_int().loc['log_bw', 1]:.2f})")
print(f"Treated vs control at the same body weight: ratio of heart weights {np.exp(b):.3f} (95% CI {np.exp(lo):.3f} to {np.exp(hi):.3f}), "
      f"t = {m.tvalues['group[T.treated]']:.2f}, df = {m.df_resid:.0f}, p {fmt(m.pvalues['group[T.treated]'])}")
