# Simulation results

library(tidyverse)
library(ggplot2)
library(sjPlot)
library(CIDAtools)
library(kableExtra)
library(pROC)
library(broom.mixed)
library(ggpubr)
# library(see)


# load(file = "JDS_Submission/sim_results_gaussian.RData")
load(file = "sim_results_gaussian.RData")

## sims is a list of simulations corresponding to each row of params
sim <- sims[[1]]
get_results <- function(sim, val = "rmse_")  {
  
  n_methods <- sum(grepl(val, names(sim)))
  
  fit_data <- sim %>%
    select(starts_with(val)) %>% 
    mutate_all(as.numeric) %>% 
    pivot_longer(starts_with(val), names_to = "Method", names_prefix = val)  %>% 
    mutate(id = rep(1:nrow(sim), each = n_methods)) %>% 
    mutate(Method = factor(ifelse(Method == "full", "Full", toupper(Method)))) %>%
    mutate(value = value*100)
  
  fit_lmm <- try(lmerTest::lmer(value ~ Method + (1|id), data = fit_data), silent = T)
  if((class(fit_lmm)[1] == "try-error"))
    fit_lmm <- lm(value ~ Method, data = fit_data)
  fit_lmm
}

params$rmse_lmm <- lapply(sims, get_results)
params$fnr_lmm <- lapply(sims, get_results, val = "fnr_")
params$fdr_lmm <- lapply(sims, get_results, val = "fdr_")

#### Tables #### 

rmse_table <- lapply(params$rmse_lmm, tidy, conf.int = T) %>% 
  lapply(function(x) {
    filter(x, effect == "fixed") %>% 
      mutate(beta = paste0(round(estimate, 2), " (", round(conf.low, 1), ", ", round(conf.high, 1), ")")) %>%
      select(-effect, -group, -estimate:-conf.high) %>% 
      mutate(term = gsub("\\(|\\)", "", term)) %>% 
      mutate(term = gsub("Method", "beta.", term)) %>% 
      pivot_wider(names_from = term, values_from = c(beta)) 
  }) %>% bind_rows()

rmse_table <- params %>% 
  rename(p.missing = p_missing) %>%
  select(-contains("_lmm"), - plot) %>% 
  select(rev(names(.))) %>% 
  bind_cols(rmse_table)

fdr_table <- lapply(params$fdr_lmm, tidy, conf.int = T) %>% 
  lapply(function(x) {
    filter(x, effect == "fixed") %>% 
      mutate(beta = paste0(round(estimate, 2), " (", round(conf.low, 1), ", ", round(conf.high, 1), ")")) %>%
      select(-effect, -group, -estimate:-conf.high) %>% 
      mutate(term = gsub("\\(|\\)", "", term)) %>% 
      mutate(term = gsub("Method", "beta.", term)) %>% 
      pivot_wider(names_from = term, values_from = c(beta)) 
  }) %>% bind_rows()

fdr_table <- params %>% 
  rename(p.missing = p_missing) %>%
  select(-contains("_lmm"), - plot) %>% 
  select(rev(names(.))) %>% 
  bind_cols(fdr_table)  

fnr_table <- lapply(params$fnr_lmm, tidy, conf.int = T) %>% 
  lapply(function(x) {
    
    if(!length(x$effect)) {
      x$effect <- "fixed"
      x$group <- ""
    }
    
    
    filter(x, effect == "fixed") %>% 
      mutate(beta = paste0(round(estimate, 2), " (", round(conf.low, 1), ", ", round(conf.high, 1), ")")) %>%
      select(-effect, -group, -estimate:-conf.high) %>% 
      mutate(term = gsub("\\(|\\)", "", term)) %>% 
      mutate(term = gsub("Method", "beta.", term)) %>% 
      pivot_wider(names_from = term, values_from = c(beta)) 
  }) %>% bind_rows()

fnr_table <- params %>% 
  rename(p.missing = p_missing) %>%
  select(-contains("_lmm"), - plot) %>% 
  select(rev(names(.))) %>% 
  bind_cols(fnr_table) 

