# hsb.R
#
# content: (1) Read and plot data
#          (2) Fit models with random school effects
#          (3) Hierarchical modeling
#
# input: data/hsbdataset.txt
# output: --
#
# last mod: 2025-06-20, NW

library(lme4)
library(lattice)

#----- (1) Read and plot data -------------------------------------------------
dat <- read.table("data/hsbdataset.txt", header = TRUE)

dat$gmmath <- mean(dat$mathach)
dat$meanmath <- with(dat, ave(mathach, school))

xyplot(mathach + meanmath + gmmath ~ cses | factor(school), data = dat,
       type = c("p", "r", "r"),
       distribute.type = TRUE,
       col = c("#91C86E", "#91C86E", "#78004B"))

# Shorter version
xyplot(gmmath + meanmath ~ cses | factor(school), data = dat, type = "r")

#----- (2) Fit models with random school effects ------------------------------

## Null model with school-specific random intercepts
m1 <- lmer(mathach ~ 1 + (1 | school), dat)

# Plot predictions
xyplot(mathach + predict(m1) + predict(m1, re.form = NA) ~ cses | factor(school),
       data = dat,
       type = c("p", "r", "r"),
       distribute.type = TRUE,
       col = c("#91C86E", "#91C86E", "#78004B"))

# ICC
VarCorr(m1)[[1]] / (VarCorr(m1)[[1]] + sigma(m1)^2)

## Model with socioeconomic status and school-specific random intercepts
xyplot(mathach ~ cses, dat)
mean(dat$cses)

m2 <- lmer(mathach ~ cses + (1 | school), dat)

xyplot(mathach + predict(m2) + predict(m2, re.form = NA) ~ cses | factor(school),
       dat, type = c("p", "r", "r"), distribute.type = TRUE,
       col = c("#91C86E", "#91C86E", "#78004B"))

# ICC
VarCorr(m2)[[1]] / (VarCorr(m2)[[1]] + sigma(m2)^2)

## Model with socioeconomic status and school-specific random slopes
m3 <- lmer(mathach ~ cses + (cses | school), dat)

xyplot(mathach + predict(m3) + predict(m3, re.form = NA) ~ cses | factor(school),
       dat, type = c("p", "r", "r"), distribute.type = TRUE,
       col = c("#91C86E", "#91C86E", "#78004B"))

## Model with socioeconomic status, sector, and school-specific random slopes
m4 <- lmer(mathach ~ cses + sector + (cses | school), data = dat)

## Model with socioeconomic status, sector, interaction, and school-specific
## random slopes
m5 <- lmer(mathach ~ cses * sector + (cses | school), data = dat)

xyplot(mathach ~ cses, data = dat, groups = sector, type = c("p", "r"))

#----- (3) Hierarchical modeling ----------------------------------------------

h1 <- lmer(mathach ~ meanses*cses + sector*cses + (1 + cses | school),
           data = dat, REML = FALSE)

h2 <- lmer(mathach ~ meanses*cses + sector*cses + (1 | school),
           data = dat, REML = FALSE)

# Likelihood-ratio test
anova(h2, h1)

pm1 <- profile(h1)
confint(pm1)
xyplot(pm1)
#densityplot(pm1)
splom(pm1, which = "theta_")


## Visualization of two way interaction of `cses` and `meanses`
c <- seq(-2, 2, length = 51)
m <- seq(-1, 1, length = 26)
ndat <- expand.grid(c, m)

colnames(ndat) <- c("cses", "meanses")

ndat$sector <- factor(0, levels = c("0", "1"))

z <- matrix(predict(lmm.2, newdata = ndat, re.form = NA), 51)

persp(c, m, z, theta = 40, phi = 20, col = "lightblue", ltheta = 60, shade = .9,
      xlab = "cses", ylab = "meanses", zlab = "mathach", main = "Model 2")

h3 <- lmer(mathach ~ meanses + sector*cses + (1 + cses | school),
           data = dat, REML = FALSE)

z <- matrix(predict(lmm.3, newdata = ndat, re.form = NA), nrow = 51)

persp(c, m, z, theta = 40, phi = 20, col = "lightblue", ltheta = 60, shade = .9,
      xlab = "cses", ylab = "meanses", zlab = "mathach", main = "Model 3")