start_time <- Sys.time()
library("trtf")
library("tram")
library("party")
library("parallel")
library("survival")
set.seed(9235)

source("tr_competitors.R")
### load data from tr_simulated.rda generated by running tr_data_simulation.R
load('tr_simulated_data.rda')

### Number of repetitions
NSim <- length(simdat[[1]])
ret <- args[rep(1:nrow(args), each = NSim),]
simdat <- do.call("c", simdat)
idx <- 1:nrow(ret)

### Log-likelihood estimates will be stored in ret
ret$tweib <- ret$tcoxph <- 0
ret$w_alpha_beta <- ret$w_theta_beta <- ret$w_theta_theta <- 0
ret$bs_alpha_beta <- ret$bs_theta_beta <- ret$bs_theta_theta <- 0
ret$loglik <- unlist(do.call("c", loglik))

### a filename to save results
res_filename <- 'tr_results_empeval.rda'

### number of cores for parallel computing
ncores <- 15

### number of randomly preselected variables for splitting
### for low-dimentional data it is set to the number of variables
### for high-dimensional data it is set to the square root of the number of variables
mtry <- lapply(idx, function(i) {
                         ifelse(ret[i, "p"] <= 5,
                         ncol(simdat[[i]]$train) - 2,
                         ceiling(sqrt(ncol(simdat[[i]]$train) - 2)))})

### folder to save Bs(theta, theta) model for each somulated dataset
tf_folder <- 'results_tf'
dir.create(tf_folder)
tf_dir <- paste(getwd(), tf_folder, sep='/')

### Bs(theta, theta)
bstt <- function(i){
    if (i %% NSim == 0){
         print(i)
     }

    tf <- bs_theta_theta(simdat[[i]]$train,
                         simdat[[i]]$test,
                         strata = simdat[[i]]$train$trt,
                         mtry = mtry[[i]])
    ## save the learned forest
    save(tf, file=paste0(tf_dir, '/tf', i, '.rda'))

    return(unclass(tf$loglik))
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, bstt, mc.cores = ncores)
ret$bs_theta_theta <- unlist(f0)
print('Model Bs_theta_theta is done!')
save(ret, file = res_filename)

### Bs(alpha, beta)
bsab <- function(i){
    if (i %% NSim == 0){
         print(i)
     }
    ## load the corresponding Bs(theta,theta) forest
    load(paste0(tf_dir, '/tf', i, '.rda'))
    ## tf is the loaded forest
    unclass(bs_alpha_beta(tf, 
                          simdat[[i]]$train,
                          simdat[[i]]$test,
                          strata = simdat[[i]]$train$trt)$loglik)
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, bsab, mc.cores = ncores)
ret$bs_alpha_beta <- unlist(f0)
print('Model Bs_alpha_beta is done!')
save(ret, file = res_filename)


### Bs(theta, beta)
bstb <- function(i){
    if (i %% NSim == 0){
         print(i)
     }

    ## load the corresponding Bs(theta,theta) forest
    load(paste0(tf_dir, '/tf', i, '.rda'))
    ## tf is the loaded forest
    unclass(bs_theta_beta(tf, 
                          simdat[[i]]$train,
                          simdat[[i]]$test,
                          strata = simdat[[i]]$train$trt)$loglik)
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, bstb, mc.cores = ncores)
ret$bs_theta_beta <- unlist(f0)
print('Model Bs_theta_beta is done!')
save(ret, file = res_filename)


### W(alpha, beta)
wab <- function(i){
    if (i %% NSim == 0){
        print(i)
    }
    ## load the corresponding Bs(theta,theta) forest
    load(paste0(tf_dir, '/tf', i, '.rda'))
    ## tf is the loaded forest
    unclass(w_alpha_beta(tf, 
                         simdat[[i]]$train,
                         simdat[[i]]$test,
                         strata = simdat[[i]]$train$trt)$loglik)
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, wab, mc.cores = ncores)
ret$w_alpha_beta <- unlist(f0)
print('Model W_alpha_beta is done!')
save(ret, file = res_filename)

### W(theta, beta)
wtb <- function(i){
    if (i %% NSim == 0){
         print(i)
     }
    ## load the corresponding Bs(theta,theta) forest
    load(paste0(tf_dir, '/tf', i, '.rda'))
    ## tf is the loaded forest
    unclass(w_theta_beta(tf, 
                         simdat[[i]]$train,
                         simdat[[i]]$test,
                         strata = simdat[[i]]$train$trt)$loglik)
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, wtb, mc.cores = ncores)
ret$w_theta_beta <- unlist(f0)
print('Model W_theta_beta is done!')
save(ret, file = res_filename)

### W(theta, theta)
wtt <- function(i){
    if (i %% NSim == 0){
         print(i)
     }
    ## load the corresponding Bs(theta,theta) forest
    load(paste0(tf_dir, '/tf', i, '.rda'))
    ## tf is the loaded forest
    unclass(w_theta_theta(tf, 
                          simdat[[i]]$train,
                          simdat[[i]]$test,
                          strata = simdat[[i]]$train$trt)$loglik)
}
### apply in parallel to all simulated datasets
f0 <- mclapply(idx, wtt, mc.cores = ncores)
ret$w_theta_theta <- unlist(f0)
print('Model W_theta_theta is done!')
save(ret, file = res_filename)

end_time <- Sys.time()
print(end_time - start_time)

sessionInfo()
