###########################################################
# Joint models for longitudinal and time-to-event data ####
# Jana Vorlickova, 2019/2020                           ####
# Priloha k diplomove praci                            ####
###########################################################

###############################
# PART A - data simulation ####
###############################

library(MASS)
library(lme4)

rm(list=ls())
#set.seed(123) # for 400
#set.seed(345) # for 200
#set.seed(567) # for 100
set.seed(789) # for 600
K <- 10  # number of planned repeated measurements per subject, per outcome
t.max <- 84 # maximum follow-up time

# number of observations
n1 <- 120 #20,40,80,120
n2 <- 300 #50,100,200,300
n3 <- 180 #30,60,120, 180
v <- rep(c(1,2,3), c(n1, n2,n3))

# parameters for the longitudinal model
betas1 <- c("(Intercept)" = 120, "lek" = 0.02, "time"=-0.2)
betas2 <- c("(Intercept)" = 120, "lek" = 5, "time"=0.5)
betas3 <- c("(Intercept)" = 120, "lek" = -10, "time"=-0.5)
# measurement error standard deviation
sd <- 6
sd.int <-  0.2
sd.slope <- 0.25

# parameters for the survival model
# coefficients for Cox model
alpha1 <- c("lek" = 0.4)
alpha2 <- c("lek" = -0.18)
alpha3 <- c("lek" = -0.59)
weibA  <- 1.1   

DFlong <- list()
DFcox <- list()
M <- 100

for(i in 1:M){
  # group 1 ---####
  # -----------####
  
  times <- c(replicate(n1, c(0, sort(runif(K-1, 0, t.max))))) # at which time points longitudinal measurements are supposed to be taken
  lek <- sample(c(0,1), n1, replace = TRUE)
  DF <- data.frame(day = times, lek = factor(rep(lek, each = K)))
  X <- model.matrix(~ lek+day, data = DF)
  Z <- model.matrix(~ 1+day, data = DF)
  # design matrix for the survival model
  W <- cbind("(Intercept)" = 1, "lek" = lek)
  
  #simulate random effects
  b0 <- rnorm(n1, 0, sd.int)
  b1 <- rnorm(n1, 0, sd.slope)
  b <- cbind(b0, b1)
  
  # simulate longitudinal responses
  id <- rep(1:n1, each = K)
  eta.y <- as.vector(X %*% betas1 + rowSums(Z * b[id,])) # linear predictor
  y <-  eta.y + rnorm(n1 * K, 0,sd)
  
  
  # Longitudinal data set
  DFlon1 <- cbind(id, y, DF)
  
  # Cox model
  lek <- unique(DFlon1[, c("id", "lek")])
  Xalpha <- alpha1*(as.numeric(lek[,2])-1)
  U      <- runif(n1, 0, 1)      
  eventT <- (-log(U)*exp(5 -Xalpha))^(1/weibA)
  # censoring
  Ctimes <-pmin(runif(n1, 60, 90),84)
  Time <- pmin(eventT, Ctimes)
  event <- as.numeric(eventT <= Ctimes) # event indicator
  DFcox1 <- cbind(Time, event, lek, Ctimes)
  
  
  # Group 2 --####
  # ----------####
  times <- c(replicate(n2, c(0, sort(runif(K-1, 0, t.max))))) # at which time points longitudinal measurements are supposed to be taken
  lek <- sample(c(0,1), n2, replace = TRUE) 
  
  DF <- data.frame(day = times, lek = factor(rep(lek, each = K)))
  X <- model.matrix(~ lek+day, data = DF)
  Z <- model.matrix(~ 1+day, data = DF)
  # design matrix for the survival model
  W <- cbind("(Intercept)" = 1, "lek" = lek)
  
  #simulate random effects
  b0 <- rnorm(n2, 0, sd.int)
  b1 <- rnorm(n2, 0, sd.slope)
  b <- cbind(b0, b1)
  
  # simulate longitudinal responses
  id <- rep(1:n2, each = K)
  eta.y <- as.vector(X %*% betas2 + rowSums(Z * b[id,])) # linear predictor
  y <-  eta.y + rnorm(n2 * K, 0,sd)
  
  
  # Longitudinal data set
  id <- id+n1
  DFlon2 <- cbind(id, y, DF)
  
  # Cox model
  lek <- unique(DFlon2[, c("id", "lek")])
  Xalpha <- alpha2*(as.numeric(lek[,2])-1)
  U      <- runif(n2, 0, 1)      
  eventT <- (-log(U)*exp(5 -Xalpha))^(1/weibA)
  # censoring
  Ctimes <-pmin(runif(n2, 60, 90),84)
  Time <- pmin(eventT, Ctimes)
  event <- as.numeric(eventT <= Ctimes) # event indicator
  DFcox2 <- cbind(Time, event, lek, Ctimes)
  
  # Group 3 --####
  # ----------####
  
  times <- c(replicate(n3, c(0, sort(runif(K-1, 0, t.max))))) # at which time points longitudinal measurements are supposed to be taken
  lek <- sample(c(0,1), n3, replace = TRUE) 
  DF <- data.frame(day = times, lek = factor(rep(lek, each = K)))
  X <- model.matrix(~ lek+day, data = DF)
  Z <- model.matrix(~ 1+day, data = DF)
  # design matrix for the survival model
  W <- cbind("(Intercept)" = 1, "lek" = lek)
  
  #simulate random effects
  b0 <- rnorm(n3, 0, sd.int)
  b1 <- rnorm(n3, 0, sd.slope)
  b <- cbind(b0, b1)
  
  # simulate longitudinal responses
  id <- rep(1:n3, each = K)
  eta.y <- as.vector(X %*% betas3 + rowSums(Z * b[id,])) # linear predictor
  y <-  eta.y + rnorm(n3 * K, 0,sd)
  
  # Longitudinal data set
  id <- id+n1+n2
  DFlon3 <- cbind(id, y, DF)
  
  # Cox model
  lek <- unique(DFlon3[, c("id", "lek")])
  Xalpha <- alpha3*(as.numeric(lek[,2])-1)
  U      <- runif(n3, 0, 1)      
  eventT <- (-log(U)*exp(5 -Xalpha))^(1/weibA)
  # censoring
  Ctimes <- pmin(runif(n3, 60, 90),84)
  Time <- pmin(eventT, Ctimes)
  event <- as.numeric(eventT <= Ctimes) # event indicator
  DFcox3 <- cbind(Time, event, lek, Ctimes)
  
  dataLon <- rbind(DFlon1, DFlon2, DFlon3)
  dataCox <- rbind(DFcox1, DFcox2, DFcox3)
  DFcox[[i]] <- dataCox
  DFlong[[i]] <- dataLon
}

