library(lme4)
library(lattice)

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

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

plot(dat$ses - dat$meanses, dat$cses)


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

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


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

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

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

sjPlot::tab_model(m1)


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)

sjPlot::tab_model(m1, m2)


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"))

sjPlot::tab_model(m1, m2, m3)


m4 <- lmer(mathach ~ cses + sector + (cses | school), data = dat)

sjPlot::tab_model(m1, m2, m3, m4)

m5 <- lmer(mathach ~ cses * sector + (cses | school), data = dat)

sjPlot::tab_model(m1, m2, m3, m4, m5)

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



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

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

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")

lmm.3 <- 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")


# TODO: Add profiling to show instability of parameter estimation?