#' Generate Violin Plots for Bulk Expression Data
#'
#' Creates violin plots displaying the distribution of gene expression values
#' across Sex and Phenotype groups for a set of selected genes. Multiple genes
#' (up to 10) are visualized using faceting.
#'
#' @param expression_data A matrix or data.frame of bulk expression values
#'   with genes in rows and samples in columns.
#' @param sex A vector indicating the sex of each sample.
#' @param phenotype A vector indicating the phenotype or condition of each sample.
#' @param genes Character vector of genes to visualize (maximum 10 recommended).
#' @param sex_labels_vector Character vector specifying sex labels.
#' @param phenotype_labels_vector Character vector specifying phenotype labels.
#'
#' @return A `ggplot` object containing the violin plots.
#'
#' @import ggplot2
#' @importFrom tidyr pivot_longer
#' @importFrom dplyr mutate
#'
#' @export

generate_violinplot_bulk <- function(expression_data, 
                                     sex, 
                                     phenotype,
                                     genes,
                                     sex_labels_vector = c("F", "M"),
                                     phenotype_labels_vector = c("WT", "TG")) {
  
  #  Validate genes
  if (length(genes) == 0) {
    stop("No genes provided")
  }
  
  genes_present <- genes[genes %in% rownames(expression_data)]
  
  if (length(genes_present) == 0) {
    stop("None of the provided genes found in expression_data")
  }
  
  if (length(genes_present) < length(genes)) {
    warning("Some genes were not found and will be ignored")
  }
  
  # Limit number of genes 
  if (length(genes_present) > 10) {
    warning("More than 10 genes provided - plotting first 10 only")
    genes_present <- genes_present[1:10]
  }
  
  #  Transpose expression (samples as rows) 
  expr_df <- as.data.frame(t(expression_data[genes_present, , drop = FALSE]))
  
  expr_df$sex <- factor(sex, levels = sex_labels_vector)
  expr_df$phenotype <- factor(phenotype, levels = phenotype_labels_vector)
  
  expr_df$group_plot <- paste(expr_df$sex, expr_df$phenotype, sep = "_")
  
  # Long format 
  expr_long <- tidyr::pivot_longer(
    expr_df,
    cols = all_of(genes_present),
    names_to = "Gene",
    values_to = "Expression"
  )
  
  # Plot
  group_plot <- Expression <- NULL
  p <- ggplot(expr_long, aes(x = group_plot, y = Expression, fill = sex)) +
    
    geom_violin(trim = FALSE, alpha = 0.6, color = NA) +
    

    geom_jitter(
     color = "black",
      width = 0.08,
      size = 1,
      alpha = 0.8,
      show.legend = FALSE
    ) +
    
    stat_summary(
      fun = mean,
      geom = "point",
      shape = 21,
      size = 2,
      fill = "red",
      color = "black"
    ) +
    
    facet_wrap(~ Gene, scales = "free_y") +
    labs(
      title = "Expression of Selected Genes Across Sex-Phenotype Groups",
      x = "Sex-Phenotype Groups",
      y = "Expression"
    )  +
    
    theme_classic() +
    
    theme(
      panel.background = element_blank(),  
      plot.background  = element_blank(),
      panel.grid       = element_blank(),   
      strip.background = element_blank(),  
      strip.text       = element_text(size = 11, face = "bold"),
      axis.text.x      = element_text(angle = 45, hjust = 1),
      legend.position  = "none"
    )
  
  return(p)
}


