###################################################
# Plotting sensitivity analysis                   #
# Input: simulation data, empirical data          #
# Output: Figure S4, S5                           #
# Original: June 16, 2022                         #
# Last modified: September 20, 2022               #
###################################################

rm(list = ls())
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))

library(ggpubr); library(reshape2)

Y_EI=c("CLINWTB_EI","MEANWST_EI","HR_EI","SBP_EI","DBP_EI","VO2_EI")
Y_sod=c("CLINWTB_sod","MEANWST_sod","HR_sod","SBP_sod","DBP_sod","VO2_sod")
Y_pot=c("CLINWTB_pot","MEANWST_pot","HR_pot","SBP_pot","DBP_pot","VO2_pot")
Y_pro=c("CLINWTB_pro","MEANWST_pro","HR_pro","SBP_pro","DBP_pro","VO2_pro")

TrueBetas <- read.csv("true_betas.csv")[,2]

TrueBeta <- c(rep(TrueBetas[1:6],4),rep(TrueBetas[7:12],4),
              rep(TrueBetas[13:18],4),rep(TrueBetas[19:24],4))

##### Eta #####
qetas = c(0.25,0.5,1,2)

EI <- c(); Sod <- c(); Pot <- c(); Pro <- c() 
for(j in 1:4){
  qeta =  qetas[j]

    for(i in 1:6){
    EI <- c(EI, paste0("SA_eta/",Y_EI[i],"_eta_",qeta,".csv"))
    Sod <- c(Sod, paste0("SA_eta/",Y_sod[i],"_eta_",qeta,".csv"))
    Pot <- c(Pot, paste0("SA_eta/",Y_pot[i],"_eta_",qeta,".csv"))
    Pro <- c(Pro, paste0("SA_eta/",Y_pro[i],"_eta_",qeta,".csv"))
  }
}

file_name <- c(EI, Sod, Pot, Pro)

NI = c(rep("EI",24),rep("sodium",24),rep("potassium",24),rep("protein",24))
HO = rep(c("BW","WC","HR","SBP","DBP","VO2"),4*4)
qeta = rep(c(rep("0.25",6),rep("0.5",6),rep("Base",6),rep("2",6)),4)

betas <- list()
for(i in 1:96){
  
  beta1 = TrueBeta[i]
  
  df <- read.csv(file_name[i],header=FALSE)
  colnames(df)<-c("SR0","SR1","SRlow","SRup","G0","G1","Glow","Gup","BIO0","BIO1","BIOlow","BIOup" )
  
  beta_hat_SR = df[,2] 
  beta_hat_G = df[,6] 
  beta_hat_BIO = df[,10]
  
  # BIO
  Mean_bias_BIO = mean(beta_hat_BIO,na.rm = T)
  (bias_BIO = Mean_bias_BIO-beta1)
  
  # SR
  Mean_bias_SR = mean(beta_hat_SR,na.rm = T)
  (bias_SR = Mean_bias_SR - beta1)
  
  # G
  Mean_bias_G = mean(beta_hat_G,na.rm = T)
  (bias_G = Mean_bias_G-beta1)
  
  betas[[i]] <- data.frame(NI = rep(NI[i],3),
                           HO = rep(HO[i],3),
                           qeta = rep(qeta[i],3),
                           Bias = c(bias_BIO,bias_SR,bias_G),
                           metric = c("BIO","SR","G"))
  
}

Betas <- do.call("rbind",betas)
Betas$qeta <- factor(Betas$qeta,levels=c("0.25","0.5","Base","2"),ordered=TRUE)
#write.csv(Betas,file="performance_measures.csv")

##### Sigma #####
qsigmas = c(0.25,0.5,1,2)

EI <- c(); Sod <- c(); Pot <- c(); Pro <- c() 
for(j in 1:4){
  qsigma =  qsigmas[j]
  
  for(i in 1:6){
    EI <- c(EI, paste0("SA_sigma/",Y_EI[i],"_sigma_",qsigma,".csv"))
    Sod <- c(Sod, paste0("SA_sigma/",Y_sod[i],"_sigma_",qsigma,".csv"))
    Pot <- c(Pot, paste0("SA_sigma/",Y_pot[i],"_sigma_",qsigma,".csv"))
    Pro <- c(Pro, paste0("SA_sigma/",Y_pro[i],"_sigma_",qsigma,".csv"))
  }
}

file_name_sigma <- c(EI, Sod, Pot, Pro)

NI = c(rep("EI",24),rep("sodium",24),rep("potassium",24),rep("protein",24))
HO = rep(c("BW","WC","HR","SBP","DBP","VO2"),4*4)
qsigma = rep(c(rep("0.25",6),rep("0.5",6),rep("Base",6),rep("2",6)),4)

