library(lme4)
library(lattice)

dat <- read.table("data/hsbdataset.txt", header = TRUE)

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

# Plot with subsample of schools
set.seed(1427)
subdat <- subset(dat, dat$school %in% sample(unique(dat$school), size = 9))

pdf("figures/hsb_model1.pdf", height = 5, width = 5)
xyplot(mathach + predict(m1, newdata = subdat) +
       predict(m1, re.form = NA, newdata = subdat) ~ cses | factor(school),
       subdat,
       pch = 16,
       cex = 0.4,
       type = c("p", "r", "r"),
       lwd = 2,
       xlab = "socioeconomic status",
       ylab = "math performance",
       distribute.type = TRUE,
       par.strip.text = list(cex = 0.8),
       col = c("gray80", "#91C86E", "#78004B"))
dev.off()

# Model 2 - random intercept model
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"))


pdf("figures/hsb_model2.pdf", height = 5, width = 5)
xyplot(mathach + predict(m2, newdata = subdat) +
       predict(m2, re.form = NA, newdata = subdat) ~ cses | factor(school),
       subdat,
       pch = 16,
       cex = 0.4,
       type = c("p", "r", "r"),
       lwd = 2,
       xlab = "socioeconomic status",
       ylab = "math performance",
       distribute.type = TRUE,
       par.strip.text = list(cex = 0.8),
       col = c("gray80", "#91C86E", "#78004B"))
dev.off()

# Model 3 - random slope model
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"))


pdf("figures/hsb_model3.pdf", height = 5, width = 5)
xyplot(mathach + predict(m3, newdata = subdat) +
       predict(m3, re.form = NA, newdata = subdat) ~ cses | factor(school),
       subdat,
       pch = 16,
       cex = 0.4,
       type = c("p", "r", "r"),
       lwd = 2,
       xlab = "socioeconomic status",
       ylab = "math performance",
       distribute.type = TRUE,
       par.strip.text = list(cex = 0.8),
       col = c("gray80", "#91C86E", "#78004B"))
dev.off()