## Simulation first go
# remotes::install_github("pbreheny/hdrm")


library(tidyverse)
library(pROC)

source("JDS_Submission/milasso.R")
source("JDS_Submission/medcoef.R")
source("JDS_Submission/genData.R")


# Simulation function
simulate_once <- function(n = 158, p = 10, p1 = 3, SNR = .3, corr = "exchangeable", rho = 0.17, 
                          signal = "homogeneous", p_missing = .3, missing_mech = "MCAR",
                          m = 20, ntest = 1500, plot = T, family = "binomial") {
  # n = 158; p = 10; p1 = 3; SNR = .26; cor = "exchangeable"; rho = 0.17;
  # signal = "homogeneous"; p_missing = .1; missing_mech = "MCAR"; m = 20; ntest = 1500; plot = T
  # df <- genData(150, 10, 3, family = "binomial", SNR = 1, corr = "autoregressive", rho = .17, signal = "homogeneous")
  
  df <- genData(n, p, p1 = p1, family = family, SNR = SNR, corr = corr, rho = rho, signal = signal)
  dfm <- mice::ampute(cbind(df$X, df$y), prop = p_missing, mech = missing_mech)$amp
  
  dfi <- mice::mice(dfm, m = m, printFlag = F)
  dfi2 <- lapply(1:m, function(i) mice::complete(dfi, i))
  
  lamvec <- (2^(seq(-1,4,by=0.05)))^2/2
  
  ## select the optimal tuning parameter by BIC
  bicvec = 0
  for (i in 1:length(lamvec)) {
    fit = MI.lasso(dfi2, m, lambda=lamvec[i])
    bicvec[i] = fit$bic
  }
  
  ## select tuning parameter with the smallest BIC
  opt = order(bicvec)[1]
  
  ## fit the model with the optimal lambda
  fit.opt = MI.lasso(dfi2, m, lambda=lamvec[opt])
  
  b_mil <- apply(fit.opt$coefficients,1,mean)
  b_mcl <- medcoeffit(dfi, family = family)
  b_mcm <- medcoeffit(dfi, penalty = "MCP", family = family)
  b_mce <- medcoeffit(dfi, penalty = "lasso", alpha = .25, family = family)
  
  fit_full <- ncvreg::cv.ncvreg(df$X, df$y, penalty = "lasso", family = family)
  b_full <- coef(fit_full)
  
  dfnew <- genData(ntest, p, beta = df$beta, family = family, SNR = SNR, corr = corr, rho = rho, signal = signal)
  
  if(family == "binomial") {
    lodds <- c(cbind(1, dfnew$X) %*% b_mcl)
    probs_mcl <- exp(lodds) / (1 + exp(lodds))
    
    lodds <- c(cbind(1, dfnew$X) %*% b_mcm)
    probs_mcm <- exp(lodds) / (1 + exp(lodds))
    
    lodds <- c(cbind(1, dfnew$X) %*% b_mce)
    probs_mce <- exp(lodds) / (1 + exp(lodds))
    
    lodds <- c(cbind(1, dfnew$X) %*% b_full)
    probs_full <- exp(lodds) / (1 + exp(lodds))
    
    probs_mil <- c(cbind(1, dfnew$X) %*% b_mil)
    
    g <- pROC::roc(dfnew$y ~ probs_mcl, percent = T, ci = T, quiet = T)
    g2 <- pROC::roc(dfnew$y ~ probs_mil, percent = T, ci = T, quiet = T)
    g3 <- pROC::roc(dfnew$y ~ probs_mcm, percent = T, ci = T, quiet = T)
    g4 <- pROC::roc(dfnew$y ~ probs_mce, percent = T, ci = T, quiet = T)
    g5 <- pROC::roc(dfnew$y ~ probs_full, percent = T, ci = T, quiet = T)
    
    if(plot) {
      par(pty="s")
      
      plot(g)    
      plot(g2, add = T, col = 2)
      plot(g3, add = T, col = 3)
      plot(g4, add = T, col = 4)
      plot(g5, add = T, col = 5)
    }
    
    results <- list(
      beta = df$beta, 
      results = data.frame(
        auc_mcl = g$auc, auc_mil = g2$auc, auc_mcm = g3$auc,
        auc_mce = g4$auc, auc_full = g5$auc,
        fnr_mil = mean(b_mil[-1] == 0 & df$beta != 0),
        fdr_mil = mean(b_mil[-1] != 0 & df$beta == 0),
        fnr_mcl = mean(b_mcl[-1] == 0 & df$beta != 0),
        fdr_mcl = mean(b_mcl[-1] != 0 & df$beta == 0),
        fnr_mcm = mean(b_mcm[-1] == 0 & df$beta != 0),
        fdr_mcm = mean(b_mcm[-1] != 0 & df$beta == 0),
        fnr_mce = mean(b_mce[-1] == 0 & df$beta != 0),
        fdr_mce = mean(b_mce[-1] != 0 & df$beta == 0),
        fnr_full = mean(b_full[-1] == 0 & df$beta != 0),
        fdr_full = mean(b_full[-1] != 0 & df$beta == 0)
      )
    )
  } else {
    
    p_mcl <- c(cbind(1, dfnew$X) %*% b_mcl)
    p_mcm <- c(cbind(1, dfnew$X) %*% b_mcm)
    p_mce <- c(cbind(1, dfnew$X) %*% b_mce)
    p_full <- c(cbind(1, dfnew$X) %*% b_full)
    p_mil <- c(cbind(1, dfnew$X) %*% b_mil)
    
    rmse_mcl <- sqrt(mean((p_mcl - dfnew$y)^2))
    rmse_mcm <- sqrt(mean((p_mcm - dfnew$y)^2))
    rmse_mce <- sqrt(mean((p_mce - dfnew$y)^2))
    rmse_full <- sqrt(mean((p_full - dfnew$y)^2))
    rmse_mil <- sqrt(mean((p_mil - dfnew$y)^2))
    
    results <- list(
      beta = df$beta, 
      results = data.frame(
        rmse_mcl = rmse_mcl, 
        rmse_mcm = rmse_mcm, 
        rmse_mce = rmse_mce,
        rmse_full = rmse_full, 
        rmse_mil = rmse_mil,
        fnr_mil = mean(b_mil[-1] == 0 & df$beta != 0),
        fdr_mil = mean(b_mil[-1] != 0 & df$beta == 0),
        fnr_mcl = mean(b_mcl[-1] == 0 & df$beta != 0),
        fdr_mcl = mean(b_mcl[-1] != 0 & df$beta == 0),
        fnr_mcm = mean(b_mcm[-1] == 0 & df$beta != 0),
        fdr_mcm = mean(b_mcm[-1] != 0 & df$beta == 0),
        fnr_mce = mean(b_mce[-1] == 0 & df$beta != 0),
        fdr_mce = mean(b_mce[-1] != 0 & df$beta == 0),
        fnr_full = mean(b_full[-1] == 0 & df$beta != 0),
        fdr_full = mean(b_full[-1] != 0 & df$beta == 0)
      )
    )
  }
  
  results
}

# set.seed(2)
# simulate_once(p_missing = .5)
# simulate_once(missing_mech = "MAR")
# simulate_once(missing_mech = "MNAR")