#' Generate dotplot of selected genes across groups
#'
#' Visualizes gene expression across Sex / Phenotype / Cell type combinations.
#'
#' @param seurat_obj A Seurat object
#' @param genes Character vector of genes to plot
#' @param sex Metadata column for sex
#' @param phenotype Metadata column for phenotype/status
#' @param celltype_col Metadata column for cell types
#' @param target_cell_type_flag Either "all" or a specific cell type
#'
#' @return A ggplot object (DotPlot)
#' 
#' @importFrom Seurat Idents "Idents<-"
#' @importFrom Seurat  DotPlot
#' @export
generate_dotplot_sc <- function(seurat_obj,
                             genes,
                             sex = "sex",
                             phenotype = "status",
                             celltype_col = "cell_type",
                             target_cell_type_flag = "all") {
  
  # Validate metadata columns
  required_cols <- c(sex, phenotype, celltype_col)
  
  if (!all(required_cols %in% colnames(seurat_obj@meta.data))) {
    stop("Required metadata columns not found in Seurat object")
  }
  # Validate genes
  if (length(genes) == 0) {
    stop("No genes provided")
  }
  
  genes_present <- genes[genes %in% rownames(seurat_obj)]
  
  if (length(genes_present) == 0) {
    stop("None of the provided genes found in Seurat object")
  }
  
  if (length(genes_present) < length(genes)) {
    warning("Some genes were not found and will be ignored")
  }
  
  # Optional cell type subset
  if (!("all" %in% target_cell_type_flag)) {
    
    missing_ct <- setdiff(
      target_cell_type_flag,
      unique(seurat_obj@meta.data[[celltype_col]])
    )
    
    if (length(missing_ct) > 0) {
      stop(paste(
        "Cell type(s) not found:",
        paste(missing_ct, collapse = ", ")
      ))
    }
    cell_meta <- seurat_obj@meta.data[
      seurat_obj@meta.data[[celltype_col]] %in% target_cell_type_flag, ,
      drop = FALSE
    ]
    
    seurat_obj <- subset(seurat_obj, cells = rownames(cell_meta))
  }
  
  # Create grouping variable
  seurat_obj$group_plot <- paste(
    seurat_obj@meta.data[[celltype_col]],
    seurat_obj@meta.data[[sex]],
    seurat_obj@meta.data[[phenotype]],
    sep = "_"
  )
  
  Idents(seurat_obj) <- "group_plot"
  # Generate dotplot
  p <- Seurat::DotPlot(
    seurat_obj,
    features = genes_present,
    cols = c("blue", "red"),
    assay = "RNA"
  ) +
    coord_flip() +
    labs(
      title =paste0("Dotplot of selected gene expression in ", target_cell_type_flag, " cells"),
      x = "Genes",
      y = "Groups"
    ) +
    theme(
      axis.text.x = element_text(angle = 90, hjust = 1)
    )
  
  return(p)
}



