#!/usr/bin/env Rscript
#
# Offline network enrichment and interaction analysis using local STRING-DB files.
# Replicates the NETWORK_ENRICHMENT Nextflow process without internet connection.
# R equivalent of the Python script using STRINGdb package.
#

suppressPackageStartupMessages({
  library(STRINGdb)
  library(data.table)
  library(argparse)
})

#' Load STRING-DB data from local files
#'
#' @param string_db_dir Directory containing STRING-DB files
#' @return List containing mappings and data frames
load_string_data <- function(string_db_dir) {
  cat("Loading STRING-DB files...\n")
  
  # Load alias mapping (ENST -> ENSP)
  aliases <- fread(
    file.path(string_db_dir, "9606.protein.aliases.enst.v12.0.tsv"),
    sep = "\t",
    header = FALSE,
    col.names = c("string_protein_id", "alias", "source")
  )
  
  enst_to_ensp <- aliases[source == "Ensembl_transcript", 
                          .(string_protein_id, alias)]
  setkey(enst_to_ensp, alias)
  
  # Load protein info (for preferred names)
  info <- fread(
    file.path(string_db_dir, "9606.protein.info.v12.0.txt"),
    sep = "\t",
    header = FALSE,
    col.names = c("string_protein_id", "preferred_name", "protein_size", "annotation")
  )
  setkey(info, string_protein_id)
  
  # Load protein links - USE DETAILED FILE
  links <- fread(
    file.path(string_db_dir, "9606.protein.links.detailed.v12.0.txt"),
    sep = " "
  )
  
  # Load enrichment terms
  enrichment <- fread(
    file.path(string_db_dir, "9606.protein.enrichment.terms.v12.0.txt"),
    sep = "\t",
    header = FALSE,
    col.names = c("string_protein_id", "category", "term", "description")
  )
  
  cat(sprintf("Loaded %d ENST->ENSP mappings\n", nrow(enst_to_ensp)))
  cat(sprintf("Loaded %d protein-protein interactions\n", nrow(links)))
  cat(sprintf("Loaded %d enrichment annotations\n", nrow(enrichment)))
  
  list(
    enst_to_ensp = enst_to_ensp,
    ensp_to_name = info,
    links = links,
    enrichment = enrichment
  )
}

#' Perform enrichment analysis using Fisher's exact test
#'
#' @param query_proteins Vector of ENSP protein IDs
#' @param enrichment_df Data table with enrichment annotations
#' @param ensp_to_name Data table mapping ENSP to preferred names
#' @param ensp_to_input_map Named vector mapping ENSP back to ENST
#' @return Data table with enrichment results
perform_enrichment_analysis <- function(query_proteins, enrichment_df, 
                                       ensp_to_name, ensp_to_input_map) {
  query_set <- unique(query_proteins)
  n_query <- length(query_set)
  
  cat(sprintf("Analyzing enrichment for %d proteins...\n", n_query))
  
  results_list <- list()
  
  # Group by category
  categories <- unique(enrichment_df$category)
  
  for (cat in categories) {
    cat_group <- enrichment_df[category == cat]
    
    # Background: Total unique proteins annotated in this category
    category_universe <- unique(cat_group$string_protein_id)
    background_size <- length(category_universe)
    
    # Iterate over terms within this category
    terms <- unique(cat_group$term)
    
    for (tm in terms) {
      term_group <- cat_group[term == tm]
      term_proteins <- unique(term_group$string_protein_id)
      
      # Intersection: proteins in both query and term
      intersection <- intersect(query_set, term_proteins)
      n_intersection <- length(intersection)
      
      if (n_intersection == 0) next
      
      # Background term count
      n_term_background <- length(term_proteins)
      
      # Fisher's exact test
      # Contingency table:
      #                  In term    Not in term
      # In query           a            b
      # Not in query       c            d
      
      a <- n_intersection
      b <- n_query - a
      c <- n_term_background - a
      d <- background_size - n_term_background - b
      
      # Ensure no negative values
      if (c < 0 || d < 0) next
      
      # One-sided test for enrichment
      contingency_matrix <- matrix(c(a, b, c, d), nrow = 2, byrow = TRUE)
      fisher_result <- fisher.test(contingency_matrix, alternative = "greater")
      pvalue <- fisher_result$p.value
      
      # Get description
      description <- term_group$description[1]
      
      # Get input genes (map ENSP back to ENST input)
      intersection_sorted <- sort(intersection)
      
      input_genes_list <- sapply(intersection_sorted, function(ensp) {
        if (ensp %in% names(ensp_to_input_map)) {
          ensp_to_input_map[ensp]
        } else {
          ensp
        }
      })
      
      input_genes <- paste(input_genes_list, collapse = ",")
      
      # Get preferred names
      preferred_names_list <- sapply(intersection_sorted, function(p) {
        name <- ensp_to_name[string_protein_id == p, preferred_name]
        if (length(name) > 0) name[1] else strsplit(p, "\\.")[[1]][length(strsplit(p, "\\.")[[1]])]
      })
      preferred_names <- paste(preferred_names_list, collapse = ",")
      
      results_list[[length(results_list) + 1]] <- data.table(
        category = cat,
        term = tm,
        number_of_genes = n_intersection,
        number_of_genes_in_background = n_term_background,
        ncbiTaxonId = 9606,
        inputGenes = input_genes,
        preferredNames = preferred_names,
        p_value = pvalue,
        description = description
      )
    }
  }
  
  if (length(results_list) == 0) {
    return(data.table())
  }
  
  # Combine results
  results_df <- rbindlist(results_list)
  
  # Benjamini-Hochberg FDR correction
  results_df <- results_df[order(p_value)]
  n_tests <- nrow(results_df)
  results_df[, fdr := p_value * n_tests / seq_len(n_tests)]
  results_df[fdr > 1.0, fdr := 1.0]
  
  # Reorder columns to match API output
  setcolorder(results_df, c(
    "category", "term", "number_of_genes", "number_of_genes_in_background",
    "ncbiTaxonId", "inputGenes", "preferredNames", "p_value", "fdr", "description"
  ))
  
  return(results_df)
}