DFlong2 <- list()
for(j in 1:100){
  dataLon <- DFlong[[j]]
  tf <- dataLon$day[dataLon$id==1]<dataCox$Ctimes[dataCox$id==1]
  for(i in 2:200){
    a <- dataLon$day[dataLon$id==i]<dataCox$Ctimes[dataCox$id==i]
    tf <- cbind(tf, a)
  }
  as.vector(tf)
  dataLonU <-  dataLon[as.vector(tf),]
  DFlong2[[j]] <- dataLonU
}

#save(DFcox,file="XXXcox.Rdata")
#save(DFlong2,file="XXXlong.Rdata")
#save(DFlong,file="XXXlongb.Rdata")



##########################################
# PART B - model + estimation process ####
##########################################

### Packages ##############
###-----------------------

library(runjags)
library(coda)
library(dplyr)
rm(list=ls())

#load("DFcox3.Rdata")
#load("DFlong3.Rdata")
#DFlong <- DFlong2
psti <-list()
# MODEL ####

modelJags <- "
      model {
  for (i in 1:n) {
    for (j in offset[i]:(offset[i+1] - 1)) { # pocet pozorovani na jeden subjekt
       muy[j] <- inprod((equals(v[i], 1) * beta1[1:ncX]) + (equals(v[i], 2) * beta2[1:ncX]) + (equals(v[i], 3) * beta3[1:ncX]), X[j, 1:ncX]) + beta*inter[j] + b[i, 1]*Z[j, 1] + b[i,2]*Z[j, 2] # tohle je vzdy stejne, fixni cast class not specific a random effects
       y[j] ~ dnorm(muy[j], tau) 
    }
    yS[i] ~ dinterval(t[i], time[i])
    regres.cox[i] <- eta + (equals(v[i], 1) * alpha1) + (equals(v[i], 2) * alpha2) +  (equals(v[i], 3) * alpha3)*lek[i]
    lambdar[i]  <- exp(regres.cox[i])
    t[i] ~ dweib(nur,lambdar[i])
    v[i] ~ dcat(pr[]) 
    b[i, 1:nb] ~ dmnorm(mu[],Wish[,])
  }
  
    nur <- nu.prior
    eta<- eta.prior 
    lambda0r <- exp(eta)
  ### Prior distributions 
  # zatim nejjednodussi zpusob
  pr[1:G] ~ ddirch(prior.cl[1:G])
  
  # Longitudinal model - fixed effects
  beta ~  dnorm(0,0.001)
  beta1[1:2] ~ dmnorm(priorMean.betas1[], priorTau.betas1[, ]) # pro 3 bety class specific, pro 1 ne
  beta2[1:2] ~ dmnorm(priorMean.betas2[], priorTau.betas2[, ])
  beta3[1:2] ~ dmnorm(priorMean.betas3[], priorTau.betas3[, ])
  tau ~ dgamma(1.0, 0.005)
  sigma = 1/tau
  
  # Random effects prior
  mu[1:2] ~ dmnorm(priorMu.mu.b[], priorTau.mu.b[,])
  Wish ~ dwish(R, 2)
  
  # parameters of Weibull distribution
  nu.prior ~  dgamma(1.0, 1) 
  eta.prior ~ dnorm(0, 0.05) 
 
  alpha1 ~  dnorm(0,0.001) 
  alpha2 ~  dnorm(0,0.001)
  alpha3 ~  dnorm(0,0.001)
  
}"