#' Create enrichment dotplots from enrichResult objects
#'
#' @description
#' Generates dotplots from enrichment results. Input must contain enrichResult
#' objects (clusterProfiler / ReactomePA). Invalid entries are skipped with messages.
#' Works for single enrichResult, flat list (DEG types), or nested list (cell type → DEG type).
#'
#' @param enrichment_results A list or nested list of enrichResult objects.
#' @param showCategory Number of categories to display (default = 10).
#' @param ncol Number of columns for arranged plots (default = 2).
#'
#' @return Named list of combined ggplots. For nested input, names correspond to cell types;
#' for flat input, a single `"all"` entry is returned.
#'
#' @importFrom clusterProfiler dotplot
#' @importFrom patchwork wrap_plots
#' @export
plot_enrichment_dotplots <- function(enrichment_results,
                                     showCategory = 10,
                                     ncol = 2) {
  
  if (is.null(enrichment_results) || length(enrichment_results) == 0) {
    message("No enrichment results provided.")
    return(NULL)
  }
  
  # Convert single enrichResult to list
  if (inherits(enrichment_results, "enrichResult")) {
    enrichment_results <- list(Result = enrichment_results)
  }
  
  create_dotplot <- function(res, label, cell_type="") {
    if (!inherits(res, "enrichResult")) {
      message("Skipping ", label, ": not an enrichResult object.")
      return(NULL)
    }
    
    df <- as.data.frame(res)
    if (nrow(df) == 0) {
      message("Skipping ", label, ": enrichResult is empty.")
      return(NULL)
    }
    
    suppressMessages(
      clusterProfiler::dotplot(res, showCategory = showCategory) +
        ggplot2::labs(title = paste("Top enriched pathways -", label, cell_type)) +
        ggplot2::scale_y_discrete(expand = ggplot2::expansion(mult = c(0.1, 0.1))) +
        ggplot2::theme(plot.margin = ggplot2::margin(10, 10, 15, 10)) +
        ggplot2::scale_size_continuous(labels = function(x) as.integer(round(x)))
    )
  }
  
  all_plots <- list()
  
  # Detect nested (cell type → DEG type) or flat
  is_nested <- any(sapply(enrichment_results, is.list))
  
  if (is_nested) {
    for (cell_type in names(enrichment_results)) {
      cell_res <- enrichment_results[[cell_type]]
      if (!is.list(cell_res)) next
      
      cell_plots <- lapply(names(cell_res), function(deg_type) {
        create_dotplot(cell_res[[deg_type]], deg_type, cell_type)
      })
      names(cell_plots) <- names(cell_res)
      
      # Remove NULLs
      cell_plots <- cell_plots[!sapply(cell_plots, is.null)]
      if (length(cell_plots) == 0) next
      
      # Combine plots side by side using patchwork
      all_plots[[cell_type]] <- patchwork::wrap_plots(cell_plots, ncol = ncol)
    }
    
  } else {
    # Flat list (bulk) or single cell
    plots <- lapply(names(enrichment_results), function(deg_type) {
      create_dotplot(enrichment_results[[deg_type]], deg_type)
    })
    plots <- plots[!sapply(plots, is.null)]
    
    if (length(plots) > 0) {
      all_plots[["all"]] <- patchwork::wrap_plots(plots, ncol = ncol)
    }
  }
  
  if (length(all_plots) == 0) {
    message("No valid enrichment plots generated.")
    return(NULL)
  }
  
  return(all_plots)
}


