# LSTAT2820 - Mémoire
# Programmation en R
# 2019/2020
# Mémoire
# By Clement Laloux
# Tool for simulation studies

# Clear Working Directory
rm(list = setdiff(ls(), c()))

# Library ####
library(invgamma)
library(rstan)
# Time data type
library(lubridate)
library(timeDate)
# Highest posterior density interval
library(HDInterval)

# Set Working Directory ####
wd <- ""
setwd(wd)
rm(wd)

# Stan model definition ####
# Define the Stan model to draw a sample of the posterior distribution of lambda
Analyse_Prior_Stan = "
data{
int<lower = 0> N;         // size of the dataset used
real y[N];                // y = InterArTm (except 0 value)
real<lower = 0> alpha_p;             // Prior on alpha for the inverse gamma distribution
real<lower = 0> beta_p;              // Prior on beta for the inverse gamma distribution
}
parameters{
real<lower = 0> lambda;    // parameter of the exponential distribution
}
model{
lambda ~ inv_gamma(alpha_p, beta_p); // Prior on lambda

for (i in 1:N)
{
  target += exponential_lpdf(y[i]|1/lambda); // No censored data, only density function
}
}"
# Transform the Stan code to be used in R then save it in an RDS file to save computational time
#Analyse_Prior_R = stan_model(model_code = Analyse_Prior_Stan)
#saveRDS(Analyse_Prior_R, paste0("Analyse_Prior_R", ".rds"))
Analyse_Prior_R = readRDS("Analyse_Prior_R.rds")

# Simulation studies ####
# n_sim: Number of simulation to be run, i.e. the simulated data that mimic real clinical trials.
# real_t_tot: The theoretical duration of the simulated trials.
# n: The number of patients to be recruited.
# Start_date: The start date of the trials to be able to compute the distributions of arrival per month.
# m: the number of subjects already in the study at the time of the predictions. There is 
#    no limit on the number of values, but the greater they are, the longer it will take to run the function.
# n_pred: the number of waiting times to be generated in the predictions.
# diff_month: the planned duration error, i.e. the difference in months between real trials and expected duration.
# V_lambda: the levels of information considered. There is no limit on the number of values, but the greater
#           they are, the longer it will take to run the function.
# Time_points: The length of the interval between time points for predictions.
# Median_precision: The method to compute the median precision. Either the Mean squared error (MSE) or
#                   the Mean absolute error (MAE).