simulace <- function (i, DFlong, DFCox, modelJags ){
  # Preparation ###
  dataLong <- DFlong[[i]]
  N <- nrow(dataLong)
  dataCox <- DFcox[[i]]
  n <- nrow(dataCox)
  names(dataCox)
  names(dataLong)
  X <-cbind(as.numeric(dataLong[, "lek"])-1, dataLong[, "day"])
  Z <- cbind(1, dataLong[, "day"])
  inter <- c(rep(1,N))
  ncX <- ncol(X)
  ID <- dataLong$id
  offset <-  as.vector(c(1, 1 + cumsum(tapply(ID, ID, length)))) 
  
  
  parameters.paq <- c("beta1", "mu", "Wish", "tau", "sigma",
                      "beta2","beta3", "beta",
                      "lambda0r","nur", "alpha1","alpha2", "alpha3",
                      "pr"
  ) 
  
  # Prior parameters ####
  ### Prior distirbutions - parameters
  G <- 3 # number of classes 
  nb <- 2 # number of random effects
  b <-  matrix(rep(NA, 2*n), nrow=n)
  for ( j in 1:nb) {
    b[,j] <- as.numeric(b[,j])}
  
  ### Longitudinal model
  priorMean.betas1 <- priorMean.betas2  <-priorMean.betas3<- rep(0,2) 
  priorTau.betas1 <- priorTau.betas2<- priorTau.betas3 <- diag(rep(0.001, 2))
  priorTau.mu.b <- diag(0.001,2)
  priorMu.mu.b <- c(0,0)
  R <- diag(0.001,2)
  
  ### Probability model
  prior.cl <- rep(1/G, G)
  
  # DATA ####
  dataJags <- list(t = as.numeric(rep(NA, n)), # vsechny casy do udalosti jsou chybejici hodnoty pokud bychom meli nejake skutecne, tak by se tu stridali ty hodnoty s NA
                   #v = as.numeric(rep(NA, n)),
                   b = b, 
                   y = dataLong[, "y"],
                   offset = offset, 
                   yS =  dataCox[, "event"],
                   time = dataCox[, "Time"],
                   ncX = ncX, 
                   G = G,
                   nb=nb,
                   n = n,
                   lek = dataCox[, "lek"],
                   Z = Z,
                   X = X, 
                   inter = inter,
                   prior.cl = prior.cl,
                   R=R, priorMu.mu.b=priorMu.mu.b, priorTau.mu.b=priorTau.mu.b,
                   priorMean.betas1=priorMean.betas1, priorMean.betas2=priorMean.betas2, priorMean.betas3=priorMean.betas3,
                   priorTau.betas1=priorTau.betas1,  priorTau.betas2=priorTau.betas2 , priorTau.betas3=priorTau.betas3
  )
  
  
  # Initial values ####
  init.t <- as.numeric(rep(NA, n))
  init.t <- dataCox[, "Time"] + runif(n, 0, 0.01)
  
  set.seed(300035)      ## this seed is used by 'rnorm'/'runif' calls below
  inits.data <- list( list(.RNG.seed = 20000901, .RNG.name = "lecuyer::RngStream", # kvuli paralelnimu - nastavujeme set.seed u kazdeho zvlast
                           t = init.t,# v=v.init,
                           beta =rnorm(1, 120 , 1), beta1 =c(rnorm(1, 0 , 1), rnorm(1, -0.2 , 1)),  beta2 =c(rnorm(1, -10 , 1), rnorm(1, 0.5 , 1)), 
                           beta3 =c(rnorm(1, -15 , 1), rnorm(1, -0.5 , 1)), tau = runif(1, 0.0001, 0.005),
                           mu =rnorm(2, 0 , 1), alpha1 = rnorm(1, 0, 1),alpha2 = rnorm(1, 0, 1),alpha3 = rnorm(1, 0, 1)),
                      list(.RNG.seed = 20010911, .RNG.name = "lecuyer::RngStream", 
                           t = init.t,# v=v.init,
                           beta =rnorm(1, 120 , 1), beta1 =c(rnorm(1, 0 , 1), rnorm(1, -0.2 , 1)),  beta2 =c(rnorm(1, -10 , 1), rnorm(1, 0.5 , 1)), 
                           beta3 =c(rnorm(1, -15 , 1), rnorm(1, -0.5 , 1)),tau = runif(1, 0.0001, 0.005),
                           mu =rnorm(2, 0 , 1), alpha1 = rnorm(1, 0, 1),alpha2 = rnorm(1, 0, 1),alpha3 = rnorm(1, 0, 1)), 
                      list(.RNG.seed = 20010921, .RNG.name = "lecuyer::RngStream", 
                           t = init.t,# v=v.init,
                           beta =rnorm(1, 120 , 1), beta1 =c(rnorm(1, 0 , 1), rnorm(1, -0.2 , 1)),  beta2 =c(rnorm(1, -10 , 1), rnorm(1, 0.5 , 1)), 
                           beta3 =c(rnorm(1, -15 , 1), rnorm(1, -0.5 , 1)), tau = runif(1, 0.0001, 0.005),
                           mu =rnorm(2, 0 , 1), alpha1 = rnorm(1, 0, 1),alpha2 = rnorm(1, 0, 1),alpha3 = rnorm(1, 0, 1)))
  
  
  ### Test implementation of JAGS
  ### -----------------------------
  # testjags(silent = FALSE)
  
  
  jagsModelsimul <- run.jags(model = modelJags, # cesta k file, pokud bych mela slozitejsi model, napsany jinde
                             monitor = parameters.paq, # to co chci ukladat
                             data = dataJags, inits = inits.data, adapt = 1000, #prvnich 1000 pozorovani se zahodi
                             n.chains = 3, n.sims = 3, thin = 2, #method = "rjparallel",  #n.sims na trech jadrech,
                             # thin = 1, znamena, aby si pamatoval kazdou hodnotu - jak snizit autokorelaci? 
                             # nepamatovat si kazdou, tak aby, ale kazdou nejakou. Ale zas tak to nefunguje.
                             burnin = 5000, sample = 25000)
  
  
  
  print(pst <- unlist(summary(jagsModelsimul, vars=c("pr","alpha", "beta", "sigma", "Wish", "lambda0r","nur"))[, 4]))
  
}

