rm(list=ls())
###

# load package
library(splines)
library(rstan)
library(loo)
library(mvtnorm)
library(ROCR)

options(mc.cores = parallel::detectCores())

#####00. Load dataset and prepare data######################
load("simdata.RData") 

#####00. Model formular#################################
#longitudinal covariates
longiForm = ~ age + X2

#longiForm_nvisit 
longiForm_nvisit =update(longiForm, ~. + numVisityr)

#longiForm_nvisitAge 
longiForm_nvisitAge =update(longiForm, ~. + numVisityr+numVisityr*age)

#Weibull covariates
survForm = ~ prev_Age + X2

#####01.1 (random intercept) longitudinal ################
X_longi <- as.matrix(model.matrix(longiForm ,
                                  simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi)

 
model_1int <- stan(file = "0_longit_int.stan",
                     data = stan_long_data,
                     iter=4000,warmup=3000, thin=5,chains=3,
                     refresh=0,
                     save_warmup = F)

print(long_0, pars=c("alpha0", "alpha_y", "sigma_y_u", "sigma_y_e", "L_y_u", "lp__"),
                         probs = c(0.025, 0.975)) 

#####01.2 (random intercept) Longidutindal with visit (random intercept)############################
X_longi <- as.matrix(model.matrix(longiForm_nvisitAge,
                                  simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi)

model_1int_nvist <- stan(file = "0_longit_int.stan",
                           data = stan_long_data,
                           iter=4000,warmup=3000, thin=5,chains=3,
                           refresh=0,
                           save_warmup = F)

#####01.3 (random intercept) joint model (random intercept)####################################
#recurrent event covariates
X_event = as.matrix(model.matrix(survForm, simdata)[, -1])

#longitudinal covaraites
X_longi <- as.matrix(model.matrix(longiForm,
                                  simdata)[,-1])

#stan dataset
stan_jm0_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                       y = simdata$y, 
                       id = simdata$id,
                       M_y = dim(X_longi)[2], X_longi= X_longi,
                       gap_time=simdata$gaptime,
                       M_x_event=dim(X_event)[2], X_event= X_event)

##
model_1int_sjm <- stan(file = "0_JM_int_HR.stan",
                         data = stan_jm0_data,
                         iter=4000,warmup=3000, thin=5,chains=3,
                         refresh=0,
                         save_warmup = F)


#####02.1 (random intercept and slope) Longitudinal ################
X_longi <- as.matrix(model.matrix(longiForm ,
                                  simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi)

#
model_2slope <- stan(file = "1_longit_slope.stan",
                       data = stan_long_data,
                       iter=4000,warmup=3000, thin=5,chains=3,
                       refresh=0,
                       save_warmup = F)

#####02.2 (random intercept and slope) Longidutindal with visit ############################
####nvisitAge
X_longi <- as.matrix(model.matrix(longiForm_nvisitAge,
                                  simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi)

#
model_2slope_nvist <- stan(file = "1_longit_slope.stan",
                             data = stan_long_data,
                            iter=4000,warmup=3000, thin=5,chains=3,
                             refresh=0,
                             save_warmup = F  )

#####02.3 (random intercept and slope) joint model (random intercept and slope)############################
#recurrent event covariates
X_event = as.matrix(model.matrix(survForm, simdata)[, -c(1)])

#longitudinal covaraites
X_longi <- as.matrix(model.matrix(longiForm,
                                  simdata)[,-1])

#stan dataset
stan_jm0_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                       y = simdata$y, 
                       id = simdata$id,
                       M_y = dim(X_longi)[2], X_longi= X_longi,
                       gap_time=simdata$gaptime,
                       M_x_event=dim(X_event)[2], X_event= X_event)

#
model_2slope_sjm <- stan(file = "1_JM_slope_HR.stan",
                           data = stan_jm0_data,
                           iter=4000,warmup=3000, thin=5,chains=3,
                           refresh=0,
                           save_warmup = F)

#####03.1 (GP) longitudinal (GP for subject)##############################
####
seq = 1:length(simdata$id)
locate0 = seq[!duplicated(simdata[,c("id")])]
locateN = seq[!duplicated(simdata[,c("id")], fromLast = T)]

X_longi <- as.matrix(model.matrix( longiForm,
                                   simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi,
                        locate0=locate0,
                        locateN=locateN)

#
model_3GP <- stan(file = "2_longit_GP.stan",
                    data = stan_long_data,
                   iter=4000,warmup=3000, thin=5,chains=3,
                    refresh=0,
                    save_warmup = F)

#####03.2 (GP) longitudinal with visit (GP for subject)################################
seq = 1:length(simdata$id)
locate0 = seq[!duplicated(simdata[,c("id")])]
locateN = seq[!duplicated(simdata[,c("id")], fromLast = T)]

X_longi <- as.matrix(model.matrix( longiForm_nvisitAge,
                                   simdata)[,-1])

stan_long_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi,
                        locate0=locate0,
                        locateN=locateN)

model_3GP_nvist <- stan(file = "2_longit_GP.stan",
                          data = stan_long_data,
                          iter=4000,warmup=3000, thin=5,chains=3,
                          refresh=0,
                          save_warmup = F)

#####03.3 (GP) joint model (GP for subject)####################################
#recurrent event covariates
X_event = as.matrix(model.matrix(survForm , simdata)[, -c(1)])

#longitudinal covaraites
seq = 1:length(simdata$id)
locate0 = seq[!duplicated(simdata[,c("id")])]
locateN = seq[!duplicated(simdata[,c("id")], fromLast = T)]

X_longi <- data.frame(model.matrix(longiForm ,
                                   simdata)[,-1])

#stan dataset
stan_jmGP_data <- list( N = length(simdata$id), Ngroups = length(unique(simdata$id)),
                        y = simdata$y, 
                        id = simdata$id,
                        M_y = dim(X_longi)[2], X_longi= X_longi,
                        gap_time=simdata$gaptime,
                        M_x_event=dim(X_event)[2], X_event= X_event,
                        locate0=locate0,
                        locateN=locateN)

##
model_3GP_sjm <- stan(file = "2_JM_GP_HR.stan",
                        data = stan_jmGP_data,
                        iter=4000,warmup=3000, thin=5,chains=3,
                        refresh=0,
                        save_warmup = F) 
