library(devtools)
library(rjags)
library(ggplot2)
library(dplyr)

ur <- "https://raw.githubusercontent.com/PerceptionCognitionLab/data0/master/contexteffects/FlankerStroopSimon/cleaning.R"
source_url(ur)




ICC_lsm <- "model{

for(j in 1:J){

# latent betas
beta_raw_l[j] ~  dnorm(0, 1)

# random effect
beta_l[j] <- fe_mu + tau_mu * beta_raw_l[j]

# cholesky
z2[j] ~ dnorm(0, 1)
beta_raw_s[j] = rho12 * beta_raw_l[j] + sqrt(1 - rho12^2) * z2[j]

beta_s[j] <- fe_sd + tau_sd * beta_raw_s[j]

}

for(i in 1:N){

# likelihood
y[i] ~ dnorm(beta_l[ID[i]], 1/exp(beta_s[ID[i]])^2)

}


# fixed effects priors
fe_mu ~ dnorm(0, 1)
fe_sd ~ dnorm(0, 1)

# random effects priors
tau_mu ~ dt(0, pow(prior_scale,-2), 10)T(0,)
tau_sd ~ dt(0, pow(prior_scale,-2), 10)T(0,)



# prior for RE correlation
fz ~ dnorm(0, 1)
rho12 = tanh(fz)

}"



incongruent <- subset(stroop, congruency == "congruent")


model <- jags.model(textConnection(ICC_lsm),
                    inits = list(fe_mu = mean(incongruent$rt)),
                    data = list(y = as.numeric((incongruent$rt)),
                                ID = incongruent$ID,
                                J = length(unique(incongruent$ID)),
                                N = length(incongruent$rt),
                                prior_scale = 1))

# fit for motivating example
fit_1 <- coda.samples(model,
                      n.iter = 5000,
                      variable.names=c("fe_mu",
                                       "fe_sd",
                                       "tau_mu",
                                       "tau_sd",
                                       "beta_s",
                                       "beta_l"))

# samples
fit_1_samples <- do.call(rbind.data.frame, fit_1)

# mean ICC
mu_icc <- mean(fit_1_samples$tau_mu^2 / (fit_1_samples$tau_mu^2 + exp(fit_1_samples$fe_sd)^2))

# colnames
col_names <- colnames(fit_1_samples)

# person specific SD
beta_s <- exp(fit_1_samples[, grep("beta_s", col_names)])

# person specific ICCs
posterior_icc <- apply(beta_s, MARGIN = 2, FUN = function(x){ fit_1_samples$tau_mu^2  /  (fit_1_samples$tau_mu^2  + x^2)}  )

# 90 % intervals for ICCs
credible_intervals <- t(apply(posterior_icc, 2, FUN = function(x)  {quantile(x, c(0.05, 0.95)) }))

posterior_means <- (colMeans(posterior_icc))


C <- cbind(posterior_means, credible_intervals) %>%
  as.data.frame() %>%
  arrange(posterior_means) %>%
  mutate(index = as.factor(1:121),
         sig = as.factor(ifelse(.$`5%` < mu_icc & .$`95%` < mu_icc, 1,
                                ifelse(.$`5%` > mu_icc & .$`95%` > mu_icc, 2, 3)))) %>%
  ggplot(aes(x = index,y  = posterior_means)) +
  geom_hline(yintercept = mu_icc, alpha = 0.5, linetype = "dotted") +
  geom_errorbar(aes(ymin = `5%`,
                    ymax = `95%`, color = sig)) +
  geom_point(color = "black") +
  scale_color_manual(values = c("#56B4E9", "#009E73", "grey70")) +
  theme_bw(base_family = "Times") +
  theme(panel.grid = element_blank(),
        panel.grid.major.x = element_blank(),
        panel.grid.minor.y = element_blank(),
        axis.text.x = element_blank(),
        axis.title = element_text(size = 14),
        title = element_text(size = 14),
        strip.background = element_rect(fill = "grey94"),
        strip.text = element_text(size = 14), legend.position = "none") +
  xlab("") +
  scale_x_discrete(expand = c(0.015, 0.015)) +
  scale_y_continuous(limits = c(.05, .6),
                     labels = seq(.1, .6, .1) ,
                     breaks = seq(.1, .6, .1)) +
  xlab("Ascending Index") +
  ylab("Intraclass Correlation Coefficient(s)")
C