data <- list()
for (i in 1:100) {
  skip_to_next <- FALSE
  # Note that print(b) fails since b doesn't exist
  tryCatch({data[[i]] <- simulace(i, DFlong, DFCox, modelJags)}, error = function(e) { skip_to_next <<- TRUE})
  if(skip_to_next) { next }     
}


#save(data,file="XXX.Rdata")

#######################
# PART C - results ####
#######################

rm(list=ls())
# load data sets
#load("XXX.Rdata")
#load("XXX.Rdata")
DFlong <- DFlong2

table(DFcox[[1]][,2])
# Longitudinal data set
a <-c()
for(i in 1:100){
  a[i] <- dim(DFlong[[i]])[1]}

mean(a)

# frequency of events
event <- c()
for(i in 1:100){
  event[i] <- table(DFcox[[i]][,2])[2]}
mean(event)

# load data
#load("XXX.Rdata")

dataUn <- matrix(c(rep(NA, 2000)), nrow=100, ncol=20)
for( i in 1:100){
  dataUn[i,] <- unlist(dataNopar[[i]])
  
}

head(dataUn)
colnames(dataUn) <- c("pr[1]","pr[2]","pr[3]","alpha1","alpha2","alpha3","beta1[1]","beta1[2]","beta2[1]",
                      "beta2[2]","beta3[1]","beta3[2]","beta","sigma","Wish[1,1]","Wish[2,1]","Wish[1,2]","Wish[2,2]", "lambda0r", "nur" )
