plot_functions.R

library(caret)
library(pROC)
library(parabar)
library(dplyr)
library(ggplot2)
library(clusterGeneration)
library(glmnet)

source("./recommend_sample_size_ML.r")

# Compute R^2
compute_population_R2 <- function(population) {
  # This function takes as input the output of generate_continious_mega_population
  
  pred_names <- setdiff(names(population$data), "y")
  X <- cbind(Intercept = 1, as.matrix(population$data[, pred_names]))
  beta_vec <- as.numeric(population$true_coefficients)
  y <- population$data$y
  y_hat <- drop(X %*% beta_vec)
  
  sst <- sum((y - mean(y))^2)
  sse <- sum((y - y_hat)^2)
  
  r2 <- if (sst == 0) NA_real_ else 1 - sse/sst
  
  return(r2)
}

plot_r2_diff <- function(pop_size=10000, num_reps=100) {
  diff <- numeric(num_reps)
  
  for (i in seq_len(num_reps)) {
    curr_n_vars <- sample(10:100, 1)
    curr_dependence <- runif(1, min = 0, max = 2)
    curr_p_dependent <- runif(1, min = 0, max = 1)
    sparsity <- runif(1)
    curr_coef_mean <- runif(1, -0.5, 0.5)
    curr_coef_sd   <- exp(rnorm(1, log(0.5), 0.5))
    curr_R2        <- 0.05 + 0.9 * rbeta(1, 2, 2)
    
    population <- generate_continious_mega_population(
      n_obs = pop_size, 
      n_vars = curr_n_vars,
      dependence = curr_dependence,
      p_dependent = curr_p_dependent, 
      sparsity = sparsity,
      coef_mean = curr_coef_mean, 
      coef_sd = curr_coef_sd,
      R2 = curr_R2)
    
    calculated_R2 <- compute_population_R2(population)
    
    curr_diff <- calculated_R2 - curr_R2
    
    diff[i] <- abs(curr_diff)
  }
  
  # plot histogram
  df <- data.frame(diff = diff)
  p <- ggplot2::ggplot(df, ggplot2::aes(x = diff)) +
    ggplot2::geom_histogram(bins = 40, alpha = 0.8) +
    ggplot2::geom_vline(xintercept = 0, linetype = "dashed") +
    ggplot2::labs(
      title = sprintf("Difference in R^2", num_reps),
      x = "R^2 absolute deviation",
      y = "Count"
    ) +
    ggplot2::theme_minimal()
  
  print(p)
}

# plot_r2_diff(num_reps=500)

