#------------------------------------------------------
# 1. Required Libraries
#------------------------------------------------------
# Required Libraries
library(MASS)
library(glmnet)
library(data.table)
library(dplyr)
library(MLmetrics)
library(writexl)
library(mccr)
library(doParallel)
library(foreach)
set.seed(2025)
#------------------------------------------------------
# 2. Sampling Functions (SRS, RSS, ERSS)
#------------------------------------------------------
srs.data <- function(data_pop, n) {
  data_pop[sample(nrow(data_pop), n), ]
}

RSS.data <- function(data_pop, n, k) {
  m <- n / k
  RSS <- vector("list", m * k)
  for (j in 1:m) {
    for (i in 1:k) {
      A <- data_pop[sample(nrow(data_pop), k, replace = TRUE), ]
      RSS[[(j - 1) * k + i]] <- A[order(A$HighCorrPredictor)[i], ]
    }
  }
  return(rbindlist(RSS))
}

ERSS.data <- function(data_pop, n, k) {
  m <- n / k
  u <- k %/% 2
  ERSS <- vector("list", m * k)
  idx <- 1
  for (j in 1:m) {
    for (i in 1:u) {
      A <- data_pop[sample(nrow(data_pop), k, replace = TRUE), ]
      ERSS[[idx]] <- A[order(A$HighCorrPredictor)[1], ]; idx <- idx + 1
    }
    for (i in 1:u) {
      A <- data_pop[sample(nrow(data_pop), k, replace = TRUE), ]
      ERSS[[idx]] <- A[order(A$HighCorrPredictor)[k], ]; idx <- idx + 1
    }
    if (k %% 2 == 1) {
      A <- data_pop[sample(nrow(data_pop), k, replace = TRUE), ]
      ERSS[[idx]] <- A[order(A$HighCorrPredictor)[(k %/% 2 + 1)], ]; idx <- idx + 1
    }
  }
  return(rbindlist(ERSS))
}

#------------------------------------------------------
# 3. Population Generator Function 
#------------------------------------------------------
pop.generator <- function(n = 10000, intercept = -5, linearVars = 1, noiseVars = 1,
                          corrVars = 2, corrType = "AR1", corrValue = 0.2,
                          mislabel = 0, highCorrCoeff = 5, imbalanceRatio = 0.5) {
  sigma <- matrix(c(2, 1.3, 1.3, 2), 2, 2)
  tmpData <- data.frame(mvrnorm(n = n, c(0, 0), sigma))
  names(tmpData) <- c("TwoFactor1", "TwoFactor2")
  
  lin_data <- matrix(rnorm(n * linearVars), ncol = linearVars)
  colnames(lin_data) <- paste0("Linear", sprintf("%02d", 1:linearVars))
  tmpData <- cbind(tmpData, lin_data)
  
  tmpData$Nonlinear1 <- runif(n, -1, 1)
  tmpData$Nonlinear2 <- runif(n)
  tmpData$Nonlinear3 <- runif(n)
  
  noise_data <- matrix(rnorm(n * noiseVars), ncol = noiseVars)
  colnames(noise_data) <- paste0("Noise", sprintf("%02d", 1:noiseVars))
  tmpData <- cbind(tmpData, noise_data)
  
  if (corrVars > 0) {
    if (corrType == "exch") {
      vc <- matrix(corrValue, ncol = corrVars, nrow = corrVars); diag(vc) <- 1
    } else {
      vc <- toeplitz(corrValue^(0:(corrVars - 1)))
    }
    corr_data <- mvrnorm(n, mu = rep(0, corrVars), Sigma = vc)
    colnames(corr_data) <- paste0("Corr", sprintf("%02d", 1:corrVars))
    tmpData <- cbind(tmpData, corr_data)
  }
  
  tmpData$HighCorrPredictor <- rnorm(n)
  
  lp <- intercept - 4 * tmpData$TwoFactor1 + 4 * tmpData$TwoFactor2 +
    2 * tmpData$TwoFactor1 * tmpData$TwoFactor2 +
    tmpData$Nonlinear1^3 + 2 * exp(-6 * (tmpData$Nonlinear1 - 0.3)^2) +
    2 * sin(pi * tmpData$Nonlinear2 * tmpData$Nonlinear3) +
    highCorrCoeff * tmpData$HighCorrPredictor
  
  lin_coefs <- seq(10, 1, length = linearVars) / 4 * rep(c(-1, 1), length.out = linearVars)
  for (i in seq_along(lin_coefs)) {
    lp <- lp + tmpData[[paste0("Linear", sprintf("%02d", i))]] * lin_coefs[i]
  }
  
  prob <- binomial()$linkinv(lp + rnorm(n, sd = 0))
  
  if (mislabel > 0 & mislabel < 1) {
    flip <- sample(n, floor(n * mislabel)); prob[flip] <- 1 - prob[flip]
  }
  
  threshold <- quantile(prob, probs = imbalanceRatio)
  tmpData$y <- as.integer(prob > threshold)
  
  tmpData <- tmpData %>% relocate(y)
  return(tmpData)
}

