#!/usr/bin/env Rscript
#
# Network enrichment and interaction analysis using STRINGdb package.
# R equivalent of the NETWORK_ENRICHMENT Nextflow process.
#

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

#' Main function
main <- function() {
  # Parse command line arguments
  parser <- ArgumentParser(
    description = "Network enrichment analysis using STRINGdb package"
  )
  
  parser$add_argument(
    "interactions_file",
    type = "character",
    help = "Input TSV file with significant interactions (conplex output)"
  )
  
  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 (STRING API limitation)"
  )
  
  parser$add_argument(
    "--output-dir",
    type = "character",
    default = ".",
    help = "Output directory for results"
  )
  
  parser$add_argument(
    "--species",
    type = "integer",
    default = 9606L,
    help = "NCBI taxon ID (default: 9606 for Homo sapiens)"
  )
  
  parser$add_argument(
    "--score-threshold",
    type = "integer",
    default = 400L,
    help = "STRING combined score threshold (0-1000, default: 400)"
  )
  
  args <- parser$parse_args()
  
  # Initialize STRINGdb
  cat("Initializing STRINGdb...\n")
  string_db <- STRINGdb$new(
    version = "12.0",
    species = args$species,
    score_threshold = args$score_threshold,
    network_type = "full",
    input_directory = "/app"
  )
  
  # Load interaction data
  cat(sprintf("Loading 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)
  interaction_data <- unique(interaction_data, by = "transcipt")
  
  if (nrow(interaction_data) > args$max_proteins) {
    cat(sprintf("Limiting to top %d proteins (STRING API limitation)\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))
  
  if (nrow(filtered_data) == 0) {
    stop("ERROR: No proteins above threshold!")
  }
  
  # Prepare data for STRING mapping
  gene_list <- data.frame(
    gene = filtered_data$transcipt,
    stringsAsFactors = FALSE
  )
  
  # Map genes to STRING IDs
  cat("Mapping genes to STRING database...\n")
  mapped <- string_db$map(gene_list, "gene", removeUnmappedRows = TRUE)
  
  if (nrow(mapped) == 0) {
    stop("ERROR: No genes could be mapped to STRING database!")
  }
  
  cat(sprintf("Successfully mapped %d/%d genes to STRING IDs\n", 
              nrow(mapped), nrow(gene_list)))
  
  # Get STRING IDs
  string_ids <- mapped$STRING_id
  
  # Get enrichment analysis
  cat("\nPerforming enrichment analysis...\n")
  enrichment_results <- string_db$get_enrichment(string_ids, category = "all", methodMT = "fdr")
  
  if (nrow(enrichment_results) == 0) {
    cat("WARNING: No enrichment results found\n")
    enrichment_df <- data.table()
  } else {
    cat(sprintf("Found %d enriched terms\n", nrow(enrichment_results)))
    
    # Convert to data.table and rename columns to match Python output
    enrichment_df <- as.data.table(enrichment_results)
    
    # The STRINGdb enrichment output has columns:
    # category, term, number_of_genes, number_of_genes_in_background, 
    # ncbiTaxonId, inputGenes, preferredNames, p_value, fdr, description
    # This should already match the Python output format
  }
  
  # Get network interactions with detailed scores
  cat("\nExtracting network interactions...\n")
  interactions <- string_db$get_interactions(string_ids)
  
  if (nrow(interactions) == 0) {
    cat("WARNING: No interactions found\n")
    interactions_df <- data.table()
  } else {
    cat(sprintf("Found %d protein-protein interactions\n", nrow(interactions)))
    
    # Convert to data.table
    interactions_df <- as.data.table(interactions)
    
    # The get_interactions() function returns:
    # from, to, combined_score, and potentially other columns
    # We need to get detailed scores separately
    
    # Get the interaction network with all score types
    # Use get_png to access the full network data including detailed scores
    network_image <- string_db$plot_network(string_ids)
    
    # Alternative: Access the graph directly for detailed scores
    # The STRINGdb object stores the full network internally
    graph <- string_db$get_graph()
    
    # Extract edge attributes which contain detailed scores
    if (!is.null(graph)) {
      edge_data <- igraph::as_data_frame(graph, what = "edges")
      edge_dt <- as.data.table(edge_data)
      
      # Filter for our query proteins
      edge_dt <- edge_dt[from %in% string_ids & to %in% string_ids]
      
      # Rename columns to match Python output format
      if ("from" %in% names(edge_dt)) setnames(edge_dt, "from", "stringId_A")
      if ("to" %in% names(edge_dt)) setnames(edge_dt, "to", "stringId_B")
      
      # Add preferred names
      edge_dt[, preferredName_A := string_db$get_aliases(stringId_A)$alias[1], by = stringId_A]
      edge_dt[, preferredName_B := string_db$get_aliases(stringId_B)$alias[1], by = stringId_B]
      
      # Normalize scores to 0-1 range if they're in 0-1000 range
      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(edge_dt)) {
          # Check if scores are in 0-1000 range (need normalization)
          max_val <- max(edge_dt[[src]], na.rm = TRUE)
          if (max_val > 1) {
            edge_dt[, (tgt) := get(src) / 1000.0]
          } else {
            edge_dt[, (tgt) := get(src)]
          }
        }
      }
      
      # Add ncbiTaxonId
      edge_dt[, ncbiTaxonId := args$species]
      
      # Select and order columns to match Python output
      output_cols <- c("stringId_A", "stringId_B", "preferredName_A", "preferredName_B",
                       "ncbiTaxonId", "score", "nscore", "fscore", "pscore",
                       "ascore", "escore", "dscore", "tscore")
      
      # Keep only columns that exist
      existing_cols <- intersect(output_cols, names(edge_dt))
      interactions_df <- edge_dt[, ..existing_cols]
      
    } else {
      # Fallback: use the interactions data from get_interactions()
      # Rename columns to match output format
      if ("from" %in% names(interactions_df)) {
        setnames(interactions_df, "from", "stringId_A")
      }
      if ("to" %in% names(interactions_df)) {
        setnames(interactions_df, "to", "stringId_B")
      }
      if ("combined_score" %in% names(interactions_df)) {
        interactions_df[, score := combined_score / 1000.0]
      }
      
      interactions_df[, ncbiTaxonId := args$species]
    }
  }
  
  # 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_df, enrichment_file, sep = "\t")
  fwrite(interactions_df, 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()
}