# Function that simulates multiple runs to compare the variance of the simulation with different repetitions
obtain_variance_per_repetitions <- function(
    num_random_settings = 5,
    reps_per_random_settings = 30,
    repetition_vector = c(10, 20, 30, 40, 50, 60, 70, 80, 90, 100),
    plot_path = "variance_by_repetitions.png",
    results_dir = "variance_results"
) {
  if (!dir.exists(results_dir)) dir.create(results_dir, recursive = TRUE)
  
  print("Start repetition test simulation.")
  print(paste0("Number of random settings rounds: ", num_random_settings))
  print(paste0("Number of repetitions per random setting: ", reps_per_random_settings))
  
  # Define variables that will not be randomized
  curr_n_obs <- 100000
  curr_model_type <- "OLS"
  curr_metric <- "RMSE"
  curr_start_size <- 500
  curr_end_size <- 1500
  curr_step_size <- 10
  
  # Results table
  results <- data.frame(
    random_setting = integer(),
    repetitions = integer(),
    run = integer(),
    conv_value = numeric(),
    stringsAsFactors = FALSE
  )
  
  for (i in seq_len(num_random_settings)) {
    print("==================================================")
    print(sprintf("This is random-setting round %s out of %s", i, num_random_settings))
    
    # Randomize some of the variables
    curr_n_vars <- sample(10:100, 1)
    curr_dependence <- runif(1, min = 0, max = 2)
    curr_p_dependent <- runif(1, min = 0, max = 1)
    curr_sparsity <- runif(1)
    curr_coef_mean <- runif(1, -0.5, 0.5)
    curr_coef_sd <- exp(rnorm(1, log(0.5), 0.5))
    curr_R2 <- 0.05 + 0.9 * rbeta(1, 2, 2)
    
    # Generate the mega population and train the population model
    print(paste0("Generating mega-population with ", curr_n_obs, " observations and ", curr_n_vars, " variables."))
    
    mega_population <- generate_continious_mega_population(
      n_obs = curr_n_obs,
      n_vars = curr_n_vars,
      dependence = curr_dependence,
      p_dependent = curr_p_dependent,
      sparsity = curr_sparsity,
      coef_mean = curr_coef_mean,
      coef_sd = curr_coef_sd,
      R2 = curr_R2
    )
    
    # Split the data in training and test sets
    train_index <- createDataPartition(mega_population$data$y, p = 0.8, list = FALSE)
    pop_train_data <- mega_population$data[train_index, ]
    pop_test_data  <- mega_population$data[-train_index, ]
    
    print(paste0("Fitting population model: ", curr_model_type, " regression on the mega-population."))
    
    # Fit population model and get the population value
    if (curr_model_type == "OLS")
      population_model <- ols_population(train_data = pop_train_data, test_data = pop_test_data)
    if (curr_model_type == "Ridge")
      population_model <- ridge_population(train_data = pop_train_data, test_data = pop_test_data)
    if (curr_model_type == "LASSO")
      population_model <- lasso_population(train_data = pop_train_data, test_data = pop_test_data)
    
    # Print the population values
    print(paste0("Population rmse: ", population_model$rmse))
    print(paste0("Population mae: ", population_model$mae))
    print(paste0("Population r squared: ", population_model$r2))
    print(paste0("Population best lambda: ", population_model$best_lambda))
    
    population_value <- population_model$rmse
    if (curr_metric == "MAE") population_value <- population_model$mae
    if (curr_metric == "Rsquared") population_value <- population_model$r2
    
    # Choose task function
    task_fn <- switch(curr_model_type,
                      "OLS"   = ordinary_least_squares,
                      "Ridge" = ridge_regression,
                      "LASSO" = lasso_regression,
                      stop("Unknown model type: ", curr_model_type)
    )
    
    # Loop over repetition counts for this random settings
    for (curr_num_repetitions in repetition_vector) {
      curr_start_size <- 800
      curr_end_size   <- 1200
      
      for (j in seq_len(reps_per_random_settings)) {
        print(sprintf("Running simulation: repetitions=%s | run %s out of %s", curr_num_repetitions, j, reps_per_random_settings))
        
        curr_recommendation <- find_convergence_sample_size(
          mega_population$data,
          task = function(sz) {
            res <- task_fn(mega_population$data, sample_size = sz, metric = curr_metric)
            if (curr_metric == "RMSE") res$rmse else if (curr_metric == "MAE") res$mae else res$r2
          },
          pop_value = population_value,
          start_size = curr_start_size,
          end_size = curr_end_size,
          step_size = curr_step_size,
          repetitions = curr_num_repetitions
        )
        
        # Extract the converged value
        conv <- if (!is.null(curr_recommendation)) curr_recommendation$converged_at_n else NULL
        rec_num <- if (is.null(conv)) NA_real_ else as.numeric(conv)
        
        # Save each run as a row
        results <- rbind(
          results,
          data.frame(
            random_setting = i,
            repetitions    = curr_num_repetitions,
            run            = j,
            conv_value     = rec_num,
            stringsAsFactors = FALSE
          )
        )
        
        print(paste0("Converged at: ", conv))
        
        # Move the minimal and maximal sample size
        # If the function converges at the minimal sample size it is possible that it will converge at a smaller sample size
        # If it does not converge, it is possible that a convergence might be achieved at a larger sample size
        if (is.null(conv)) {
          curr_start_size <- max(100, curr_start_size + 300)
          curr_end_size   <- min(curr_n_obs - 1000, curr_end_size + 300)
        } else {
          if (!is.na(rec_num) && rec_num == curr_start_size) {
            curr_start_size <- max(100, curr_start_size - 300)
            curr_end_size   <- min(curr_n_obs - 1000, curr_end_size - 300)
            
            if (curr_end_size > curr_start_size) {
              curr_end_size <- curr_start_size + 400
            }
          }
        }
      }
      
      # Print the variance
      x <- results$conv_value[
        results$random_setting == i &
          results$repetitions == curr_num_repetitions &
          !is.na(results$conv_value)
      ]
      curr_variance <- if (length(x) >= 2) var(x) else NA_real_
      
      print(sprintf(
        "Variance (monitor only) for repetitions=%s (random-setting round %s): %s",
        curr_num_repetitions, i, curr_variance
      ))
    }
    
    # Save the results
    round_tag <- sprintf("round_%02d", i)
    
    round_results <- results[results$random_setting == i, , drop = FALSE]
    utils::write.csv(round_results, file = file.path(results_dir, paste0("results_", round_tag, ".csv")), row.names = FALSE)
    
    print(paste0("Saved results after random-setting round ", i, " to folder: ", results_dir))
  }
  
  # Calculate summary tables
  variance_tbl <- aggregate(
    conv_value ~ random_setting + repetitions,
    data = results,
    FUN = function(v) {
      v <- v[!is.na(v)]
      if (length(v) >= 2) var(v) else NA_real_
    }
  )
  names(variance_tbl)[names(variance_tbl) == "conv_value"] <- "variance"
  
  summary_tbl <- aggregate(variance ~ repetitions, data = variance_tbl, FUN = mean, na.rm = TRUE)
  names(summary_tbl)[names(summary_tbl) == "variance"] <- "mean_variance"
  
  # Final save of full results and summary
  utils::write.csv(results, file = file.path(results_dir, "final_results_raw.csv"), row.names = FALSE)
  utils::write.csv(variance_tbl, file = file.path(results_dir, "final_results_variance_by_setting.csv"), row.names = FALSE)
  utils::write.csv(summary_tbl, file = file.path(results_dir, "final_results_summary.csv"), row.names = FALSE)
  
  invisible(list(raw = results, variance_by_setting = variance_tbl, summary = summary_tbl))
}

# Run it
# obtain_variance_per_repetitions(num_random_settings = 3, reps_per_random_settings = 20, repetition_vector = c(10, 20, 30, 40, 50))