dataUn <- dataUn[complete.cases(dataUn),]
apply(dataUn, 2, summary)
dataUn <- as.data.frame(dataUn)

# class-specific parameters 
skupina1 <- dataUn[,c("beta1[1]", "beta1[2]", "alpha1", "pr[1]")]
skupina2 <-dataUn[,c("beta2[1]", "beta2[2]", "alpha2", "pr[2]")]
skupina3 <- dataUn[,c("beta3[1]", "beta3[2]", "alpha3","pr[3]")]
names(skupina1) <-names(skupina2) <- names(skupina3) <- var<- c("beta1", "beta2", "alpha", "pr")


round((t1mean <-apply(skupina1, 2, mean)),2)
round((t1sd <- apply(skupina1, 2, sd)),2)
for(i in 1:4){
  hist(skupina1[, i])}

round((t2mean <-apply(skupina2, 2, mean)),2)
round((t2sd <- apply(skupina2, 2, sd)),2)
for(i in 1:4){
  hist(skupina2[, i])}

round((t3mean <-apply(skupina3, 2, mean)),2)
round((t3sd <- apply(skupina3, 2, sd)), 2)
for(i in 1:4){
  hist(skupina3[, i])}

# NON-SPECIFIC parameters
# Longitudinal model - non-specific
sigma <- sqrt(dataUn$sigma)
hist(sigma)
(meanS <- mean(sigma))
(sdS <- sd(sigma))

beta <- (dataUn$beta)
hist(beta)
(meanB <- mean(beta))
(sdB <- sd(beta))

# Parameters of Weibull distribution
nu <- (dataUn$nur)
hist(nu)
(meanNu<- mean(nu))
(sdNu <- sd(nu))

eta <- dataUn$lambda0r
hist(eta)
(meanE <- mean(eta))
(sdE <- sd(eta))

# Random effects - covariance matrix

wish11 <- dataUn[,"Wish[1,1]"]
wish11in <- 1/wish11
sqrt(wish11in)
(meanW11 <- mean(sqrt(wish11in)))
(sdW11 <- sd(sqrt(wish11in)))
hist(sqrt(wish11in))

wish21 <- dataUn[,"Wish[2,1]"]
wish21in <- 1/wish21
(meanW21 <-mean(wish21in))
(sdW21 <-sd(wish21in))
hist(wish21in)

wish22 <- dataUn[,"Wish[2,2]"]
wish22in <- 1/wish22
hist(sqrt(wish22in))
(meanW22 <- mean(sqrt(wish22in)))
(sdW22 <- sd(sqrt(wish22in)))

round(parMean <- c(meanS, meanB, meanE, meanNu, meanW11, meanW21, meanW22),2)
(parSD <- round(c(sdS, sdB, sdE, sdNu, sdW11, sdW21, sdW22),3))