#' Extract network interactions for query proteins
#'
#' @param query_proteins Vector of ENSP protein IDs
#' @param links_df Data table with protein-protein interactions
#' @param ensp_to_name Data table mapping ENSP to preferred names
#' @return Data table with network interactions
get_network_interactions <- function(query_proteins, links_df, ensp_to_name) {
  query_set <- unique(query_proteins)
  
  # Filter links where both proteins are in query set
  interactions <- links_df[protein1 %in% query_set & protein2 %in% query_set]
  
  if (nrow(interactions) == 0) {
    return(data.table())
  }
  
  # Normalize scores (STRING files are integers 0-1000, API is float 0-1)
  score_cols <- c("combined_score", "neighborhood", "fusion", "cooccurence",
                  "coexpression", "experimental", "database", "textmining")
  target_cols <- c("score", "nscore", "fscore", "pscore",
                   "ascore", "escore", "dscore", "tscore")
  
  for (i in seq_along(score_cols)) {
    src <- score_cols[i]
    tgt <- target_cols[i]
    if (src %in% names(interactions)) {
      interactions[, (tgt) := get(src) / 1000.0]
    } else {
      interactions[, (tgt) := 0.0]
    }
  }
  
  # Add preferred names
  interactions[, preferredName_A := sapply(protein1, function(p) {
    name <- ensp_to_name[string_protein_id == p, preferred_name]
    if (length(name) > 0) name[1] else strsplit(p, "\\.")[[1]][length(strsplit(p, "\\.")[[1]])]
  })]
  
  interactions[, preferredName_B := sapply(protein2, function(p) {
    name <- ensp_to_name[string_protein_id == p, preferred_name]
    if (length(name) > 0) name[1] else strsplit(p, "\\.")[[1]][length(strsplit(p, "\\.")[[1]])]
  })]
  
  # Rename columns to match STRING API output
  setnames(interactions, c("protein1", "protein2"), c("stringId_A", "stringId_B"))
  
  interactions[, ncbiTaxonId := 9606]
  
  # Select and order columns to match original output
  column_order <- c(
    "stringId_A", "stringId_B", "preferredName_A", "preferredName_B",
    "ncbiTaxonId", "score", "nscore", "fscore", "pscore",
    "ascore", "escore", "dscore", "tscore"
  )
  
  interactions <- interactions[, ..column_order]
  
  return(interactions)
}