obtain_variance_per_consecutive_n <- function(
    num_random_settings = 5,
    reps_per_random_settings = 30,
    consecutive_n_vector = c(2, 3, 4, 5, 6, 7, 8, 9, 10),
    repetitions_inside_search = 10,
    plot_path = "variance_by_consecutive_n.png",
    results_dir = "variance_results_consecutive_n"
) {
  if (!dir.exists(results_dir)) dir.create(results_dir, recursive = TRUE)
  
  print("Start consecutive_n test simulation.")
  print(paste0("Number of random settings rounds: ", num_random_settings))
  print(paste0("Number of runs per random setting: ", reps_per_random_settings))
  print(paste0("Fixed repetitions inside convergence search: ", repetitions_inside_search))
  
  # Fixed variables
  curr_n_obs <- 100000
  curr_model_type <- "OLS"
  curr_metric <- "RMSE"
  curr_start_size <- 500
  curr_end_size <- 1500
  curr_step_size <- 10
  
  # Results table
  results <- data.frame(
    random_setting = integer(),
    consecutive_n = integer(),
    run = integer(),
    conv_value = numeric(),
    stringsAsFactors = FALSE
  )
  
  for (i in seq_len(num_random_settings)) {
    print("==================================================")
    print(sprintf("This is random-setting round %s out of %s", i, num_random_settings))
    
    # Randomize population parameters
    curr_n_vars <- sample(10:100, 1)
    curr_dependence <- runif(1, min = 0, max = 2)
    curr_p_dependent <- runif(1, min = 0, max = 1)
    curr_sparsity <- runif(1)
    curr_coef_mean <- runif(1, -0.5, 0.5)
    curr_coef_sd <- exp(rnorm(1, log(0.5), 0.5))
    curr_R2 <- 0.05 + 0.9 * rbeta(1, 2, 2)
    
    # Generate the mega population and train the population model
    print(paste0("Generating mega-population with ", curr_n_obs," observations and ", curr_n_vars, " variables."))
    
    mega_population <- generate_continious_mega_population(
      n_obs = curr_n_obs,
      n_vars = curr_n_vars,
      dependence = curr_dependence,
      p_dependent = curr_p_dependent,
      sparsity = curr_sparsity,
      coef_mean = curr_coef_mean,
      coef_sd = curr_coef_sd,
      R2 = curr_R2
    )
    
    # Split the data in training and test sets
    train_index <- createDataPartition(mega_population$data$y, p = 0.8, list = FALSE)
    pop_train_data <- mega_population$data[train_index, ]
    pop_test_data  <- mega_population$data[-train_index, ]
    
    print(paste0("Fitting population model: ", curr_model_type, " regression on the mega-population."))
    
    # Fit population model and get the population value
    if (curr_model_type == "OLS")
      population_model <- ols_population(train_data = pop_train_data, test_data = pop_test_data)
    if (curr_model_type == "Ridge")
      population_model <- ridge_population(train_data = pop_train_data, test_data = pop_test_data)
    if (curr_model_type == "LASSO")
      population_model <- lasso_population(train_data = pop_train_data, test_data = pop_test_data)
    
    # Print the population values
    print(paste0("Population rmse: ", population_model$rmse))
    print(paste0("Population mae: ", population_model$mae))
    print(paste0("Population r squared: ", population_model$r2))
    print(paste0("Population best lambda: ", population_model$best_lambda))
    
    population_value <- population_model$rmse
    if (curr_metric == "MAE") population_value <- population_model$mae
    if (curr_metric == "Rsquared") population_value <- population_model$r2
    
    # Choose task function
    task_fn <- switch(curr_model_type,
                      "OLS"   = ordinary_least_squares,
                      "Ridge" = ridge_regression,
                      "LASSO" = lasso_regression,
                      stop("Unknown model type: ", curr_model_type))
    
    # Loop over consecutive_n values for this random settings
    for (curr_consecutive_n in consecutive_n_vector) {
      curr_start_size <- 800
      curr_end_size   <- 1200
      
      for (j in seq_len(reps_per_random_settings)) {
        print(sprintf("Running simulation: consecutive_n=%s | run %s out of %s", curr_consecutive_n, j, reps_per_random_settings))
        
        curr_recommendation <- find_convergence_sample_size(
          mega_population$data,
          task = function(sz) {
            res <- task_fn(mega_population$data, sample_size = sz, metric = curr_metric)
            if (curr_metric == "RMSE") res$rmse else if (curr_metric == "MAE") res$mae else res$r2
          },
          pop_value = population_value,
          start_size = curr_start_size,
          end_size = curr_end_size,
          step_size = curr_step_size,
          repetitions = repetitions_inside_search,
          consecutive_n = curr_consecutive_n
        )
        
        # Extract the converged value
        conv <- if (!is.null(curr_recommendation)) curr_recommendation$converged_at_n else NULL
        rec_num <- if (is.null(conv)) NA_real_ else as.numeric(conv)
        
        # Save each row
        results <- rbind(
          results,
          data.frame(
            random_setting = i,
            consecutive_n  = curr_consecutive_n,
            run            = j,
            conv_value     = rec_num,
            stringsAsFactors = FALSE
          )
        )
        
        print(paste0("Converged at: ", conv))
        
        # Move the minimal and maximal sample size
        # If the function converges at the minimal sample size it is possible that it will converge at a smaller sample size
        # If it does not converge, it is possible that a convergence might be achieved at a larger sample size
        if (is.null(conv)) {
          curr_start_size <- max(100, curr_start_size + 300)
          curr_end_size   <- min(curr_n_obs - 1000, curr_end_size + 300)
        } else {
          if (!is.na(rec_num) && rec_num == curr_start_size) {
            curr_start_size <- max(100, curr_start_size - 300)
            curr_end_size   <- min(curr_n_obs - 1000, curr_end_size - 300)
            
            if (curr_end_size > curr_start_size) {
              curr_end_size <- curr_start_size + 400
            }
          }
        }
      }
      
      x <- results$conv_value[
        results$random_setting == i &
          results$consecutive_n == curr_consecutive_n &
          !is.na(results$conv_value)
      ]
      curr_variance <- if (length(x) >= 2) var(x) else NA_real_
      
      print(sprintf("Variance (monitor only) for consecutive_n=%s (random-setting round %s): %s", curr_consecutive_n, i, curr_variance))
    }
    
    # Save the results
    round_tag <- sprintf("round_%02d", i)
    round_results <- results[results$random_setting == i, , drop = FALSE]
    
    utils::write.csv(
      round_results,
      file = file.path(results_dir, paste0("results_", round_tag, ".csv")),
      row.names = FALSE)
    
    print(paste0("Saved results after random-setting round ", i, " to folder: ", results_dir))
  }
  
  # Compute variance
  variance_tbl <- aggregate(
    conv_value ~ random_setting + consecutive_n,
    data = results,
    FUN = function(v) {
      v <- v[!is.na(v)]
      if (length(v) >= 2) var(v) else NA_real_
    }
  )
  names(variance_tbl)[names(variance_tbl) == "conv_value"] <- "variance"
  
  # Mean variance across settings for each consecutive_n
  summary_tbl <- aggregate(variance ~ consecutive_n, data = variance_tbl, FUN = mean, na.rm = TRUE)
  names(summary_tbl)[names(summary_tbl) == "variance"] <- "mean_variance"
  
  # Final save
  utils::write.csv(results, file = file.path(results_dir, "final_results_raw.csv"), row.names = FALSE)
  utils::write.csv(variance_tbl, file = file.path(results_dir, "final_results_variance_by_setting.csv"), row.names = FALSE)
  utils::write.csv(summary_tbl, file = file.path(results_dir, "final_results_summary.csv"), row.names = FALSE)
  
  invisible(list(raw = results, variance_by_setting = variance_tbl, summary = summary_tbl, plot = p))
}

