suggest.R

#' Regression analysis for vulnerability factors.
#'
#' This script performs logistic regression analysis on vulnerability factors
#' using survey-weighted data to identify significant predictors of health outcomes.

library(dplyr)
library(duckdb)
library(survey)
library(stringr)
library(parallel)

# Input file paths
MAIN_DATA_PATH <- "Senegal_2019DHS8_1.0/data/input/SEN_2019DHS8.parquet"
OUTCOMES_PATH <- "Senegal_2019DHS8_1.0/data/input/SEN_2019DHS8_outcomes.parquet"
VULNERABILITIES_PATH <- "Senegal_2019DHS8_1.0/data/input/SEN_2019DHS8_vulnerabilities.parquet"

# Analysis thresholds
P_VALUE_THRESHOLD_1 <- 0.001
P_VALUE_THRESHOLD_2 <- 0.01
P_VALUE_THRESHOLD_3 <- 0.05
COEF_THRESHOLD <- 4
STD_ERROR_THRESHOLD <- 1
SAMPLE_SIZE_THRESHOLD <- 3000

#' Read parquet file using DuckDB
#'
#' @param file_path Character string specifying the path to the parquet file
#' @return A data frame containing the parquet data
read_parquet_file <- function(file_path) {
  con <- dbConnect(duckdb::duckdb())
  on.exit(dbDisconnect(con))
  
  query <- sprintf("CREATE TABLE data AS SELECT * FROM read_parquet('%s')", file_path)
  dbExecute(con, query)
  dbGetQuery(con, "SELECT * FROM data")
}

#' Validate variable for regression analysis
#'
#' @param var_data Vector containing the variable data
#' @param variable_name Character string of the variable name for error messages
#' @return Logical indicating if variable is valid for analysis
validate_variable <- function(var_data, variable_name) {
  if (all(is.na(var_data))) {
    warning(paste("Variable", variable_name, "has all missing values"))
    return(FALSE)
  }
  
  if (is.logical(var_data) || is.character(var_data) || is.factor(var_data)) {
    unique_values <- length(unique(var_data[!is.na(var_data)]))
    if (unique_values < 2) {
      warning(paste("Variable", variable_name, "has insufficient variation"))
      return(FALSE)
    }
  } else if (is.numeric(var_data)) {
    if (var(var_data, na.rm = TRUE) == 0) {
      warning(paste("Variable", variable_name, "has no variation"))
      return(FALSE)
    }
  } else {
    warning(paste("Variable", variable_name, "has unsupported type:", typeof(var_data)))
    return(FALSE)
  }
  
  TRUE
}

#' Prepare variable for analysis by converting to appropriate type
#'
#' @param var_data Vector containing the variable data
#' @return Converted variable vector
prepare_variable <- function(var_data) {
  if (is.logical(var_data) || is.character(var_data)) {
    return(as.factor(var_data))
  }
  var_data
}

#' Fit survey-weighted logistic regression model
#'
#' @param design Pre-created survey design object
#' @param outcome Character string of outcome variable name
#' @param predictor Character string of predictor variable name
#' @return Survey GLM model object or NULL if fitting fails
fit_survey_model <- function(design, outcome, predictor) {
  formula_str <- paste(outcome, "~", predictor)
  model_formula <- as.formula(formula_str)
  
  tryCatch({
    svyglm(model_formula, design = design, family = quasibinomial())
  }, error = function(e) {
    warning(paste("Model fitting failed for", predictor, ":", e$message))
    NULL
  })
}

#' Extract model statistics from survey GLM
#'
#' @param model Survey GLM model object
#' @param outcome_data Vector of outcome variable data
#' @param predictor_data Vector of predictor variable data
#' @param predictor_name Character string of predictor variable name
#' @return Named list of model statistics
extract_model_stats <- function(model, outcome_data, predictor_data, predictor_name) {
  if (is.null(model)) return(NULL)
  
  model_summary <- summary(model)
  coef_table <- model_summary$coefficients
  
  if (nrow(coef_table) < 2) return(NULL)
  
  coefficient <- coef_table[2, 1]
  std_error <- coef_table[2, 2]
  p_value <- coef_table[2, 4]
  
  ci_lower <- coefficient - 1.96 * std_error
  ci_upper <- coefficient + 1.96 * std_error
  
  odds_ratio <- exp(coefficient)
  or_ci_lower <- exp(ci_lower)
  or_ci_upper <- exp(ci_upper)
  
  complete_cases <- complete.cases(outcome_data, predictor_data)
  n_samples <- sum(complete_cases)
  
  list(
    variable = predictor_name,
    p_value = p_value,
    coefficient = coefficient,
    std_error = std_error,
    conf_interval_95 = paste0("(", round(ci_lower, 3), ", ", round(ci_upper, 3), ")"),
    odds_ratio_95ci = paste0(round(odds_ratio, 3), " (", round(or_ci_lower, 3), ", ", round(or_ci_upper, 3), ")"),
    sample_size = n_samples
  )
}

