# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 
#               R code for running Empirical Application 2 in paper                   #
#   "Generalized Latent Variable Models for Location, Scale, and Shape parameters"    #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 

rm(list = ls())
set.seed(1234)

# Set working directory
# Save file "glvmlss_0.0.1.tar.gz" in working directory
# Uncomment below for installing the glvmlss package
# if(!"glvmlss" %in% installed.packages()) install.packages("glvmlss_0.0.1.tar.gz", type = "source", repos = NULL)

library(glvmlss)
data <- PISA2018

# Uncomment to load estimation results and avoid running the code
# Save file "Blind Submission - Auxiliary files/bsub_Results_PISA2018.RData" in working directory
# load("bsub_Results_PISA2018.RData")

# Summary statistics (Table A2 in Online Supplementary Material)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
vtable::st(data[,1:9], summ = c('notNA(x)','mean(x, na.rm = T)','sd(x, na.rm = T)','PerformanceAnalytics::skewness(x, na.rm = T)','PerformanceAnalytics::kurtosis(x, na.rm = T,method = "excess")'),
           summ.names = c('N','Mean','Std. Dev.','Skewness','Kurtosis'))
round(colMeans(as.matrix(data[,10:18])),2)

# Estimation
# ~~~~~~~~~~
p = ncol(data)
famSN = c(rep("sn",p/2), rep("binomial",p/2))
famNO = c(rep("normal",p/2), rep("binomial",p/2))

irR0 <- matrix(NA,2,2); irR0[2,2] <- 1
irR1 <- matrix(NA,2,2); diag(irR1) <- 1

irNO <- array(NA,dim = c(p,3,2))
irNO[1:9,3,1:2] <- 0; irNO[10:18,2,1] <- 0
irNO[10:18,,2] <- 0

# Model 1 (van der Linden, 2007)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
irmod1 <- irNO
irmod1[1:9,2,2] <- 0; irmod1[1:9,2,1] <- -1
mod1 <- glvmlss::glvmlss(data = data, family = famNO, mu.eq = ~ Z1+Z2, sg.eq = ~ 1,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod1, restr.R = irR0, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)

# Model 2 (Molenaar et al, 2015)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
irmod2 <- irNO
irmod2[1:9,2,2] <- 0;
mod2 <- glvmlss::glvmlss(data = data, family = famNO, mu.eq = ~ Z1+Z2, sg.eq = ~ 1,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod2, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)

# Model 3
# ~~~~~~~
irmod3 <- irNO
mod3 <- glvmlss::glvmlss(data = data, family = famNO, mu.eq = ~ Z1+Z2, sg.eq = ~ Z1+Z2,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod3, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)

irSN <- array(NA,dim = c(p,3,3))
irSN[1:9,3,1:3] <- 0; irSN[10:18,2,1] <- 0
irSN[10:18,,2:3] <- 0

# Model 4
# ~~~~~~~
irmod4 <- irSN
irmod4[1:9,2,2:3] <- 0
mod4 <- glvmlss::glvmlss(data = data, family = famSN, mu.eq = ~ Z1+Z2, sg.eq = ~ 1, nu.eq = ~ 1,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod4, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)
# Model 5
# ~~~~~~~
irmod5 <- irSN
irmod5[1:9,2,3] <- 0
mod5 <- glvmlss::glvmlss(data = data, family = famSN, mu.eq = ~ Z1+Z2, sg.eq = ~ Z1+Z2, nu.eq = ~ 1,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod5, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)
# Model 6
# ~~~~~~~
irmod6 <- irSN
irmod6[1:9,2,2] <- 0
mod6 <- glvmlss::glvmlss(data = data, family = famSN, mu.eq = ~ Z1+Z2, sg.eq = ~ 1, nu.eq = ~ Z1+Z2,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod6, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)
# Model 7
# ~~~~~~~
irmod7 <- irSN
mod7 <- glvmlss::glvmlss(data = data, family = famSN, mu.eq = ~ Z1+Z2, sg.eq = ~ Z1+Z2, nu.eq = ~ Z1+Z2,
                         EM.useGD = T, LR = 0.0001,
                         EM.maxit = 750, DM.maxit = 300, nQP = 45, verbose = T,
                         restr.b = irmod7, restr.R = irR1, sign.restr.b = c(-1,2),
                         corLV = T, tolerance = 1e-5)

# # Uncomment to save estimation results.
# save.image("bsub_Results_PISA2018.RData")

# Table 3 (Section 3.2)
# ~~~~~~~~~~~~~~~~~~~~~
xtable::xtable(
  rbind(c(mod1$AIC, mod1$BIC, mod1$K),
        c(mod2$AIC, mod2$BIC, mod2$K),
        c(mod3$AIC, mod3$BIC, mod3$K),
        c(mod4$AIC, mod4$BIC, mod4$K),
        c(mod5$AIC, mod5$BIC, mod5$K),
        c(mod6$AIC, mod6$BIC, mod6$K),
        c(mod7$AIC, mod7$BIC, mod7$K)))

# Table 4 (Section 3.2)
# ~~~~~~~~~~~~~~~~~~~~~
printT <- cbind(mod7$b[10:18,1,1], mod7$seb[10:18,1,1],
                mod7$b[10:18,3,1], mod7$seb[10:18,3,1],
                mod7$b[1:9,1,1], mod7$seb[1:9,1,1],
                mod7$b[1:9,2,1], mod7$seb[1:9,2,1],
                mod7$b[1:9,1,2], mod7$seb[1:9,1,2],
                mod7$b[1:9,2,2], mod7$seb[1:9,2,2],
                mod7$b[1:9,1,3], mod7$seb[1:9,1,3],
                mod7$b[1:9,2,3], mod7$seb[1:9,2,3])
rownames(printT) <- NULL
xtable::xtable(printT, digits = 2)

# Figures 4 and 5 (Section 3.2)
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Save file "Blind Submission - Auxiliary files/plotAuxE2.R" in working directory
source("plotAuxE2.R")

# Change item as necessary
gg_Prob(item = 8, mod = mod7, data = PISA2018)
gg_Dist(item = 8, mod = mod7, data = PISA2018) # Save 450 x 400