# Run it
# obtain_variance_per_consecutive_n(
#   num_random_settings = 3,
#   reps_per_random_settings = 20,
#   consecutive_n_vector = c(3, 5, 10, 15, 20),
#   repetitions_inside_search = 10
# )


obtain_mean_sample_size_per_model <- function(
    num_populations = 5,
    n_runs_per_model = 20,
    n_obs = 100000,
    metric = "RMSE", # "RMSE", "MAE" or "Rsquared"
    start_size = 100,
    end_size = 800,
    step_size = 10,
    repetitions_inside_search = 10,
    convergence_threshold = 0.1,
    consecutive_n = 5,
    n_cores = 7,
    plot_path = "mean_sample_size_by_model.png",
    results_dir = "mean_sample_size_results"
) {
  if (!dir.exists(results_dir)) dir.create(results_dir, recursive = TRUE)
  
  models <- c("OLS", "Ridge", "LASSO")
  
  # Results table
  results <- data.frame(
    population_id = integer(),
    model_type    = character(),
    run           = integer(),
    conv_value    = numeric(),
    stringsAsFactors = FALSE
  )
  
  pick_metric <- function(res, metric) {
    if (metric == "RMSE") res$rmse else if (metric == "MAE") res$mae else res$r2
  }
  
  pop_fit_fn <- list(
    OLS   = ols_population,
    Ridge = ridge_population,
    LASSO = lasso_population
  )
  
  task_fn <- list(
    OLS   = ordinary_least_squares,
    Ridge = ridge_regression,
    LASSO = lasso_regression
  )
  
  print("Start model mean sample size simulation.")
  print(paste0("Populations: ", num_populations))
  print(paste0("Runs per model per population: ", n_runs_per_model))
  print(paste0("Repetitions inside convergence search: ", repetitions_inside_search))
  
  for (p in seq_len(num_populations)) {
    print("==================================================")
    print(sprintf("Population %s / %s", p, num_populations))
    
    # Random population parameters
    curr_n_vars      <- sample.int(100, 1)
    curr_dependence  <- runif(1, 0, 2)
    curr_p_dependent <- runif(1, 0, 1)
    curr_sparsity    <- runif(1)
    curr_coef_mean   <- runif(1, -0.5, 0.5)
    curr_coef_sd     <- exp(rnorm(1, log(0.5), 0.5))
    curr_R2          <- 0.05 + 0.9 * rbeta(1, 2, 2)
    
    # Generate mega population
    mega_population <- generate_continious_mega_population(
      n_obs = n_obs,
      n_vars = curr_n_vars,
      dependence = curr_dependence,
      p_dependent = curr_p_dependent,
      sparsity = curr_sparsity,
      coef_mean = curr_coef_mean,
      coef_sd = curr_coef_sd,
      R2 = curr_R2
    )
    
    train_index <- createDataPartition(mega_population$data$y, p = 0.8, list = FALSE)
    pop_train_data <- mega_population$data[train_index, ]
    pop_test_data  <- mega_population$data[-train_index, ]
    
    population_values <- list()
    
    for (m in models) {
      print(paste0("Fitting population model for: ", m))
      fit <- pop_fit_fn[[m]](train_data = pop_train_data, test_data = pop_test_data)
      population_values[[m]] <- pick_metric(fit, metric)
      
      print(paste0("Population ", metric, " (", m, "): ", population_values[[m]]))
      if (!is.null(fit$best_lambda)) {
        print(paste0("Population best lambda (", m, "): ", fit$best_lambda))
      }
    }
    
    for (m in models) {
      print("--------------------------------------------------")
      print(paste0("Now running convergence searches for model: ", m))
      
      curr_start_size <- start_size
      curr_end_size   <- end_size
      
      for (j in seq_len(n_runs_per_model)) {
        print(sprintf("Model=%s | run %s / %s", m, j, n_runs_per_model))
        
        rec <- find_convergence_sample_size(
          mega_population$data,
          task = function(sz) {
            res <- task_fn[[m]](mega_population$data, sample_size = sz, metric = metric)
            pick_metric(res, metric)
          },
          pop_value = population_values[[m]],
          start_size = curr_start_size,
          end_size = curr_end_size,
          step_size = step_size,
          repetitions = repetitions_inside_search,
          convergence_threshold = convergence_threshold,
          consecutive_n = consecutive_n,
          n_cores = n_cores
        )
        
        conv <- if (!is.null(rec)) rec$converged_at_n else NULL
        conv_num <- if (is.null(conv)) NA_real_ else as.numeric(conv)
        
        results <- rbind(
          results,
          data.frame(
            population_id = p,
            model_type    = m,
            run           = j,
            conv_value    = conv_num,
            stringsAsFactors = FALSE
          )
        )
        
        print(paste0("Converged at: ", conv))
        
        # Move the minimal and maximal sample size
        # If the function converges at the minimal sample size it is possible that it will converge at a smaller sample size
        # If it does not converge, it is possible that a convergence might be achieved at a larger sample size
        if (is.null(conv)) {
          curr_start_size <- max(100, curr_start_size + 300)
          curr_end_size   <- min(n_obs - 1000, curr_end_size + 300)
        } else {
          if (!is.na(conv_num) && conv_num == curr_start_size) {
            curr_start_size <- max(100, curr_start_size - 300)
            curr_end_size   <- min(n_obs - 1000, curr_end_size - 300)
            if (curr_end_size > curr_start_size) curr_end_size <- curr_start_size + 400
          }
        }
      }
    }
    
    # Save after each population
    pop_tag <- sprintf("population_%02d", p)
    pop_results <- results[results$population_id == p, , drop = FALSE]
    
    utils::write.csv(pop_results, file = file.path(results_dir, paste0("results_", pop_tag, ".csv")), row.names = FALSE)
    print(paste0("Saved results after population ", p, " to folder: ", results_dir))
  }
  
  # Calculate summary tables
  summary_tbl <- aggregate(
    conv_value ~ model_type,
    data = results,
    FUN = function(x) mean(x, na.rm = TRUE)
  )
  names(summary_tbl)[names(summary_tbl) == "conv_value"] <- "mean_recommended_n"
  
  # Final save
  utils::write.csv(results, file = file.path(results_dir, "final_results_raw.csv"), row.names = FALSE)
  utils::write.csv(summary_tbl, file = file.path(results_dir, "final_results_summary.csv"), row.names = FALSE)
  
  invisible(list(raw = results, summary = summary_tbl))
}

