library(readstata13)
library(lme4)
data <- read.dta13("trial_data/chapter 6_1a.dta")
source('crt.R')

#Trial 1

id <- NULL
y <- NULL
s <- NULL
num_homes <- 35
num_steps <- 4
alpha <-0.10

data_homes <- split(data, f = data$home)
data_arms <- split(data, f = data$arm)

for (i in 1:num_homes){
  id <- c(id,rep(i,num_steps+1))
}

#MCRTs+Z

for (i in 1:num_homes){
  data_i <- data_homes[[i]]
  for (t in 1:(num_steps+1)){
    data_it <- subset(data_i,time==t)
    data_it <- subset(data_it,is.na(data_it$Qol)==FALSE)
    y <- c(y, mean(data_it$Qol))
    s<-c(s,unique(data_i$arm))
  }
}
t <- rep((0:num_steps),num_homes)
data_ <- data.frame(id, t, y, s)

for (i in 1:num_homes){
  data_[id==i,]$y <- data_[id==i,]$y - data_[id==i & t==0,]$y
}

mcrtsz <- matrix(0, 3,2)
for (l in 0:2){
  data_list <- data_swd(data_,l,sample_split=TRUE)
  p_list <- crt_swd(data_list, type='ci',b=10000,n_beta=10000,const=4,print=TRUE)
  ci <- global_ci(data_list,p_list,method='zscore',alpha=0.1,const=4)
  cat("\n")
  cat("lag:", l)
  cat("\n")
  cat("Confidence interval:", ci)
  cat("\n")
  mcrtsz[l+1,1:2] = ci
}

#Mixed-effects models with and without time effects

data$time_ <- data$time -1 

lags <- data$time_ - data$arm
num <- NROW(data)
z_matrix <-  matrix(0,nrow = num,ncol = num_steps)
for (j in 1:num){
  if (lags[j]>=0){
    z_matrix[j,lags[j]+1] <- 1
  }
}
z_data <- data.frame(z_matrix)
for(i in 0:(NCOL(z_data)-1)) {                   
  colnames(z_data)[i+1] <- paste0("z_", i)
}
data_xyz <- cbind(data,z_data)

model <- lmer(Qol ~  (1 | home) + z_0 + z_1 + z_2+z_3,  data=data_xyz, REML=FALSE)
result<- confint(model, method = "Wald",level = 1-alpha)
mix_1 <- result[4:6,1:2]

model <- lmer(Qol ~  (1 | home) + time_ + z_0 + z_1 + z_2+z_3,  data=data_xyz, REML=FALSE)
result<- confint(model, method = "Wald",level = 1-alpha)
mix_2 <- result[5:7,1:2]

all_ci_1 <- rbind(mix_1,mix_2,mcrtsz)
rownames(all_ci_1) <- NULL
save(all_ci_1, file= paste0("results/real_1.RData"))
