################################
### HW4 Q12-Q15: hw4_sleep.csv (18 subjects x 10 days, Reaction in ms)
### Q12: |SE(Age) OLS - SE(Age) LMM|, model Reaction ~ Age*Sex, random intercept per Subject.
### Q13: ICC (%) from that LMM.
### Q14: GEE exchangeable, Reaction ~ Days + Age*Sex; marginal mean, male, Days=0, Age=30.
### Q15: same with AR(1) working correlation.
### Needs: install.packages(c("lme4", "lmerTest", "geepack"))
### ! Rscript "/mnt/c/Users/kirch/OneDrive/Documents/dev/ut-11-hdi/docs/part2/hw4_q12_15.R"
################################

suppressMessages({library(lme4); library(lmerTest); library(geepack)})
setwd("/mnt/c/Users/kirch/OneDrive/Documents/dev/ut-11-hdi/docs/part2")
sleep = read.csv("hw4_sleep.csv")
sleep$Sex = factor(sleep$Sex)              # levels F (reference), M
sleep$Subject = factor(sleep$Subject)
sleep = sleep[order(sleep$Subject, sleep$Days), ]   # geepack needs rows ordered within cluster
cat("n rows:", nrow(sleep), " subjects:", nlevels(sleep$Subject), " Sex levels:", levels(sleep$Sex), "\n")

cat("\n==== Q12: SE of Age, OLS vs LMM (Reaction ~ Age*Sex) ====\n")
ols = lm(Reaction ~ Age*Sex, data = sleep)
lmm = lmer(Reaction ~ Age*Sex + (1 | Subject), data = sleep)
cat("OLS coefficients:\n"); print(round(summary(ols)$coefficients, 4))
cat("LMM fixed effects:\n"); print(round(summary(lmm)$coefficients, 4))
se.ols = summary(ols)$coefficients["Age", "Std. Error"]
se.lmm = summary(lmm)$coefficients["Age", "Std. Error"]
cat("SE(Age) OLS   :", se.ols, "\n")
cat("SE(Age) LMM   :", se.lmm, "\n")
cat("|difference|  :", abs(se.ols - se.lmm), "\n")
cat("Q12 ANSWER (2 dp):", round(abs(se.ols - se.lmm), 2), "\n")

cat("\n==== Q13: ICC from the LMM ====\n")
vc = as.data.frame(VarCorr(lmm))
s2b = vc$vcov[vc$grp == "Subject"]; s2e = vc$vcov[vc$grp == "Residual"]
cat("sigma_b^2 (Subject intercept):", s2b, "\n")
cat("sigma^2   (Residual)         :", s2e, "\n")
icc = s2b / (s2b + s2e)
cat("ICC = sigma_b^2/(sigma_b^2+sigma^2):", icc, "\n")
cat("Q13 ANSWER (%, 2 dp):", round(100*icc, 2), "\n")

newd = data.frame(Days = 0, Age = 30, Sex = factor("M", levels = levels(sleep$Sex)))
cat("\n==== Q14: GEE exchangeable, Reaction ~ Days + Age*Sex ====\n")
gee.ex = geeglm(Reaction ~ Days + Age*Sex, data = sleep, id = Subject, corstr = "exchangeable")
print(summary(gee.ex)$coefficients); cat("alpha:", summary(gee.ex)$corr$Estimate, "\n")
b = coef(gee.ex)
mu.ex = b["(Intercept)"] + b["Days"]*0 + b["Age"]*30 + b["SexM"] + b["Age:SexM"]*30
cat("by hand: b0 + 30*bAge + bSexM + 30*bAge:SexM =", mu.ex, "\n")
cat("predict():", predict(gee.ex, newdata = newd), "\n")
cat("Q14 ANSWER (2 dp):", round(mu.ex, 2), "\n")

cat("\n==== Q15: GEE AR(1), Reaction ~ Days + Age*Sex ====\n")
gee.ar = geeglm(Reaction ~ Days + Age*Sex, data = sleep, id = Subject, corstr = "ar1")
print(summary(gee.ar)$coefficients); cat("alpha:", summary(gee.ar)$corr$Estimate, "\n")
b = coef(gee.ar)
mu.ar = b["(Intercept)"] + b["Age"]*30 + b["SexM"] + b["Age:SexM"]*30
cat("by hand:", mu.ar, "   predict():", predict(gee.ar, newdata = newd), "\n")
cat("Q15 ANSWER (2 dp):", round(mu.ar, 2), "\n")