# obtain_mean_sample_size_per_model(num_populations = 5, n_runs_per_model = 5)
plot_results <- function(
    results_dir,
    file = "final_results_summary.csv",
    plot_path = NULL,
    title = NULL,
    x_label = NULL,
    y_label = NULL,
    value_label = NULL,
    sort_bars = FALSE,
    descending = TRUE,
    digits = 2,
    rotate_x = TRUE,
    width = 9,
    height = 5,
    dpi = 150,
    convert_variance_to_sd = FALSE,          # NEW
    variance_colname_pattern = "(?i)variance" # NEW (optional safety): only auto-rename y if it looks like variance
) {
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("ggplot2 not installed; install with install.packages('ggplot2').")
  }
  
  csv_path <- file.path(results_dir, file)
  if (!file.exists(csv_path)) {
    stop("File not found: ", csv_path)
  }
  
  df <- utils::read.csv(csv_path, stringsAsFactors = FALSE, check.names = FALSE)
  
  # Must contain at least 2 columns: labels (A) and values (B)
  if (ncol(df) < 2) {
    stop("CSV must have at least two columns. Found: ", ncol(df))
  }
  
  # Column A (labels) and Column B (values)
  labels <- df[[1]]
  values <- suppressWarnings(as.numeric(df[[2]]))
  
  # Drop rows with missing labels or values
  keep <- !is.na(labels) & labels != "" & !is.na(values)
  labels <- labels[keep]
  values <- values[keep]
  
  # Convert variance -> standard deviation if requested
  if (convert_variance_to_sd) {
    if (any(values < 0, na.rm = TRUE)) {
      stop("convert_variance_to_sd=TRUE but some values are negative. Cannot take sqrt of negative variance.")
    }
    values <- sqrt(values)
  }
  
  plot_df <- data.frame(
    label = as.character(labels),
    value = as.numeric(values),
    stringsAsFactors = FALSE
  )
  
  # Optional sorting
  if (sort_bars) {
    ord <- order(plot_df$value, decreasing = descending)
    plot_df <- plot_df[ord, , drop = FALSE]
  }
  
  # Default text
  if (is.null(title))   title   <- tools::file_path_sans_ext(basename(file))
  if (is.null(x_label)) x_label <- names(df)[1]
  
  if (is.null(y_label)) {
    # Default to column name, but if converting variance->sd and the column name suggests variance, rename to SD-like label
    y_label <- names(df)[2]
    if (convert_variance_to_sd && grepl(variance_colname_pattern, y_label, perl = TRUE)) {
      y_label <- gsub(variance_colname_pattern, "sd", y_label, perl = TRUE)
      # common case: "mean_variance" -> "mean_sd"
    }
  }
  
  # Preserve display order
  plot_df$label <- factor(plot_df$label, levels = plot_df$label)
  
  p <- ggplot2::ggplot(plot_df, ggplot2::aes(x = label, y = value)) +
    ggplot2::geom_col() +
    ggplot2::labs(title = title, x = x_label, y = y_label) +
    ggplot2::theme_minimal()
  
  if (rotate_x) {
    p <- p + ggplot2::theme(axis.text.x = ggplot2::element_text(angle = 45, hjust = 1))
  }
  
  # Values on top of bars
  p <- p + ggplot2::geom_text(
    ggplot2::aes(label = format(round(value, digits), nsmall = digits)),
    vjust = -0.3,
    size = 3
  )
  
  # Output filename
  if (is.null(plot_path)) {
    plot_path <- paste0(tools::file_path_sans_ext(basename(file)), ".png")
  }
  
  out_path <- file.path(results_dir, plot_path)
  ggplot2::ggsave(filename = out_path, plot = p, width = width, height = height, dpi = dpi, bg = "white")
  
  message("Saved plot to: ", out_path)
  return(p)
}


