# Edps 590BAY   ---- post on web-site
# C.J.Anderson
#
# Estimating random intercept model with 2 level 1 
#  predictors to 23 schools from the nels data set.
#
# Takes a long time (brms is considerly faster),
#  but it is still an useful example.
#
library(rstan)
library(lme4)
library(brms)

setwd("D:\\Dropbox\\edps 590BAY\\Lectures\\9 Ham_Stan_brms")

##################################
# Data and Data list             #
##################################

nels <- read.table("school23_data.txt",header=T)
nels <- nels[order(nels$school,nels$student) , ]

# Get information on number schools & create consecutive 
# integers that I will use an id numbers in model block in 
# stan model.

school <- unique(nels$school)# original school ids
N <- length(school)          # number schools  
school.int <- as.data.frame(cbind(school,seq(1:N)))  
names(school.int) <- c("school", "id")        
nels <- merge(nels,school.int,by="school")           

n <- length(nels$math)    # total number of students

###############################################
# construct x design matrix for fixed effects #
# and create data list for rstan              #
###############################################
one <- matrix(1,nrow=n,ncol=1)     # for the intercept
x <- cbind(one,nels$homew,nels$ses)# design matrix
K <- ncol(x)                       # number predictors
y <- nels$math                     # outcome variable
id <- nels$id  # id for schools within which student attend

#
# Putting data together into a data list
#
hlmList = list(n=n,N=N,K=K,y=y,x=x,id=id)

#############################################
# The model code for random intercept model #
#  with weakly informative priors           #
# The model is vectorized to some extend    #
#############################################

hlm.1inf <- "	
data {
  int n;           // number of students (total)
  int N;	       // number of schools
  int<lower=0> K;  // number of predictors
  matrix[n,K] x;   // matrix of predicotrs values 
                   //  including column of 1s for intercept
  vector[n] y;     // outcome vector
  int id[n];       // id that indicates what school a student
                   //   attended
}
parameters {
  vector[3] gamma;    // fixed-effects parameters for
                      //   the intercept & 2 predictors
  vector[N] u0j;      // random intercepts 
  real<lower=0> sigma;// level 1 standard deviation
  real<lower=0> tau0; // level 2 standard deviation
}
model {
                         // weakly informative
  tau0  ~ cauchy(0, 1);  // prior for sd between schools
                         //   which is a hyper-paramter
  sigma ~ cauchy(0, 1);  // prior for residual sd
                         // vectorized for speed
  gamma ~ normal(0, 20); // prior for fixed-effect intercept
  u0j   ~ normal(0,tau0);// prior for random intercepts

  //likelihood
  for (i in 1:n){
     y[i] ~ normal(x*gamma + u0j[id[i]], sigma);
  }
}
"

#################################################
# Running hlm random intercept:                 #   
# Laptop computer issue:                        #
# o This runs when have just 1 chain and 1 core #
# o Trying 2 chains and more iterations         #   
#                                               #
#################################################                             
hlm1 <- stan(
     model_code = hlm.1inf, 
     model_name = "Nels23: math ~ 1 + homework + ses+ (1|school)",
     data = hlmList, 
	 iter = 4000,
	 chains = 2,
	 cores = 1,
	 verbose = FALSE)
	   
###############################################
# Output (don't want to see the u0js)         # 
###############################################	   	   
print(hlm1, 
      pars = c("gamma","sigma","tau0"), 
	  probs = c(0.025, 0.975), 
	  digits = 4)

#######################################################
# Alternative estimation options that are much faster #
#######################################################
hlm1.brm <- brm(math ~ homew + ses + (1 | id), nels)


hlm1.lmer <- lmer(math ~ homew + ses + (1 | id), data=nels,REML=FALSE)