#' Add significance stars based on p-values
#'
#' @param p_values Numeric vector of p-values
#' @return Character vector of significance indicators
add_significance_stars <- function(p_values) {
  case_when(
    p_values < P_VALUE_THRESHOLD_1 ~ "***",
    p_values < P_VALUE_THRESHOLD_2 ~ "**",
    p_values < P_VALUE_THRESHOLD_3 ~ "*",
    TRUE ~ ""
  )
}

#' Extract numeric odds ratio from formatted string
#'
#' @param or_string Character string containing odds ratio and CI
#' @return Numeric odds ratio value
extract_odds_ratio <- function(or_string) {
  as.numeric(str_extract(or_string, "^[0-9.]+"))
}

#' Identify potential issues with model results
#'
#' @param results_df Data frame containing model results
#' @return Character vector of red flags for each variable
identify_red_flags <- function(results_df) {
  flags <- character(nrow(results_df))
  
  # Check for separation
  separation_idx <- abs(results_df$coefficient) > COEF_THRESHOLD
  flags[separation_idx] <- paste0(flags[separation_idx], "SEPARATION;")
  
  # Check for large standard errors
  large_se_idx <- results_df$std_error > STD_ERROR_THRESHOLD
  flags[large_se_idx] <- paste0(flags[large_se_idx], "LARGE_SE;")
  
  # Check for small sample sizes
  small_n_idx <- results_df$sample_size < SAMPLE_SIZE_THRESHOLD
  flags[small_n_idx] <- paste0(flags[small_n_idx], "SMALL_N;")
  
  str_remove(flags, ";$")
}

#' Perform vulnerability factor analysis (optimized)
#'
#' @param vulnerability_variables Character vector of variable names to analyze
#' @param main_data Data frame containing all survey data
#' @param outcome Character string of outcome variable name
#' @param use_parallel Logical, whether to use parallel processing
#' @return Data frame with analysis results sorted by p-value
analyze_vulnerability_factors <- function(vulnerability_variables, main_data, outcome, use_parallel = TRUE) {
  # Pre-filter valid variables to avoid redundant checks
  valid_vars <- vulnerability_variables[vulnerability_variables %in% colnames(main_data)]
  if (length(valid_vars) == 0) {
    warning("No valid variables found")
    return(data.frame())
  }
  
  # Pre-filter complete cases for outcome
  outcome_complete <- !is.na(main_data[[outcome]])
  if (sum(outcome_complete) == 0) {
    warning("No complete cases for outcome")
    return(data.frame())
  }
  
  # Prepare all variables at once
  prepared_data <- main_data
  for (var in valid_vars) {
    var_data <- main_data[[var]]
    if (validate_variable(var_data, var)) {
      prepared_data[[var]] <- prepare_variable(var_data)
    } else {
      valid_vars <- valid_vars[valid_vars != var]
    }
  }
  
  if (length(valid_vars) == 0) {
    warning("No valid variables after preparation")
    return(data.frame())
  }
  
  # Create survey design once
  design <- svydesign(
    ids = ~psu,
    strata = ~strata,
    weights = ~weight,
    data = prepared_data,
    nest = TRUE
  )
  
  # Define analysis function for single variable
  analyze_single_var <- function(variable) {
    model <- fit_survey_model(design, outcome, variable)
    extract_model_stats(model, main_data[[outcome]], prepared_data[[variable]], variable)
  }
  
  # Choose processing method
  if (use_parallel && length(valid_vars) > 4) {
    num_cores <- min(detectCores() - 1, length(valid_vars), 8)
    cat("Using", num_cores, "cores for parallel processing\n")
    all_results <- mclapply(valid_vars, analyze_single_var, mc.cores = num_cores)
    names(all_results) <- valid_vars
  } else {
    all_results <- list()
    for (variable in valid_vars) {
      cat("Processing:", variable, "\n")
      all_results[[variable]] <- analyze_single_var(variable)
    }
  }
  
  # Filter valid results
  valid_results <- all_results[!sapply(all_results, is.null)]
  
  if (length(valid_results) == 0) {
    warning("No valid results obtained")
    return(data.frame())
  }
  
  results_df <- bind_rows(valid_results) %>%
    arrange(p_value) %>%
    mutate(
      significance = add_significance_stars(p_value),
      red_flags = identify_red_flags(.)
    )
  
  results_df
}