# Run the function
# plot_results(
#   results_dir = "variance_results",
#   file = "final_results_summary.csv",
#   plot_path = "sd_by_repetitions.png",
#   title = "Mean SD of recommended sample size vs. repetitions",
#   x_label = "Repetitions used inside convergence search",
#   y_label = "Mean standard deviation (across random settings)",
#   convert_variance_to_sd = TRUE
# )
# 
# plot_results(
#   results_dir = "variance_results_consecutive_n",
#   file = "final_results_summary.csv",
#   plot_path = "sd_by_consecutive_n.png",
#   title = "Mean SD of recommended sample size vs. consecutive n",
#   x_label = "Consecutive n required",
#   y_label = "Mean standard deviation (across random settings)",
#   convert_variance_to_sd = TRUE
# )
# 
# plot_results(
#   results_dir = "mean_sample_size_results",
#   file = "final_results_summary.csv",
#   plot_path = "mean_sample_size_by_model.png",
#   title = "Mean recommended sample size by model",
#   x_label = "Model type",
#   y_label = "Mean recommended sample size"
# )


obtain_mean_sample_size_vs_dependence_by_model <- function(
    dependence_vector = c(0, 0.2, 0.5, 0.8, 1, 10, 100),
    num_random_settings = 5,
    n_obs = 100000,
    start_size = 200,
    end_size = 1200,
    step_size = 10,
    repetitions_inside_search = 10,
    consecutive_n = 5,
    convergence_threshold = 0.1,
    n_cores = 7,
    results_dir = "mean_recommended_n_by_dependence_results"
) {
  if (!dir.exists(results_dir)) dir.create(results_dir, recursive = TRUE)
  
  models <- c("OLS", "Ridge", "LASSO")
  
  pick_rmse <- function(res) res$rmse
  
  pop_fit_fn <- list(
    OLS   = ols_population,
    Ridge = ridge_population,
    LASSO = lasso_population
  )
  
  task_fn <- list(
    OLS   = ordinary_least_squares,
    Ridge = ridge_regression,
    LASSO = lasso_regression
  )
  
  # Results table
  raw <- data.frame(
    dependence      = numeric(),
    random_setting  = integer(),
    model_type      = character(),
    conv_value      = numeric(),
    stringsAsFactors = FALSE
  )
  
  print("Start mean recommended sample size vs dependence simulation.")
  print(paste0("Dependence values: ", paste(dependence_vector, collapse = ", ")))
  print(paste0("Random settings per dependence: ", num_random_settings))
  print(paste0("Fixed repetitions inside search: ", repetitions_inside_search))
  print(paste0("Fixed consecutive_n: ", consecutive_n))
  
  for (dep in dependence_vector) {
    print("==================================================")
    print(paste0("Dependence = ", dep))
    
    for (s in seq_len(num_random_settings)) {
      print(sprintf("Random setting %s / %s", s, num_random_settings))
      
      # Randomize parameters
      curr_n_vars <- sample(10:100, 1)
      curr_p_dependent <- 1
      curr_sparsity <- runif(1)
      curr_coef_mean <- runif(1, -0.5, 0.5)
      curr_coef_sd <- exp(rnorm(1, log(0.5), 0.5))
      curr_R2 <- 0.05 + 0.9 * rbeta(1, 2, 2)
      
      mega_population <- generate_continious_mega_population(
        n_obs = n_obs,
        n_vars = curr_n_vars,
        dependence = dep,
        p_dependent = curr_p_dependent,
        sparsity = curr_sparsity,
        coef_mean = curr_coef_mean,
        coef_sd = curr_coef_sd,
        R2 = curr_R2
      )
      
      # Train/test split for the population model
      safe_createDataPartition <- function(y, p = 0.8, groups_max = 5) {
        n <- length(y)
        if (n < 2) stop("y is too short")
        
        y2 <- y[is.finite(y)]
        if (length(y2) < 2 || length(unique(y2)) < 2) {
          return(sample.int(n, size = floor(p * n)))
        }
        
        for (g in seq(groups_max, 2L, by = -1L)) {
          idx <- tryCatch(
            caret::createDataPartition(y, p = p, list = FALSE, groups = g),
            error = function(e) NULL
          )
          if (!is.null(idx)) return(idx)
        }
        
        sample.int(n, size = floor(p * n))
      }
      
      train_index <- safe_createDataPartition(mega_population$data$y, p = 0.8)
      
      pop_train_data <- mega_population$data[train_index, ]
      pop_test_data  <- mega_population$data[-train_index, ]
      
      # Population RMSE per model
      pop_rmse <- list()
      for (m in models) {
        fit <- pop_fit_fn[[m]](train_data = pop_train_data, test_data = pop_test_data)
        pop_rmse[[m]] <- pick_rmse(fit)
      }
      
      # One convergence search per model
      for (m in models) {
        rec <- find_convergence_sample_size(
          mega_population$data,
          task = function(sz) {
            res <- task_fn[[m]](mega_population$data, sample_size = sz, metric = "RMSE")
            res$rmse
          },
          pop_value = pop_rmse[[m]],
          start_size = start_size,
          end_size = end_size,
          step_size = step_size,
          repetitions = repetitions_inside_search,
          convergence_threshold = convergence_threshold,
          consecutive_n = consecutive_n,
          n_cores = n_cores
        )
        
        conv <- if (!is.null(rec)) rec$converged_at_n else NULL
        conv_num <- if (is.null(conv)) NA_real_ else as.numeric(conv)
        
        raw <- rbind(
          raw,
          data.frame(
            dependence     = dep,
            random_setting = s,
            model_type     = m,
            conv_value     = conv_num,
            stringsAsFactors = FALSE
          )
        )
        
        print(sprintf("dependence=%s | setting=%s | model=%s | converged_at_n=%s", dep, s, m, conv_num))
      }
    }
  }
  
  # Calculate summary
  summary_tbl <- aggregate(
    conv_value ~ dependence + model_type,
    data = raw,
    FUN = function(x) mean(x, na.rm = TRUE)
  )
  names(summary_tbl)[names(summary_tbl) == "conv_value"] <- "mean_recommended_n"
  
  # Save
  utils::write.csv(raw, file = file.path(results_dir, "final_results_raw.csv"), row.names = FALSE)
  utils::write.csv(summary_tbl, file = file.path(results_dir, "final_results_summary.csv"), row.names = FALSE)
  
  invisible(list(raw = raw, summary = summary_tbl))
}