#' Main function
main <- function() {
  # Parse command line arguments
  parser <- ArgumentParser(
    description = "Offline network enrichment analysis using local STRING-DB files"
  )
  
  parser$add_argument(
    "interactions_file",
    type = "character",
    help = "Input TSV file with significant interactions (conplex output)"
  )
  
  parser$add_argument(
    "--string-db-dir",
    type = "character",
    default = "/app",
    help = "Directory containing STRING-DB files"
  )
  
  parser$add_argument(
    "--threshold",
    type = "double",
    default = 0.65,
    help = "ConPlex score threshold for protein network"
  )
  
  parser$add_argument(
    "--max-proteins",
    type = "integer",
    default = 450L,
    help = "Maximum number of proteins to analyze"
  )
  
  parser$add_argument(
    "--output-dir",
    type = "character",
    default = ".",
    help = "Output directory for results"
  )
  
  args <- parser$parse_args()
  
  # Load STRING-DB data
  string_data <- load_string_data(args$string_db_dir)
  
  # Load interaction data
  cat(sprintf("\nLoading interaction data from %s...\n", args$interactions_file))
  interaction_data <- fread(args$interactions_file, sep = "\t")
  interaction_data <- interaction_data[order(-conplex_score)]
  
  # Remove duplicates and limit
  # Replicates Nextflow logic: drops duplicates on transcript BEFORE limit
  interaction_data <- unique(interaction_data, by = "transcipt")
  
  if (nrow(interaction_data) > args$max_proteins) {
    cat(sprintf("Limiting to top %d proteins\n", args$max_proteins))
    interaction_data <- interaction_data[1:args$max_proteins]
  }
  
  # Filter by threshold
  filtered_data <- interaction_data[conplex_score > args$threshold]
  cat(sprintf("Found %d proteins above threshold %.2f\n", 
              nrow(filtered_data), args$threshold))
  
  # Map ENST to ENSP and keep track of mapping for output
  enst_list <- filtered_data$transcipt
  ensp_list <- character()
  ensp_to_input_map <- character()
  
  for (enst in enst_list) {
    ensp <- string_data$enst_to_ensp[alias == enst, string_protein_id]
    if (length(ensp) > 0) {
      ensp <- ensp[1]
      ensp_list <- c(ensp_list, ensp)
      ensp_to_input_map[ensp] <- enst
    }
  }
  
  cat(sprintf("Mapped %d ENST IDs to ENSP IDs\n", length(ensp_list)))
  
  if (length(ensp_list) == 0) {
    stop("ERROR: No valid ENSP mappings found!")
  }
  
  # Perform enrichment analysis
  cat("\nPerforming enrichment analysis...\n")
  enrichment_results <- perform_enrichment_analysis(
    ensp_list, 
    string_data$enrichment, 
    string_data$ensp_to_name,
    ensp_to_input_map
  )
  cat(sprintf("Found %d enriched terms\n", nrow(enrichment_results)))
  
  # Get network interactions
  cat("\nExtracting network interactions...\n")
  network_interactions <- get_network_interactions(
    ensp_list, 
    string_data$links, 
    string_data$ensp_to_name
  )
  cat(sprintf("Found %d protein-protein interactions\n", nrow(network_interactions)))
  
  # Generate output filename
  output_name <- gsub("_significant_interactions", "", 
                     tools::file_path_sans_ext(basename(args$interactions_file)))
  
  # Create output directory if needed
  dir.create(args$output_dir, showWarnings = FALSE, recursive = TRUE)
  
  # Save results
  enrichment_file <- file.path(args$output_dir, 
                               sprintf("%s_network_enrichment.tsv", output_name))
  interactions_file <- file.path(args$output_dir, 
                                sprintf("%s_network_interactions.tsv", output_name))
  
  fwrite(enrichment_results, enrichment_file, sep = "\t")
  fwrite(network_interactions, interactions_file, sep = "\t")
  
  cat("\nResults saved:\n")
  cat(sprintf("  Enrichment: %s\n", enrichment_file))
  cat(sprintf("  Interactions: %s\n", interactions_file))
}

# Run main function
if (!interactive()) {
  main()
}