#------------------------------------------------------
# 4. Define Metrics Function 
#------------------------------------------------------
calculate_metrics_lowD <- function(train, test) {
  y_train <- train$y
  y_test  <- test$y
  
  # Sanity check: both classes must be present
  if (any(table(y_train) < 2) || any(table(y_test) < 2)) return(NULL)
  
  # Fit logistic regression on remaining variables
  model_log <- tryCatch(glm(y ~ ., data = train, family = "binomial"), error = function(e) NULL)
  
  if (is.null(model_log)) return(NULL)
  
  prob_log <- predict(model_log, newdata = test, type = "response")
  class_log <- ifelse(prob_log > 0.5, 1, 0)
  
  # confusion matrix components
  TP <- sum(class_log == 1 & y_test == 1)
  TN <- sum(class_log == 0 & y_test == 0)
  FP <- sum(class_log == 1 & y_test == 0)
  FN <- sum(class_log == 0 & y_test == 1)
  
  sensitivity <- if ((TP + FN) > 0) TP / (TP + FN) else NA
  specificity <- if ((TN + FP) > 0) TN / (TN + FP) else NA
  gmean <- if (!is.na(sensitivity) && !is.na(specificity)) sqrt(sensitivity * specificity) else NA
  
  
  # Wrapper to safely calculate metrics that might error out
  safe_metric <- function(expr) {
    tryCatch(expr, error = function(e) NA)
  }
  
  # Check if prob_log has variation (needed for AUC etc.)
  has_variation <- length(unique(prob_log)) > 1 && !any(is.na(prob_log))
  
  result_df <- data.frame(
    accuracy = safe_metric(Accuracy(class_log, y_test)),
    precision = safe_metric(Precision(y_pred = class_log, y_true = y_test, positive = "1")),
    recall = safe_metric(Recall(y_pred = class_log, y_true = y_test, positive = "1")),
    f1 = safe_metric(F1_Score(y_pred = class_log, y_true = y_test, positive = "1")),
    auc = if (has_variation) safe_metric(AUC(prob_log, y_test)) else NA,
    pr_auc = if (has_variation) safe_metric(PRAUC(prob_log, y_test)) else NA,
    mcc = safe_metric(mccr::mccr(class_log, y_test)),
    logloss = if (has_variation) safe_metric(LogLoss(prob_log, y_test)) else NA,
    specificity = specificity,
    gmean = gmean
  )
  
  return(result_df)
}

#------------------------------------------------------
# 5. Simulation 
#------------------------------------------------------
n_pop <- 5000
n_values <- c( 60, 120, 210, 300)
k_values <- c(3, 5, 10)
iterations <- 1000

coeff_0.3_LD <- data.frame(
  imbalanceRatio = c(0.5, 0.6, 0.7, 0.8),
  mean_highCorrCoeff = c(2.72, 2.73, 2.91, 3.41)
)

coeff_0.5_LD <- data.frame(
  imbalanceRatio = c(0.5, 0.6, 0.7, 0.8),
  mean_highCorrCoeff = c(5.38, 5.52, 6.13, 7.60)
)

coeff_0.7_LD <- data.frame(
  imbalanceRatio = c(0.5, 0.6, 0.7, 0.8),
  mean_highCorrCoeff = c(12.70 , 13.68 , 17.39 , 28.89)
)

coeff_baselines <- list("0.3" = coeff_0.3_LD,"0.5" = coeff_0.5_LD,"0.7" = coeff_0.7_LD)

imbalanceRatios <- c(0.5, 0.6, 0.7, 0.8)

sampling_methods <- c("SRS", "RSS", "ERSS") 


start_time <- Sys.time()

simulation_results <- list()
param_results <- list()  