# obtain_mean_sample_size_vs_dependence_by_model(
#   dependence_vector = c(0, 0.2, 0.5, 0.8, 1, 10, 100),
#   num_random_settings = 5,
#   repetitions_inside_search = 10,
#   consecutive_n = 5,
#   results_dir = "dependence_mean_sample_size"
# )

plot_mean_recommended_sample_size_vs_dependence_by_model <- function(
    results_dir,
    file = "final_results_summary.csv",
    plot_path = "mean_recommended_sample_size_vs_dependence.png",
    title = "Mean recommended sample size vs. dependence",
    x_label = "Dependence",
    y_label = "Mean recommended sample size",
    width = 9,
    height = 5,
    dpi = 150
) {
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("ggplot2 not installed; install with install.packages('ggplot2').")
  }
  
  csv_path <- file.path(results_dir, file)
  if (!file.exists(csv_path)) stop("File not found: ", csv_path)
  
  df <- utils::read.csv(csv_path, stringsAsFactors = FALSE, check.names = FALSE)
  
  required <- c("dependence", "model_type", "mean_recommended_n")
  if (!all(required %in% names(df))) {
    stop("Expected columns: ", paste(required, collapse = ", "), " in ", file)
  }
  
  df$dependence <- factor(df$dependence, levels = unique(df$dependence))
  
  p <- ggplot2::ggplot(
    df,
    ggplot2::aes(x = dependence, y = mean_recommended_n, color = model_type)
  ) +
    ggplot2::geom_line(ggplot2::aes(group = model_type), linewidth = 1) +
    ggplot2::geom_point(size = 2) +
    ggplot2::labs(title = title, x = x_label, y = y_label, color = "Model") +
    ggplot2::theme_minimal()
  
  out_path <- file.path(results_dir, plot_path)
  ggplot2::ggsave(filename = out_path, plot = p, width = width, height = height, dpi = dpi, bg = "white")
  
  message("Saved plot to: ", out_path)
  return(p)
}

# plot_mean_recommended_sample_size_vs_dependence_by_model(
#   results_dir = "dependence_mean_sample_size",
#   plot_path = "dependence_mean_recommended_n.png"
# )