robustness_sim <- function(n_sim, real_t_tot, n, Start_date = date("2019/12/1"), 
                           m, n_pred, diff_month, V_lambda, Time_points,
                           Median_precision = c("MSE", "MAE")){
  
  # First Part: Generate the n_sim simulations ####
  Real_trials <- lapply(X = 1:n_sim, FUN = function(jj){
    # Define real lambda
    lambda_real <- (real_t_tot*(365.25/12))/n
    
    # Compute the waiting time based on the real value of lambda for the simulation
    Waiting_Time_real <- rexp(n, 1/lambda_real)
    Waiting_Time_real_cumsum <- cumsum(Waiting_Time_real)
    
    # Create real subject dataset
    subject_real <- data.frame(id = 1:n, arrival_date = Start_date, inter_arrival = Waiting_Time_real, 
                               cum_inter_arrival = Waiting_Time_real_cumsum)
    subject_real$arrival_date <- subject_real$arrival_date+subject_real$cum_inter_arrival
    
    # Trial duration in month ####
    Tot_day <- as.numeric(subject_real$arrival_date[n]-Start_date)
    correct_label <- Start_date+1:round(Tot_day)
    correct_label <- unique(paste(month(correct_label), year(correct_label)))
    t_tot_observed <- length(correct_label)
    
    return(list(t_tot_observed, subject_real))
  })
  
  # Get Real_trials dataset
  Real_trials_dataset <- lapply(Real_trials, "[[", 2)
  
  # Get real trial duration
  Real_trials_duration <- unlist(lapply(Real_trials, "[[", 1))
  print(summary(Real_trials_duration))
  
  # Second Part: Define general summary tables for results ####
  # All the remaining tables will resume information from tables in each simulation
  # We will compute the mean of results
  # General IC width ####
  # Create base table and fill it in the loops
  General_IC_width <- as.data.frame(matrix(0, ncol = 1+length(V_lambda), nrow = 5))
  rownames(General_IC_width) <- c(3, 6, 9, 12, "t_tot")
  colnames(General_IC_width) <- c("Time Point (in month)", paste("V = ", V_lambda))
  
  # Create a list of m elements for all 
  General_IC_width_Evol_list <- rep(list(General_IC_width), (length(m)+1))
  
  # General Median Precision ####
  # Create base table and fill it in the loops
  General_Med_precision <- as.data.frame(matrix(0, ncol = 1+length(V_lambda), nrow = 5))
  rownames(General_Med_precision) <- c(3, 6, 9, 12, "t_tot")
  colnames(General_Med_precision) <- c("Time Point (in month)", paste("V = ", V_lambda)) 
  
  # Create a list of m elements for all 
  General_Median_precision_list <- rep(list(General_Med_precision), (length(m)+1))
  
  # General real recruitment covered by 90% IC ####
  # Create base table and fill it in the loops
  General_IC_cover <- as.data.frame(matrix(0, ncol = 1+length(V_lambda), nrow = 5))
  rownames(General_IC_cover) <- c(3, 6, 9, 12, "t_tot")
  colnames(General_IC_cover) <- c("Time Point (in month)", paste("V = ", V_lambda)) 
  
  # Create a list of m elements for all 
  General_IC_coverage_list <- rep(list(General_IC_cover), (length(m)+1))
  
  # Third Part: Perform the simulation studies ####
  # Store the results of each simulation
  simulations <- lapply(1:n_sim, FUN = function(jj){
    print(jj)
    # Real data ####
    # Import real trial data
    subject_real <- Real_trials_dataset[[jj]]
    
    # Define real_t_tot
    real_t_tot <- Real_trials_duration[jj]
    
    # Define t_tot prior ####
    # The expected duration is the real duration minus the planned duration error
    t_tot <- real_t_tot - diff_month
    print(t_tot)
    
    # Define Time point for each m ####
    # They represent a sequence of values at regular intervals, which is define in the function by
    # the Time_points argument
    time_point <- lapply(X = 1:(length(m)+1), FUN = function(jj){
      set2 <- seq(Time_points, by = Time_points, length.out = 5)
      if(jj==1) set2
      else{
        # Compute the number of months between Start_date and last arrival
        Nbr_months <- Start_date+1:ceiling(subject_real$arrival_date[m[jj-1]]-Start_date)
        Nbr_months <- length(unique(paste(month(Nbr_months), year(Nbr_months))))
        
        if(sum((Nbr_months+set2) > real_t_tot) == 0) (Nbr_months+set2)
        else{
          # Take the month pred before real_t_tot
          lower <- (Nbr_months+set2)[which((Nbr_months+set2) < real_t_tot)]
          # Save
          c(lower, real_t_tot)
        }
      } 
    })
    
    # IC width ####
    # Create tables and lists to stock the tables
    # IC_width_Evol_list gives the different widths of 90% IC for one m value and multiple variance of lambda
    # at several time points
    IC_width_Evol_list <- list()
    # Create base table and fill it in the loops
    IC_width <- as.data.frame(matrix(ncol = 1+length(V_lambda), nrow = length(time_point[[1]])))
    colnames(IC_width) <- c("Time Point (in month)", paste("V = ", V_lambda))
    
    # Median Precision ####
    # Median_precision_list compute either the MSE or the MAE
    # between the real recruitment number and the prediction for one m value
    # and multiple variance of lambda
    Median_precision_list <- list()
    # Create base table and fill it in the loops
    Med_precision <- as.data.frame(matrix(ncol = 1+length(V_lambda), nrow = length(time_point[[1]])))
    colnames(Med_precision) <- c("Time Point (in month)", paste("V = ", V_lambda)) 
    
    # Prediction details ####
    # IC_details gives all the 90% IC informations and median prediction estimate
    # for one m and one variance of lambda at several time points
    IC_details <- list()
    # Create base table and fill it in the loops
    Pred_comparison <- as.data.frame(matrix(ncol = 6, nrow = length(time_point[[1]])))
    colnames(Pred_comparison) <- c("Time Point (in month)", "Real recruitment", 
                                   "Lower 90% IC post", "Median post", "Upper 90% IC post", "IC width")
    
    # Real recruitment covered by 90% IC ####
    # IC_coverage_list tells us whether or not the 90% IC contains the true recruitment numbers for one value of m and 
    # multiple variance of lambda
    IC_coverage_list <- list()
    # Create base table and fill it in the loops
    IC_cover <- as.data.frame(matrix(ncol = 1+length(V_lambda), nrow = length(time_point[[1]])))
    colnames(IC_cover) <- c("Time Point (in month)", paste("V = ", V_lambda)) 
    
    # Expected waiting time of the study ####
    # Compute prior on lambda knowing waiting time has an exponential distribution with lambda, the scale parameter.
    # Prior mean on lambda is measured by the expected waiting time between patient (in days)
    lambda_0 <- (t_tot*(365.25/12))/n
    
    # Bayesian setting, an inverse gamma distribution is considered for the parameter lambda of the exponential distribution
    # Find prior value for alpha and beta considering E(lambda) = lambda_0 and V(lambda) = constant introduced in the function
    E_lambda <- lambda_0
    cat("E(lambda) =", E_lambda)
    
    # Stan Model ####
    # Introduce the Stan model define above
    Analyse_Prior_R = readRDS("Analyse_Prior_R.rds")
    
    # Loop on V(lambda) ####
    # For each V(lambda) introduced in the function, compute the different results of interest
    for(e in 1:length(V_lambda)){
      # Compute prior predictions on the trial
      # Compute prior paramater values
      alpha_0 <- E_lambda^2/V_lambda[e]+2
      beta_0 <- E_lambda*(alpha_0-1)
      
      # Sample 1000 different waiting time prior prediction, with for each waiting time, 
      # a new lambda from the inverse gamma distribution (in order to have IC)
      # Then with the trial start date, define subjects arrival dates and compute recruitment per month
      Prior_pred <- lapply(1:1000, FUN = function(jj){
        # For each waiting time w_i, sample a lambda from the prior distribution and use it in the exponential
        lambda_0 <- rinvgamma(n_pred, alpha_0, rate = beta_0)
        w <- rexp(n_pred, 1/lambda_0)
        
        # Then compute arrival date from trial start date
        d <- Start_date+cumsum(w)
        
        # Create a new variable that includes only month and year
        m_y <- paste(month(d), year(d))
        
        # Count the number of subject recruited per month
        # Problem, 1 2021 comes before 2 2020
        # Thus create a prevariable with label in correct order and work in another lapply
        # Create a variable for the maximum number of days to recruit n_pred subjects
        t_tot_1000 <- as.numeric(d[n_pred]-Start_date)
        
        # Return the arrival dates, the month and year and the trial length
        list(d, m_y, t_tot_1000)
      })
      
      # Get the maximum difference in number of days
      max_t_tot <- max(unlist(lapply(Prior_pred, "[[", 3)))
      # To ensure the real t_tot is included in the maximum time duration
      max_t_tot <- max(max_t_tot, as.numeric(subject_real$arrival_date[n]-Start_date))
      
      # Create the correct label order
      # Compute a variable with all the dates from trial start date to trial last date possible in prediction
      # Then only keep one month and year combination in correct order
      correct_label <- Start_date+1:round(max_t_tot)
      correct_label <- unique(paste(month(correct_label), year(correct_label)))
      
      # Count the cumulative sum of subject recruited per month in each prediction
      Prior_pred <- lapply(1:1000, FUN = function(jj){
        d <- unlist(Prior_pred[[jj]][2])
        d <- factor(d, levels = correct_label)
        d <- table(d)
        list(cumsum(d))
      })
      
      # For each month, get the 1000 cumulative sum of subject recruited in the prediction
      Prior_pred <- lapply(1:length(correct_label), FUN = function(jj){
        d <- lapply(Prior_pred, "[[", 1)
        d <- unlist(unlist(lapply(d, "[[", jj)))
        # Get the median value of subjects recruited in each month and the 90% HDI
        list(hdi(d, credMass = .90)[1], median(d), hdi(d, credMass = .90)[2])
      })
      
      # Retrieve the quantities of interest from the prior predictive distribution (median and 90% HDI)
      q05_prior <- round(unlist(lapply(Prior_pred, "[[", 1)), 0)
      q5_prior <- round(unlist(lapply(Prior_pred, "[[", 2)), 0)
      q95_prior <- round(unlist(lapply(Prior_pred, "[[", 3)), 0)
      
      # Count real data per month
      # Count per month
      real_d <- paste(month(subject_real$arrival_date), year(subject_real$arrival_date))
      real_d <- factor(real_d, levels = correct_label)
      real_d <- as.numeric(table(real_d))
      
      # IC_details
      # Create a list inside the variance
      IC_details_V <- list()
      
      # Update table
      Pred_comparison[,1] <- time_point[[1]]
      Pred_comparison[,2] <- cumsum(real_d)[time_point[[1]]]
      Pred_comparison[,3] <- q05_prior[time_point[[1]]]
      Pred_comparison[,4] <- q5_prior[time_point[[1]]]
      Pred_comparison[,5] <- q95_prior[time_point[[1]]]
      Pred_comparison[,6] <- Pred_comparison[,5] - Pred_comparison[,3]
      
      IC_details_V[[1]] <- Pred_comparison
      
      # IC_width_Evol_list
      # Store the different results in their respective tables
      # if first column to be fulfill, just assigns IC_width
      if(e == 1){
        # IC_width
        IC_width[,1] <- time_point[[1]]
        IC_width[,1+e] <- Pred_comparison[,6]
        IC_width_Evol_list[[1]] <- IC_width
        
        # Median Precision
        Med_precision[,1] <- time_point[[1]]
        
        # Compute the Median precision depending on the type chosen
        if(Median_precision == "MSE") Med_precision[,1+e] <- (Pred_comparison[,2] - Pred_comparison[,4])^2
        else if(Median_precision == "MAE") Med_precision[,1+e] <- abs(Pred_comparison[,2] - Pred_comparison[,4])
        else return(print("No Median precision selected, error"))
        
        Median_precision_list[[1]] <- Med_precision
        
        # 90% IC Coverage
        IC_cover[,1] <- time_point[[1]]
        # if the true recuitment value is covered by 90% IC, value is 1, otherwise 0
        # index gives time points where true value is included
        index <- which(Pred_comparison[,2] >= Pred_comparison[,3] & Pred_comparison[,2] <= Pred_comparison[,5])
        # if no time points, all values are 0
        IC_cover[,1+e] <- 0
        # if some time points, put their values at 1
        if(length(index) != 0) IC_cover[index,1+e] <- 1
        
        IC_coverage_list[[1]] <- IC_cover
        
      } # end if(e == 1)
      # else fulfill the right IC_width_Evol_list
      else{
        # IC_width
        #IC_width_Evol_list[[1]][,1] <- time_point[[1]]
        IC_width_Evol_list[[1]][,1+e] <- Pred_comparison[,6]
        
        # Median Precision
        # Compute the Median precision depending on the type chosen
        if(Median_precision == "MSE") Median_precision_list[[1]][,1+e] <- (Pred_comparison[,2] - Pred_comparison[,4])^2
        else if(Median_precision == "MAE") Median_precision_list[[1]][,1+e] <- abs(Pred_comparison[,2] - Pred_comparison[,4])
        else return(print("No Median precision selected, error"))
        
        # 90% IC Coverage
        #IC_coverage_list[[1]][,1] <- time_point[[1]]
        # index gives time points where true value is included
        index <- which(Pred_comparison[,2] >= Pred_comparison[,3] & Pred_comparison[,2] <= Pred_comparison[,5])
        # if no time points, all values are 0
        IC_coverage_list[[1]][,1+e] <- 0
        
        # if some time points, put their values at 1
        if(length(index) != 0) IC_coverage_list[[1]][index,1+e] <- 1
      } #end else of if(e == 1)
      
      # Then consider the predictions when m subjects are in the trial
      # Model 1 : Model with informative prior
      # Fit the model depending on the information we have on real data to adjust prediction based on prior
      # Inter-Arrival Time Model ==> Time to recruit a new patient
      # Here we will consider n data at a time, we will use a exponential to predict the arrival of n-m patients
      # within a certain period
      
      for(i in 1:length(m)){
        # Model 1 Estimation
        # Define the number of subjects in the study
        m_test <- m[i]
        
        # Define model parameters 
        # Waiting Time
        y <- subject_real$inter_arrival[1:m_test]
        # Size of the dataset used
        N_y <- length(y)
        # Prior on alpha and beta (already defined above)
        
        # Create list for Stanfitting
        arrlist = list(y=y, N=N_y, alpha_p=alpha_0, beta_p=beta_0)
        
        # Fit the model
        Analyse_Prior_Fit = sampling(object = Analyse_Prior_R, data = arrlist, chains = 4, 
                                     iter = 2000, warmup = 1000, thin = 4)
        
        # Extract chains
        Analyse_Prior_Para1 = extract(Analyse_Prior_Fit)
        print(summary(Analyse_Prior_Para1$lambda))
        
        # Compute Prediction
        # Sample 1000 different waiting time posterior prediction, with for each waiting time, 
        # a new lambda from the posterior lambda distribution (in order to have IC)
        # Then with the trial start date, define subjects arrival dates and compute recruitment per month
        Post_pred <- lapply(X = 1:1000, FUN = function(jj){
          # Sample one lambda from posterior
          lambda_post <- Analyse_Prior_Para1$lambda[jj]
          w <- rexp(n_pred, 1/lambda_post)
          
          # Then compute arrival date from trial start date
          d <- c(subject_real$arrival_date[1:m_test], subject_real$arrival_date[m_test]+cumsum(w))
          
          # Create a new variable that includes only month and year
          m_y <- paste(month(d), year(d))
          
          # Create a variable for the maximum number of days to recruit n_pred subjects
          t_tot_1000 <- as.numeric(d[n_pred]-Start_date)
          
          # Return the arrival dates, the month and year and the trial length
          list(d, m_y, t_tot_1000)
        })
        
        # Get the maximum difference in number of days
        max_t_tot <- max(unlist(lapply(Post_pred, "[[", 3)))
        # To ensure the real t_tot is included in the maximum time duration
        max_t_tot <- max(max_t_tot, as.numeric(subject_real$arrival_date[n]-Start_date))
        
        # Create the correct label order
        # Compute a variable with all the dates from trial start date to trial last date possible in prediction
        # Then only keep one month and year combination in correct order
        correct_label <- Start_date+1:round(max_t_tot)
        correct_label <- unique(paste(month(correct_label), year(correct_label)))
        
        # Count the cumulative sum of subject recruited per month in each prediction
        Post_pred <- lapply(1:1000, FUN = function(jj){
          d <- unlist(Post_pred[[jj]][2])
          d <- factor(d, levels = correct_label)
          d <- table(d)
          list(cumsum(d))
        })
        
        # For each month, get the 1000 cumulative sum of subject recruited in the prediction
        Post_pred <- lapply(1:length(correct_label), FUN = function(jj){
          d <- lapply(Post_pred, "[[", 1)
          d <- unlist(unlist(lapply(d, "[[", jj)))
          # Get the median value of subjects recruited in each month and the 90% HDI
          list(hdi(d, credMass = .90)[1], median(d), hdi(d, credMass = .90)[2])
        })
        
        # Retrieve the quantities of interest from the posterior predictive distribution (median and 90% HDI)
        q05_post <- unlist(lapply(Post_pred, "[[", 1))
        q5_post <- round(unlist(lapply(Post_pred, "[[", 2)), 0)
        q95_post <- unlist(lapply(Post_pred, "[[", 3))
        
        # IC_details
        # Update table
        Pred_comparison[,1] <- time_point[[1+i]]
        Pred_comparison[,2] <- cumsum(real_d)[time_point[[1+i]]]
        Pred_comparison[,3] <- q05_post[time_point[[1+i]]]
        Pred_comparison[,4] <- q5_post[time_point[[1+i]]]
        Pred_comparison[,5] <- q95_post[time_point[[1+i]]]
        Pred_comparison[,6] <- Pred_comparison[,5] - Pred_comparison[,3]
        
        IC_details_V[[i+1]] <- Pred_comparison
        
        # IC_width_Evol_list
        # Store the different results in their respective tables
        # if first column to be fulfill, just assigns IC_width
        if(e == 1){
          # IC_width
          IC_width[,1] <- time_point[[1+i]]
          IC_width[,1+e] <- Pred_comparison[,6]
          IC_width_Evol_list[[1+i]] <- IC_width
          
          # Median Precision
          Med_precision[,1] <- time_point[[1+i]]
          
          # Compute the Median precision depending on the type chosen
          if(Median_precision == "MSE") Med_precision[,1+e] <- (Pred_comparison[,2] - Pred_comparison[,4])^2
          else if(Median_precision == "MAE") Med_precision[,1+e] <- abs(Pred_comparison[,2] - Pred_comparison[,4])
          else return(print("No Median precision selected, error"))
          
          Median_precision_list[[1+i]] <- Med_precision
          
          # 90% IC Coverage
          IC_cover[,1] <- time_point[[1+i]]
          # if the true recuitment value is covered by 90% IC, value is 1, otherwise 0
          # index gives time points where true value is included
          index <- which(Pred_comparison[,2] >= Pred_comparison[,3] & Pred_comparison[,2] <= Pred_comparison[,5])
          # if no time points, all values are 0
          IC_cover[,1+e] <- 0
          # if some time points, put their values at 1
          if(length(index) != 0) IC_cover[index,1+e] <- 1
          
          IC_coverage_list[[1+i]] <- IC_cover
        } 
        # else fulfill the right IC_width_Evol_list
        else{
          # IC_width
          #IC_width_Evol_list[[1+i]][,1] <- time_point[[1+i]]
          IC_width_Evol_list[[1+i]][,1+e] <- Pred_comparison[,6]
          
          # Median Precision
          # Median Precision
          # Compute the Median precision depending on the type chosen
          if(Median_precision == "MSE") Median_precision_list[[1+i]][,1+e] <- (Pred_comparison[,2] - Pred_comparison[,4])^2
          else if(Median_precision == "MAE") Median_precision_list[[1+i]][,1+e] <- abs(Pred_comparison[,2] - Pred_comparison[,4])
          else return(print("No Median precision selected, error"))
          
          # 90% IC Coverage
          # index gives time points where true value is included
          index <- which(Pred_comparison[,2] >= Pred_comparison[,3] & Pred_comparison[,2] <= Pred_comparison[,5])
          # if no time points, all values are 0
          IC_coverage_list[[1+i]][,1+e] <- 0
          
          # if some time points, put their values at 1
          if(length(index) != 0) IC_coverage_list[[1+i]][index,1+e] <- 1
        }
        
      } # fin boucle m
      
      # Update list IC_details
      IC_details[[e]] <- IC_details_V
      
    } #end V(lambda) loop
    
    # Update General Tables ####
    # General tables with a table by m value
    for(i in 1:(length(m)+1)){
      # General IC width
      General_IC_width_Evol_list[[i]] <<- General_IC_width_Evol_list[[i]] + IC_width_Evol_list[[i]]
      
      # General median Precision
      General_Median_precision_list[[i]] <<- General_Median_precision_list[[i]] + Median_precision_list[[i]]
      
      # General real recruitment covered by 90% IC
      General_IC_coverage_list[[i]] <<- General_IC_coverage_list[[i]] + IC_coverage_list[[i]]
      
      if(jj == n_sim){
        # General IC width
        General_IC_width_Evol_list[[i]] <<- General_IC_width_Evol_list[[i]]/n_sim
        #General_IC_width_Evol_list[[i]][,1] <<- c(3, 6, 9, 12, "t_tot")
        
        # General median Precision
        General_Median_precision_list[[i]][] <<- General_Median_precision_list[[i]]/n_sim
        #General_Median_precision_list[[i]][,1] <<- c(3, 6, 9, 12, "t_tot")
        
        # General real recruitment covered by 90% IC
        General_IC_coverage_list[[i]][] <<- General_IC_coverage_list[[i]]/n_sim
        #General_IC_coverage_list[[i]][,1] <<- c(3, 6, 9, 12, "t_tot")
      }
    }
    
    # Return list at simulation level ####
    list(IC_details, IC_width_Evol_list, Median_precision_list, IC_coverage_list)
  }) # end lapply
  
  # Return all the all the results from simulation and all the general tables
  return(list(simulations, General_IC_width_Evol_list, General_Median_precision_list, General_IC_coverage_list))
  
} # end function robustness_sim