# rmse_table %>%
#   filter(missing_mech != "MCAR") %>%
#   select(-missing_mech, - rho) %>%
#   select(-beta_MCE, -beta_MCM) %>%
#   kable(digits = 2, escape = F, align = "c") %>%
#   kable_styling(c("striped"), full_width = FALSE)
#   
# fdr_table %>%
#   filter(missing_mech != "MCAR" & p_missing != .3) %>%
#   select(-missing_mech, - rho) %>%
#   select(-beta_MCE) %>%
#   kable(digits = 2, escape = F, align = "c") %>% 
#   kable_styling(c("striped"), full_width = FALSE)
# 
# fnr_table %>%
#   filter(missing_mech != "MCAR" & p_missing != .3) %>%
#   select(-missing_mech, - rho) %>%
#   select(-beta_MCE) %>%
#   kable(digits = 2, escape = F, align = "c") %>%
#   kable_styling(c("striped"), full_width = FALSE)
# 
# rmse_table %>%
#   filter(missing_mech != "MCAR" & p_missing != .3) %>%
#   select(-missing_mech, - rho) %>%
#   select(-beta_MCE, -beta_MCM) %>%
#   kable(digits = 2, escape = F, align = "c") %>% 
#   kable_styling(c("striped"), full_width = FALSE)


## Look at plots of LMMs
set_theme(base =theme_light(), panel.minor.gridcol = "white")

# Fix rmses
for(i in 1:nrow(params)) {
  sims[[i]] <- mutate_at(sims[[i]], vars(starts_with("rmse_")), as.numeric)
  sims[[i]]$run <- paste0(i, "_",sims[[i]]$sim)
}

full <- sims %>% 
  bind_rows()

#### PLOTS ####

simfits <- list()

simfits$rmse <- full %>% 
  filter(corr == "exchangeable") %>% 
  group_by(p, p_missing) %>% 
  select(-starts_with("fnr"), -starts_with("fdr"), - missing_mech) %>% 
  pivot_longer(starts_with("rmse_"), names_to = "Method", names_prefix = "rmse_")  %>% 
  mutate(Method = factor(ifelse(Method == "full", "Full", toupper(Method)))) %>%
  ungroup() %>%
  mutate(p_missing = factor(p_missing), 
         p = factor(p, labels = paste0("p=", unique(p))), 
         SNR = factor(SNR), 
         RMSE = value) %>% 
  lme4::lmer(RMSE ~ Method*p_missing*p*SNR + (1|run), data = .)

 # plot_model(simfits$rmse, type = "int")


#### FDR PLOTS ####

simfits$fdr <- full %>% 
  filter(corr == "exchangeable") %>% 
  group_by(p, p_missing) %>% 
  select(-starts_with("fnr"), -starts_with("rmse"), - missing_mech) %>% 
  pivot_longer(starts_with("fdr_"), names_to = "Method", names_prefix = "fdr_")  %>% 
  mutate(Method = factor(ifelse(Method == "full", "Full", toupper(Method)))) %>%
  ungroup() %>%
  mutate(p_missing = factor(p_missing), 
         p = factor(p, labels = paste0("p=", unique(p))), 
         SNR = factor(SNR), 
         FDR = value*100) %>% 
  lme4::lmer(FDR ~ Method*p_missing*p*SNR + (1|run), data = .)
 
# plot_model(simfits$fdr, type = "int")
 
simfits$fnr <- full %>% 
  filter(corr == "exchangeable") %>% 
  group_by(p, p_missing) %>% 
  select(-starts_with("rmse"), -starts_with("fdr"), - missing_mech) %>% 
  pivot_longer(starts_with("fnr"), names_to = "Method", names_prefix = "fnr_")  %>% 
  mutate(Method = factor(ifelse(Method == "full", "Full", toupper(Method)))) %>%
  ungroup() %>%
  mutate(p_missing = factor(p_missing), 
         p = factor(p, labels = paste0("p=", unique(p))), 
         SNR = factor(SNR), 
         FNR = value*100) %>% 
  lme4::lmer(FNR ~ Method*p_missing*p*SNR + (1|run), data = .)

# plot_model(simfits$fnr, type = "int")

simfits$rmse_boxplot <- full %>% 
  filter(corr == "exchangeable") %>% 
  group_by(p, p_missing) %>% 
  select(-starts_with("fnr"), -starts_with("fdr"), - missing_mech) %>% 
  pivot_longer(starts_with("rmse_"), names_to = "Method", names_prefix = "rmse_")  %>% 
  mutate(Method = factor(ifelse(Method == "full", "Full", toupper(Method)))) %>%
  ungroup() %>%
  mutate(p_missing = factor(p_missing), 
         p = factor(p, labels = paste0("p=", unique(p))), 
         SNR = factor(paste0("SNR=", SNR)), 
         RMSE = value) 

# ggboxplot(simfits$rmse_boxplot, x = "Method", y = "RMSE", facet.by = c("SNR", "p"), color = "p_missing",
#           scales = "free", width = .5, legend.title = "Prop. Missing", outlier.size = .5) +
#   theme_light()