#' Generate volcano plots for categorized DEGs
#'
#' @description
#' Creates volcano plots for all DEG categories (male-specific, female-specific,
#' sex-dimorphic, sex-neutral). Accepts either a single data.frame (bulk) or a
#' list of data.frames per cell type.
#'
#' @param de_results A data.frame or named list of data.frames containing categorized DEGs.
#'                   Required columns: DEG_Type, Gene_Symbols, Male_avg_logFC, Male_FDR,
#'                   Female_avg_logFC, Female_FDR
#' @param top_n Number of top genes to highlight per sex (default = 5).
#' @param logfc_thresh Minimum absolute logFC to consider a gene for highlighting (default = 1.5).
#' @param ncol Number of columns for arranging plots (default = 2).
#'
#' @return Named list of combined volcano plots per cell type or "all" for bulk.
#' @importFrom ggplot2 ggplot geom_point scale_color_manual theme_minimal labs theme element_text
#' @importFrom ggrepel geom_text_repel
#' @importFrom dplyr filter mutate slice_min
#' @importFrom patchwork wrap_plots
#' @export
plot_volcano_deg <- function(de_results,
                             top_n = 5,
                             logfc_thresh = 1.5,
                             ncol = 2) {
  
  create_volcano_plot <- function(df, deg_type, x_col, y_col, title_suffix) {
    
    # prevent warning when checking package
    DEG_Type <- GeneName <- Gene_Symbols <- Male_avg_logFC <- Male_FDR <- Female_avg_logFC <- Female_FDR <- logFC <- Sex <- . <- Significant <- NULL 
    
    df_filtered <- df %>% dplyr::filter(DEG_Type == deg_type) %>%
      dplyr::mutate(GeneName = Gene_Symbols)  # ensure gene names column
    
    # Identify top significant genes per sex
    df_filtered <- df_filtered %>%
      dplyr::mutate(Significant = GeneName %in% c(
        df_filtered %>%
          dplyr::filter(Male_FDR < 0.05, abs(Male_avg_logFC) > logfc_thresh) %>%
          dplyr::slice_min(Male_FDR, n = top_n, with_ties = FALSE) %>%
          .$Gene_Symbols,
        df_filtered %>%
          dplyr::filter(Female_FDR < 0.05, abs(Female_avg_logFC) > logfc_thresh) %>%
          dplyr::slice_min(Female_FDR, n = top_n, with_ties = FALSE) %>%
          .$Gene_Symbols
      ))
    
    # Create mapped columns 
    plot_x <- plot_y <- NULL
    df_filtered$plot_x <- eval(parse(text = x_col), envir = df_filtered)
    
    df_filtered$plot_y <- eval(parse(text = y_col), envir = df_filtered)
    
    
    
    # Create volcano plot
    
    ggplot(df_filtered, aes(x = plot_x, y = plot_y, color = Significant)) +
      geom_point(alpha = 0.6) +
      scale_color_manual(values = c("grey70", "firebrick")) +
      theme_minimal() +
      labs(
        x = "log2 Fold Change",
        y = "-log10(p-value)",
        title = paste("Volcano Plot -", title_suffix, "-", deg_type)
      ) +
      theme(
        legend.position = "bottom",
        plot.title = element_text(size = 10, hjust = 0.5, margin = ggplot2::margin(b = 10))
      ) +
      ggrepel::geom_text_repel(
        data = df_filtered %>% dplyr::filter(Significant),
        aes(label = Gene_Symbols),
        size = 3,
        max.overlaps = 30
      )
  }
  
  all_plots <- list()
  
  # Detect nested input (list of cell types) or bulk (single data.frame)
  is_nested <- is.list(de_results) && !inherits(de_results, "data.frame")
  
  if (is_nested) {
    for (cell_type in names(de_results)) {
      df <- de_results[[cell_type]]
      
      plots <- list(
        create_volcano_plot(df, "female-specific", "Female_avg_logFC", "-log10(Female_FDR)", cell_type),
        create_volcano_plot(df, "male-specific", "Male_avg_logFC", "-log10(Male_FDR)", cell_type),
        create_volcano_plot(df, "sex-dimorphic", "Male_avg_logFC - Female_avg_logFC","-log10(pmin(Male_FDR, Female_FDR))", cell_type),
        create_volcano_plot(df, "sex-neutral","(Male_avg_logFC + Female_avg_logFC)/2", "-log10(pmin(Male_FDR, Female_FDR))", cell_type)
      )
      
      # Remove NULLs
      plots <- plots[!sapply(plots, is.null)]
      if (length(plots) == 0) next
      
      all_plots[[cell_type]] <- patchwork::wrap_plots(plots, ncol = ncol)
    }
    
  } else {
    # Bulk input
    plots <- list(
      create_volcano_plot(de_results, "female-specific", "Female_avg_logFC", "-log10(Female_FDR)", "all"),
      create_volcano_plot(de_results, "male-specific", "Male_avg_logFC", "-log10(Male_FDR)", "all"),
      create_volcano_plot(de_results, "sex-dimorphic", "Male_avg_logFC - Female_avg_logFC","-log10(pmin(Male_FDR, Female_FDR))", "all"),
      create_volcano_plot(de_results, "sex-neutral","(Male_avg_logFC + Female_avg_logFC)/2", "-log10(pmin(Male_FDR, Female_FDR))",  "all")
    )
    
    plots <- plots[!sapply(plots, is.null)]
    if (length(plots) > 0) {
      all_plots[["all"]] <- patchwork::wrap_plots(plots, ncol = ncol)
    }
  }
  
  if (length(all_plots) == 0) {
    message("No valid volcano plots generated.")
    return(NULL)
  }
  
  return(all_plots)
}