obtain_mean_sample_size_vs_sparsity_by_model <- function(
    sparsity_vector = c(0, 0.25, 0.5, 0.75, 1),
    dependence_candidates = c(0, 0.2, 0.5, 0.8, 1, 10, 100),
    p_dependent_sampler = function() runif(1, 0.1, 1),
    num_random_settings = 5,
    n_obs = 100000,
    start_size = 800,
    end_size = 1600,
    step_size = 10,
    repetitions_inside_search = 10,
    consecutive_n = 5,
    convergence_threshold = 0.1,
    n_cores = 7,
    results_dir = "mean_recommended_n_by_sparsity_results"
) {
  if (!dir.exists(results_dir)) dir.create(results_dir, recursive = TRUE)
  
  checkpoint_dir <- file.path(results_dir, "checkpoints")
  if (!dir.exists(checkpoint_dir)) dir.create(checkpoint_dir, recursive = TRUE)
  
  if (!requireNamespace("caret", quietly = TRUE)) {
    stop("caret not installed; install with install.packages('caret').")
  }
  
  models <- c("OLS", "Ridge", "LASSO")
  
  pick_rmse <- function(res) res$rmse
  
  pop_fit_fn <- list(
    OLS   = ols_population,
    Ridge = ridge_population,
    LASSO = lasso_population
  )
  
  task_fn <- list(
    OLS   = ordinary_least_squares,
    Ridge = ridge_regression,
    LASSO = lasso_regression
  )
  
  # Collect rows efficiently
  rows <- vector("list", length = 0L)
  row_i <- 0L
  
  message("Start mean recommended sample size vs sparsity simulation.")
  message("Sparsity values: ", paste(sparsity_vector, collapse = ", "))
  message("Random settings per sparsity: ", num_random_settings)
  
  for (sp in sparsity_vector) {
    message("==================================================")
    message("Sparsity = ", sp)
    
    sp_row_start <- row_i + 1L
    
    for (s in seq_len(num_random_settings)) {
      message(sprintf("Random setting %s / %s", s, num_random_settings))
      
      curr_n_vars       <- sample(10:100, 1)
      curr_dependence   <- sample(dependence_candidates, 1)
      curr_p_dependent  <- p_dependent_sampler()
      curr_sparsity     <- sp
      curr_coef_mean    <- runif(1, -0.5, 0.5)
      curr_coef_sd      <- exp(rnorm(1, log(0.5), 0.5))
      curr_R2           <- 0.05 + 0.9 * rbeta(1, 2, 2)
      
      mega_population <- generate_continious_mega_population(
        n_obs       = n_obs,
        n_vars      = curr_n_vars,
        dependence  = curr_dependence,
        p_dependent = curr_p_dependent,
        sparsity    = curr_sparsity,
        coef_mean   = curr_coef_mean,
        coef_sd     = curr_coef_sd,
        R2          = curr_R2
      )
      
      train_index <- caret::createDataPartition(
        mega_population$data$y, p = 0.8, list = FALSE
      )
      pop_train_data <- mega_population$data[train_index, ]
      pop_test_data  <- mega_population$data[-train_index, ]
      
      pop_rmse <- list()
      for (m in models) {
        fit <- pop_fit_fn[[m]](train_data = pop_train_data, test_data = pop_test_data)
        pop_rmse[[m]] <- pick_rmse(fit)
      }
      
      for (m in models) {
        rec <- find_convergence_sample_size(
          mega_population$data,
          task = function(sz) {
            res <- task_fn[[m]](mega_population$data, sample_size = sz, metric = "RMSE")
            res$rmse
          },
          pop_value = pop_rmse[[m]],
          start_size = start_size,
          end_size = end_size,
          step_size = step_size,
          repetitions = repetitions_inside_search,
          convergence_threshold = convergence_threshold,
          consecutive_n = consecutive_n,
          n_cores = n_cores
        )
        
        conv <- if (!is.null(rec)) rec$converged_at_n else NULL
        conv_num <- if (is.null(conv)) NA_real_ else as.numeric(conv)
        
        row_i <- row_i + 1L
        rows[[row_i]] <- data.frame(
          sparsity       = sp,
          dependence     = curr_dependence,
          p_dependent    = curr_p_dependent,
          random_setting = s,
          model_type     = m,
          conv_value     = conv_num,
          stringsAsFactors = FALSE
        )
        
        message(sprintf(
          "sparsity=%s | model=%s | converged_at_n=%s",
          sp, m, conv_num
        ))
      }
    }
    
    # Create a checkpoint for each sparcity value
    if (row_i >= sp_row_start) {
      sp_tag <- gsub("\\.", "p", as.character(sp))  # e.g. 0.25 -> "0p25"
      sp_checkpoint_path <- file.path(checkpoint_dir, paste0("raw_sparsity_", sp_tag, ".csv"))
      
      sp_raw <- do.call(rbind, rows[sp_row_start:row_i])
      utils::write.csv(sp_raw, file = sp_checkpoint_path, row.names = FALSE)
      
      message("Checkpoint saved for sparsity=", sp, " to: ", sp_checkpoint_path)
    }
    
    # Also save the latest raw file
    latest_path <- file.path(checkpoint_dir, "raw_latest.csv")
    utils::write.csv(do.call(rbind, rows[seq_len(row_i)]), file = latest_path, row.names = FALSE)
    
  }
  
  raw <- do.call(rbind, rows)
  
  # Calculate summary
  mean_safe <- function(x) {
    if (all(is.na(x))) NA_real_ else mean(x, na.rm = TRUE)
  }
  
  summary_tbl <- aggregate(
    conv_value ~ sparsity + model_type,
    data = raw,
    FUN = mean_safe
  )
  names(summary_tbl)[names(summary_tbl) == "conv_value"] <- "mean_recommended_n"
  
  # Save
  utils::write.csv(raw, file = file.path(results_dir, "final_results_raw.csv"), row.names = FALSE)
  utils::write.csv(summary_tbl, file = file.path(results_dir, "final_results_summary.csv"), row.names = FALSE)
  
  invisible(list(raw = raw, summary = summary_tbl))
}

# obtain_mean_sample_size_vs_sparsity_by_model(
#   sparsity_vector = c(0, 0.2, 0.5, 0.7, 0.9),
#   num_random_settings = 5,
#   repetitions_inside_search = 10,
#   consecutive_n = 5,
#   results_dir = "sparsity_mean_sample_size"
# )

plot_mean_recommended_sample_size_vs_sparsity_by_model <- function(
    results_dir,
    file = "final_results_summary.csv",
    plot_path = "mean_recommended_sample_size_vs_sparsity.png",
    title = "Mean recommended sample size vs. sparsity",
    x_label = "Sparsity",
    y_label = "Mean recommended sample size",
    width = 9,
    height = 5,
    dpi = 150
) {
  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("ggplot2 not installed; install with install.packages('ggplot2').")
  }
  
  csv_path <- file.path(results_dir, file)
  if (!file.exists(csv_path)) stop("File not found: ", csv_path)
  
  df <- utils::read.csv(csv_path, stringsAsFactors = FALSE, check.names = FALSE)
  
  required <- c("sparsity", "model_type", "mean_recommended_n")
  if (!all(required %in% names(df))) {
    stop("Expected columns: ", paste(required, collapse = ", "), " in ", file)
  }
  
  df$sparsity <- factor(df$sparsity, levels = unique(df$sparsity))
  
  p <- ggplot2::ggplot(
    df,
    ggplot2::aes(x = sparsity, y = mean_recommended_n, color = model_type)
  ) +
    ggplot2::geom_line(ggplot2::aes(group = model_type), linewidth = 1) +
    ggplot2::geom_point(size = 2) +
    ggplot2::labs(title = title, x = x_label, y = y_label, color = "Model") +
    ggplot2::theme_minimal()
  
  out_path <- file.path(results_dir, plot_path)
  ggplot2::ggsave(filename = out_path, plot = p, width = width, height = height, dpi = dpi, bg = "white")
  
  message("Saved plot to: ", out_path)
  return(p)
}

# plot_mean_recommended_sample_size_vs_sparsity_by_model(results_dir = "sparsity_mean_sample_size")