# Follow instructions at link below to install Rtools
# https://github.com/stan-dev/rstan/wiki/Installing-RStan-on-Windows#toolchain

# Install brms package and all required packages
install.package("brms")

# Load brms package and all required packages
library(brms)

# Read autism data, and examine
autism <- read.csv("X:\\autism.csv", h=T)

dim(autism)

autism[1:5,]

# Remove NAs

autism <- autism[!is.na(autism$vsae),]
dim(autism)

# lme() approach from Chapter 6 (see prep syntax on book web page)

model6.1.fit <- lme(vsae ~ age.2 + I(age.2^2) + sicdegp2.f +
+ age.2:sicdegp2.f + I(age.2^2):sicdegp2.f, 
+ random = ~ age.2 + I(age.2^2), method="REML", 
+ data = autism.grouped)

# Note lack of convergence
Error in lme.formula(vsae ~ age.2 + I(age.2^2) + sicdegp2.f + age.2:sicdegp2.f +  : 
  nlminb problem, convergence error code = 1
  message = iteration limit reached without convergence (10)

# lmer() approach from Chapter 6

model6.1.fit <- lmer(vsae ~ age.2 + I(age.2^2) + sicdegp2.f +
+ age.2*sicdegp2.f + I(age.2^2)*sicdegp2.f + (age.2 + I(age.2^2) | childid), 
+ REML = T, data = autism.updated)

# NOTE: NO WARNING MESSAGE!!

summary(model6.1.fit)

Linear mixed model fit by REML ['lmerMod']
Formula: vsae ~ age.2 + I(age.2^2) + sicdegp2.f + age.2 * sicdegp2.f +  
    I(age.2^2) * sicdegp2.f + (age.2 + I(age.2^2) | childid)
   Data: autism.updated

REML criterion at convergence: 4608

Scaled residuals: 
    Min      1Q  Median      3Q     Max 
-4.3015 -0.3905 -0.0392  0.2837  6.8667 

Random effects:
 Groups   Name        Variance Std.Dev. Corr       
 childid  (Intercept)  1.4332  1.1972              
          age.2       14.2647  3.7769    0.06      
          I(age.2^2)   0.1562  0.3952    0.89 -0.40
 Residual             37.6414  6.1353              
Number of obs: 610, groups:  childid, 158

Fixed effects:
                       Estimate Std. Error t value
(Intercept)            13.78643    0.82251  16.762
age.2                   5.59928    0.78666   7.118
I(age.2^2)              0.20417    0.08537   2.392
sicdegp2.f1            -5.43483    1.11108  -4.891
sicdegp2.f2            -4.03533    1.04597  -3.858
age.2:sicdegp2.f1      -3.28366    1.08196  -3.035
age.2:sicdegp2.f2      -2.75204    1.01795  -2.704
I(age.2^2):sicdegp2.f1 -0.13881    0.11833  -1.173
I(age.2^2):sicdegp2.f2 -0.13396    0.11023  -1.215

# Fit example normal model from Chapter 6 using Bayesian approach, with specified priors for fixed effects, standard deviations of random effects, and correlation of random effects

fit1 <- brm(vsae ~ age.2 + I(age.2^2) + sicdegp2.f + age.2 * sicdegp2.f + I(age.2^2) * sicdegp2.f + (age.2 + I(age.2^2) | childid), data = autism.updated, family = gaussian("identity"), prior = c(set_prior("normal(0,10)", class = "b"), set_prior("cauchy(0,2)", class = "sd"), set_prior("lkj(2)", class = "cor")), warmup = 500, iter = 2000, chains = 2)

# Analyze the results based on draws from the relevant posterior distributions
summary(fit1, waic=T)

 Family: gaussian 
  Links: mu = identity; sigma = identity 
Formula: vsae ~ age.2 + I(age.2^2) + sicdegp2.f + age.2 * sicdegp2.f + I(age.2^2) * sicdegp2.f + (age.2 + I(age.2^2) | childid) 
   Data: autism.updated (Number of observations: 610) 
Samples: 2 chains, each with iter = 2000; warmup = 500; thin = 1; 
         total post-warmup samples = 3000
    ICs: LOO = NA; WAIC = 4219.42; R2 = NA
 
Group-Level Effects: 
~childid (Number of levels: 158) 
                        Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
sd(Intercept)               0.90      0.52     0.05     1.95        172 1.01
sd(age.2)                   3.75      0.36     3.07     4.49        835 1.00
sd(Iage.2E2)                0.38      0.04     0.30     0.45        810 1.00
cor(Intercept,age.2)        0.05      0.30    -0.49     0.65        101 1.01
cor(Intercept,Iage.2E2)     0.44      0.35    -0.46     0.88         35 1.03
cor(age.2,Iage.2E2)        -0.31      0.12    -0.52    -0.07        744 1.00

Population-Level Effects: 
                     Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
Intercept               13.74      0.83    12.12    15.41       2025 1.00
age.2                    5.59      0.77     4.11     7.04        952 1.00
Iage.2E2                 0.20      0.08     0.05     0.36        952 1.00
sicdegp2.f1             -5.36      1.11    -7.53    -3.23       3000 1.00
sicdegp2.f2             -3.99      1.09    -6.19    -1.91       3000 1.00
age.2:sicdegp2.f1       -3.31      1.07    -5.38    -1.24       1067 1.00
age.2:sicdegp2.f2       -2.71      1.02    -4.70    -0.67       1131 1.00
Iage.2E2:sicdegp2.f1    -0.13      0.12    -0.36     0.10        977 1.00
Iage.2E2:sicdegp2.f2    -0.13      0.10    -0.34     0.07       1022 1.00

Family Specific Parameters: 
      Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
sigma     6.21      0.25     5.76     6.74       1140 1.00

Samples were drawn using sampling(NUTS). For each parameter, Eff.Sample 
is a crude measure of effective sample size, and Rhat is the potential 
scale reduction factor on split chains (at convergence, Rhat = 1).

# The results are quite similar to those from the frequentist approach, without any convergence issues!


# Plot the chains and simulated draws from the posterior distributions
plot(fit1)

# Plot marginal effects (need to create dummies first)

autism$gp2[autism$sicdegp == 2] <- 1
autism$gp2[autism$sicdegp != 2] <- 0

autism$gp3[autism$sicdegp == 3] <- 1
autism$gp3[autism$sicdegp != 3] <- 0

fit1 <- brm(vsae ~ age + gp2 + gp3 + (1 + age | childid), data = autism, family = gaussian("identity"), prior = c(set_prior("normal(0,10)", class = "b"), set_prior("cauchy(0,2)", class = "sd"), set_prior("lkj(2)", class = "cor")), warmup = 500, iter = 2000, chains = 2)

summary(fit1, waic = T)

plot(marginal_effects(fit1), points = TRUE, ask = TRUE)

# Fit model removing the correlation between the random intercepts and random slopes

fit2 <- brm(formula = vsae ~ age + factor(sicdegp) + (1 + age || childid), data = autism, family = gaussian("identity"), prior = c(set_prior("normal(0,10)", class = "b"), set_prior("cauchy(0,2)", class = "sd")), warmup = 500, iter = 2000, chains = 2)

summary(fit2, waic = T)

# Compare fits of two models

LOO(fit1, fit2)

# Fit a model with heterogeneous variance of conditional residuals across groups

fit3 <- brm(bf(vsae ~ age + factor(sicdegp) + (1 | childid), sigma ~ factor(sicdegp)), data = autism, family = gaussian("identity"), prior = c(set_prior("normal(0,10)", class = "b"), set_prior("cauchy(0,2)", class = "sd")), warmup = 500, iter = 2000, chains = 2)

summary(fit3, waic = T)

# LATEST GITHUB VERSION (3/5/18)

# Fit a model with heterogeneous variance components for random effects across groups

fit <- brm(count ~ Trt + (1|gr(patient, by = Trt)), data = epilepsy)
summary(fit)