#-----------------------------------------------------------------------------
# packages
# install.packages("nlme")
library(lme4)
# install.packages("mlmRev")
library(mlmRev)
# install.packages("sjPlot")
library(sjPlot)
# install.packages("msm")
library(msm)
# install.packages("ggplot2")
library(ggplot2)

#-----------------------------------------------------------------------------
# data
data(egsingle)
head(egsingle)
summary(egsingle)

#-----------------------------------------------------------------------------
# Data for two years only
df0 <- egsingle[egsingle$year == -0.5, c("childid", "schoolid", "math")]
df1 <- egsingle[egsingle$year == 0.5, c("childid", "schoolid", "math")]
df <- merge(df0, df1, by = c("childid", "schoolid"))
colnames(df) <- c("childid", "schoolid", "math0", "math1")
df$childid <- as.factor(paste(df$childid))
df$schoolid <- as.factor(paste(df$schoolid))

head(df)
summary(df)

n <- as.vector(table(df$schoolid))

# VAM model with no effect of school (Linear regression VAM)
fit1 <- lm(math1 ~ math0, data = df)
VAM1 <- aggregate(residuals(fit1) ~ df$schoolid, FUN = mean)[, 2]

resvar1<- hatvalues(fit1)
sigma1 <- sigma(fit1)

var1 <- aggregate(sigma1^2*(1 - resvar1) ~ df$schoolid, FUN = sum)[, 2]/n^2
sd1 <- sqrt(var1)

VAM1_low <- VAM1 - 1.96*sd1/sqrt(n)
VAM1_upp <- VAM1 + 1.96*sd1/sqrt(n)

VAM1

# VAM model with fixed effect of school
fit2 <- lm(math1 ~ 0 + math0 + schoolid, data = df)
VAM2 <- coef(fit2)[-1] - mean(coef(fit2)[-1])

estmean <- coef(fit2)[-1]
estvar <- vcov(fit2)[-1, -1]

sd2 <- c()
for (i in 1:59){
  sd2 <- c(sd2, 
           deltamethod(as.formula(paste0("~ ", paste0("x", i), " + (", paste(paste0("x", 1:59), collapse = " + "), ")/59")),
                       estmean, estvar))
}

var2 <- sd2^2
VAM2_low <- VAM2 - 1.96*sd2/sqrt(n)
VAM2_upp <- VAM2 + 1.96*sd2/sqrt(n)

VAM2

# VAM model with random effect of school
fit3 <- lmer(math1 ~ math0 + (1|schoolid), data = df)
VAM3 <- unlist(ranef(fit3)$schoolid)

var3 <- as.vector(VarCorr(fit3)$schoolid)
sd3 <- sqrt(var3)
VAM3_low <- VAM3 - 1.96*sd3/sqrt(n)
VAM3_upp <- VAM3 + 1.96*sd3/sqrt(n)

VAM3

cor(VAM1, VAM2)
cor(VAM1, VAM3)
cor(VAM2, VAM3)

# PLOTS
dat <- data.frame(ID = levels(df$schoolid),
                  VAM1 = VAM1,
                  VAM1_low = VAM1_low,
                  VAM1_upp = VAM1_upp,
                  VAM2 = VAM2,
                  VAM2_low = VAM2_low,
                  VAM2_upp = VAM2_upp,
                  VAM3 = VAM3,
                  VAM3_low = VAM3_low,
                  VAM3_upp = VAM3_upp)

(g1 <- ggplot(dat, 
              aes(y = VAM1, x = ID, ymin = VAM1_low, ymax = VAM1_upp)) +
    geom_point() + 
    geom_errorbar(width = 0) + 
    coord_flip() + 
    xlab("") + 
    ylab("") + 
    ggtitle("Model 1") +
    geom_hline(yintercept = 0) +
    theme_bw()) 

(g2 <- ggplot(dat, 
              aes(y = VAM2, x = ID, ymin = VAM2_low, ymax = VAM2_upp)) +
    geom_point() + 
    geom_errorbar(width = 0) + 
    coord_flip() + 
    xlab("") + 
    ylab("") + 
    ggtitle("Model 2") +
    geom_hline(yintercept = 0) +
    theme_bw()) 

(g3 <- ggplot(dat, 
              aes(y = VAM3, x = ID, ymin = VAM3_low, ymax = VAM3_upp)) +
    geom_point() + 
    geom_errorbar(width = 0) + 
    coord_flip() + 
    xlab("") + 
    ylab("") + 
    ggtitle("Model 3") +
    geom_hline(yintercept = 0) +
    theme_bw()) 

#-----------------------------------------------------------------------------
# Models from McCaffrey et al. (2004)
#-----------------------------------------------------------------------------
# Examining intercepts and slopes
set.seed(1)
samp <- sample(levels(egsingle$childid), 50)
egsingle.sub <- subset(egsingle, childid %in% samp)

level2 <- lmer(math ~ year + (year|childid), data = egsingle.sub)
level3 <- lmer(math ~ year + (year|schoolid), data = egsingle)

sjp.lmer(level2)
sjp.lmer(level3)
#-----------------------------------------------------------------------------
# Fitting unconditional VAM
(fit_uncon <- lmer(math ~ year + (year|schoolid/childid), 
            data = egsingle))
summary(fit_uncon)
ranef(fit_uncon)$schoolid
#-----------------------------------------------------------------------------
# Adding covariates as fixed effects
fit_covar1 <- lmer(math ~ year*size + female + (year|schoolid/childid), 
                   data = egsingle)
summary(fit_covar1)
VarCorr(fit_covar1)
# Uncorrelated random slopes and intercept
fit_covar2 <- lmer(math ~ year*size + female + (year||schoolid/childid), 
                  data = egsingle)

anova(fit_covar1, fit_covar2)
anova(fit_covar1, fit_uncon)