#' Check if outcome variable is suitable for binary logistic regression
#'
#' @param outcome_data Vector containing the outcome variable data
#' @param outcome_name Character string of outcome variable name
#' @return Logical indicating if outcome is valid for binary analysis
validate_outcome <- function(outcome_data, outcome_name) {
  if (all(is.na(outcome_data))) {
    warning(paste("Outcome", outcome_name, "has all missing values"))
    return(FALSE)
  }
  
  unique_values <- unique(outcome_data[!is.na(outcome_data)])
  if (length(unique_values) != 2) {
    warning(paste("Outcome", outcome_name, "is not binary"))
    return(FALSE)
  }
  
  TRUE
}

#' Analyze vulnerability factors for all outcomes
#'
#' @param vulnerability_variables Character vector of variable names to analyze
#' @param main_data Data frame containing all survey data
#' @param outcome_variables Character vector of outcome variable names
#' @param use_parallel Logical, whether to use parallel processing
#' @return Data frame with analysis results for all outcomes
analyze_all_outcomes <- function(vulnerability_variables, main_data, outcome_variables, use_parallel = TRUE) {
  valid_outcomes <- character()
  for (outcome in outcome_variables) {
    if (outcome %in% colnames(main_data)) {
      outcome_data <- main_data[[outcome]]
      if (validate_outcome(outcome_data, outcome)) {
        valid_outcomes <- c(valid_outcomes, outcome)
      }
    } else {
      warning(paste("Outcome", outcome, "not found in dataset"))
    }
  }
  
  if (length(valid_outcomes) == 0) {
    warning("No valid outcomes found")
    return(data.frame())
  }
  
  cat("Analyzing", length(valid_outcomes), "valid outcomes\n")
  
  analyze_single_outcome <- function(outcome) {
    cat("Analyzing outcome:", outcome, "\n")
    results <- analyze_vulnerability_factors(vulnerability_variables, main_data, outcome, use_parallel = FALSE)
    if (nrow(results) > 0) {
      results$outcome <- outcome
      return(results)
    }
    return(NULL)
  }
  
  all_outcome_results <- list()
  for (outcome in valid_outcomes) {
    result <- analyze_single_outcome(outcome)
    if (!is.null(result)) {
      all_outcome_results[[outcome]] <- result
    }
  }
  
  if (length(all_outcome_results) == 0) {
    warning("No valid results obtained for any outcome")
    return(data.frame())
  }
  
  bind_rows(all_outcome_results) %>%
    arrange(outcome, p_value)
}

#' Main analysis function with file output
#'
#' @param main_path Path to main data parquet file
#' @param vulnerabilities_path Path to vulnerabilities parquet file
#' @param outcomes_path Path to outcomes parquet file
#' @param output_path Path for output CSV file
#' @param use_parallel Logical, whether to use parallel processing
run_analysis <- function(main_path, vulnerabilities_path, outcomes_path, output_path = "results.csv", use_parallel = TRUE) {
  cat("Loading data files...\n")
  main_data <- read_parquet_file(main_path)
  vulnerabilities <- read_parquet_file(vulnerabilities_path)
  outcomes <- read_parquet_file(outcomes_path)
  
  cat("Starting analysis...\n")
  results <- analyze_all_outcomes(vulnerabilities$variable, main_data, outcomes$variable, use_parallel)
  
  write.csv(results, output_path, row.names = FALSE)
  cat("Total results:", nrow(results), "\n")
  results
}

args <- commandArgs(trailingOnly = TRUE)

if (length(args) == 0) {
  results <- run_analysis(MAIN_DATA_PATH, VULNERABILITIES_PATH, OUTCOMES_PATH)
} else if (length(args) >= 3) {
  main_path <- args[1]
  vulnerabilities_path <- args[2]
  outcomes_path <- args[3]
  output_path <- if (length(args) >= 4) args[4] else "results.csv"
  results <- run_analysis(main_path, vulnerabilities_path, outcomes_path, output_path)
} else {
  cat("Usage: Rscript suggest.R [main_data_path] [vulnerabilities_path] [outcomes_path] [output_path]\n")
}