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

#Trial 2

id <- NULL
y <- NULL
s <- NULL
num_hosp <- 17
num_steps <- 5
alpha <-0.10

data_hosp <- split(data, f = data$hospital)
data_arms <- split(data, f = data$arm)

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

#MCRTs+Z

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

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

mcrtsz <- matrix(0,4,2)
for (l in 0:3){
  cat("\n")
  cat("lag:", l)
  data_list <- data_swd(data_,l,sample_split=TRUE)
  p_list <- crt_swd(data_list, type='ci',exact=TRUE,n_beta=10000,const=1,print=TRUE)
  ci <- global_ci(data_list,p_list,method='zscore',alpha=0.1,const=1)
  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(pain ~    (1 | hospital)  + z_0 + z_1 + z_2 +z_3+z_4,  data=data_xyz, REML=FALSE)
result<- confint(model, method = "Wald",level = 1-alpha)
mix_1 <- result[4:7,1:2]

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

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