for (coeff_baseline_name in names(coeff_baselines)) {
  
  coeff_df <- coeff_baselines[[coeff_baseline_name]]
  
  for (imbalanceRatio in imbalanceRatios) {
    highCorrCoeff <- coeff_df$mean_highCorrCoeff[coeff_df$imbalanceRatio == imbalanceRatio]
    
    data_pop <- pop.generator(n = n_pop,highCorrCoeff = highCorrCoeff,imbalanceRatio = imbalanceRatio)
    
    true_beta <- coef(glm(y ~ ., data = data_pop, family = "binomial", offset = NULL))
    
    for (n in n_values) {
      
      for (i in 1:iterations) {
        
        SRS_train <- srs.data(data_pop, n)
        SRS_test  <- srs.data(data_pop, n / 2)
        
        for (k in k_values) {
          cat(paste0("Running for n = ", n, ", k = ", k, ", iter = ", i, "\n"))
          
          for (method in sampling_methods) {
            
            #-----------------------------------------
            # Step 1: Create training/testing data
            #-----------------------------------------
            if (method == "SRS") {
              train <- SRS_train
              test  <- SRS_test
            } else {
              train <- switch(method,
                              RSS = RSS.data(data_pop, n, k),
                              ERSS = ERSS.data(data_pop, n, k)
                              )
              
              test <- switch(method,
                             RSS = RSS.data(data_pop, n / 2, k),
                             ERSS = ERSS.data(data_pop, n / 2, k)
                             )
            }
            
            if (length(unique(train$y)) < 2 || length(unique(test$y)) < 2) {
              next
            }
            # -------------------------------
            # Fit logistic regression
            # -------------------------------
            model_log <- tryCatch(
              glm(y ~ ., data = train, family = "binomial"),
              error = function(e) NULL
            )
            
            if (!is.null(model_log) && all(!is.na(coef(model_log)))) {
              
              est <- coef(model_log)
              se  <- coef(summary(model_log))[, "Std. Error"]
              
              low <- est - 1.96 * se
              up  <- est + 1.96 * se
              
              OR_est <- exp(est)
              OR_low <- exp(low)
              OR_up  <- exp(up)
              
              true_beta_matched <- true_beta[names(est)]
              TrueOR <- exp(true_beta_matched)
              
              Beta_Bias <- est - true_beta_matched       
              OR_Bias   <- OR_est - TrueOR               
              
              Coverage <- (TrueOR >= OR_low) & (TrueOR <= OR_up)
              
              coefs_df <- data.frame(
                Parameter      = names(est),
                TrueBeta       = true_beta_matched,
                Estimate       = est,
                Beta_Bias      = Beta_Bias,
                Std.Error      = se,
                Lower95_logit  = low,
                Upper95_logit  = up,
                TrueOR         = TrueOR,
                OR             = OR_est,
                OR_Lower95     = OR_low,
                OR_Upper95     = OR_up,
                OR_Bias        = OR_Bias,
                Coverage       = Coverage,
                SamplingMethod = method,
                n              = n,
                k              = k,
                imbalanceRatio = imbalanceRatio,
                CoeffBaseline  = coeff_baseline_name,
                iter           = i,
                Failure        = FALSE,
                stringsAsFactors = FALSE
              )
              
              # Save
              param_results[[length(param_results) + 1]] <- coefs_df
              
            } else {
              
              param_results[[length(param_results) + 1]] <- data.frame(
                Parameter      = names(true_beta),
                TrueBeta       = true_beta,
                Estimate       = NA,
                Beta_Bias      = NA,
                Std.Error      = NA,
                Lower95_logit  = NA,
                Upper95_logit  = NA,
                TrueOR         = exp(true_beta),
                OR             = NA,
                OR_Lower95     = NA,
                OR_Upper95     = NA,
                OR_Bias        = NA,
                Coverage       = NA,
                SamplingMethod = method,
                n              = n,
                k              = k,
                imbalanceRatio = imbalanceRatio,
                CoeffBaseline  = coeff_baseline_name,
                iter           = i,
                Failure        = TRUE,
                stringsAsFactors = FALSE
              )
            }
            
            #-----------------------------------------
            # Step 2: Evaluate model performance metrics
            #-----------------------------------------
            t0 <- Sys.time()
            result_metrics <- calculate_metrics_lowD(train, test)
            t1 <- Sys.time()
            elapsed_time <- as.numeric(difftime(t1, t0, units = "secs"))
            
            if (!is.null(result_metrics)) {
              result_metrics <- result_metrics %>%
                mutate(
                  SamplingMethod = method,
                  n = n, k = k, iter = i,
                  imbalanceRatio = imbalanceRatio,
                  CoeffBaseline = coeff_baseline_name,
                  Failure = FALSE,
                  elapsed_time = elapsed_time
                )
              
              simulation_results[[length(simulation_results) + 1]] <- result_metrics
            }
            
            
          }  
        }  
      }  
    }  
  }  
}  


results_df <- rbindlist(simulation_results, fill = TRUE)
param_df   <- rbindlist(param_results, fill = TRUE)

saveRDS(results_df, "all_iterations_results.rds")
saveRDS(param_df, "all_iterations_logit_parames.rds")