betas_sigma <- list()
for(i in 1:96){
  
  beta1 = TrueBeta[i]
  
  df <- read.csv(file_name_sigma[i],header=FALSE)
  colnames(df)<-c("SR0","SR1","SRlow","SRup","G0","G1","Glow","Gup","BIO0","BIO1","BIOlow","BIOup" )
  
  beta_hat_SR = df[,2] 
  beta_hat_G = df[,6] 
  beta_hat_BIO = df[,10]
  
  # BIO
  Mean_bias_BIO = mean(beta_hat_BIO,na.rm = T)
  (bias_BIO = Mean_bias_BIO-beta1)
  
  # SR
  Mean_bias_SR = mean(beta_hat_SR,na.rm = T)
  (bias_SR = Mean_bias_SR - beta1)
  
  # G
  Mean_bias_G = mean(beta_hat_G,na.rm = T)
  (bias_G = Mean_bias_G-beta1)
  
  betas_sigma[[i]] <- data.frame(NI = rep(NI[i],3),
                           HO = rep(HO[i],3),
                           qsigma = rep(qsigma[i],3),
                           Bias = c(bias_BIO,bias_SR,bias_G),
                           metric = c("BIO","SR","G"))
}

Betas_sigma <- do.call("rbind",betas_sigma)
Betas_sigma$qsigma <- factor(Betas_sigma$qsigma,levels=c("0.25","0.5","Base","2"),ordered=TRUE)
#write.csv(Betas,file="performance_measures.csv")

##### Plot the results #####
Font1 = 40; Font2 = 27; Font3 = 28; SIZE=5; SIZE2=10; ALPHA = 0.7

theme_box = function() {
  theme_classic()+
    theme(aspect.ratio=1,
      axis.text.x = element_text(size = Font3),
      plot.title = element_text(face = "bold",size=Font1),
      plot.subtitle = element_text(face = "bold",size=Font2),
      axis.text = element_text(size =Font3),
      axis.title = element_text(size =Font2, face = "bold"),
      legend.title=element_blank(),
      legend.text = element_text(size=Font3))
}

##### Sample size #####
s100 <- read.csv("PerformanceMeasure100.csv")
s100$size = "Base"
s200 <- read.csv("PerformanceMeasure200.csv")
s200$size = "200"
s50 <- read.csv("PerformanceMeasure50.csv")
s50$size = "50"
s300 <- read.csv("PerformanceMeasure300.csv")
s300$size = "300"

bias_Size <- rbind(s100[,c("Bias", "metric","group","size")],
                   s200[,c("Bias", "metric","group","size")],
                   s50[,c("Bias", "metric","group","size")],
                   s300[,c("Bias", "metric","group","size")])

CP_Size <- rbind(s100[,c("CP", "metric","group","size")],
                 s200[,c("CP", "metric","group","size")],
                 s50[,c("CP", "metric","group","size")],
                 s300[,c("CP", "metric","group","size")])

data2plot <- melt(bias_Size)
data2plot$size <- factor(data2plot$size,levels=c("50","Base","200","300"),ordered=TRUE)

# Supplemental figure 4: Sample size
p1 <- ggplot(data=data2plot,aes(x=size,y=value,color=metric,shape=metric))+
  geom_point(position = position_jitter(width = 0.2),size=SIZE,alpha=ALPHA)+
  geom_hline(yintercept=0, size=1,color="grey")+#linetype = "dashed")+
  scale_shape_manual(values=c(15,16,17))+
  labs(title = "A", y = "Bias", x = "Sample size")+
  theme_box()

data2plot2 <- melt(CP_Size)
data2plot2$size <- factor(data2plot2$size,levels=c("50","Base","200","300"),ordered=TRUE)

p2 <- ggplot(data=data2plot2,aes(x=size,y=value,color=metric,shape=metric))+
  geom_point(position = position_jitter(width = 0.2),size=SIZE,alpha=ALPHA)+
  geom_hline(yintercept=95, size=1,color="grey")+#linetype = "dashed")+
  scale_shape_manual(values=c(15,16,17))+
  labs(title = "B", y = "Coverage Probability", x = "Sample size")+
  theme_box()

g2 <- ggarrange(p1, p2, ncol=2, common.legend = TRUE, legend="bottom")

#ggsave(filename = "Figures/FigureS4.png", g2, dpi = 400,width =600,height = 300,units = c("mm"))


# Supplemental figure 5: Varied eta and sigma
p1 <- ggplot(data=Betas,aes(x=qeta,y=Bias,color=metric,shape=metric))+
  geom_point(position = position_jitter(width = 0.2),size=SIZE,alpha=ALPHA)+
  geom_hline(yintercept=0, size=1,color="grey")+#linetype = "dashed")+
  scale_shape_manual(values=c(15,16,17))+
  labs(title = "A", y = "Bias", x = "Varied eta")+
  theme_box()


p2 <- ggplot(data=Betas_sigma,aes(x=qsigma,y=Bias,color=metric,shape=metric))+
  geom_point(position = position_jitter(width = 0.2),size=SIZE,alpha=ALPHA)+
  geom_hline(yintercept=0, size=1,color="grey")+#linetype = "dashed")+
  scale_shape_manual(values=c(15,16,17))+
  labs(title = "B", y = "Bias", x = "Varied sigma")+
  theme_box()

g <- ggarrange(p1, p2, ncol=2, common.legend = TRUE, legend="bottom")

ggsave(filename = "Figures/FigureS5.png", g, dpi = 400,width =600,height = 300,units = c("mm"))