#!/usr/bin/env Rscript
# Ka/Ks Analysis Visualization - Enhanced Version
# Ka/Ks分析可视化脚本（增强版）

suppressPackageStartupMessages({
  library(ggplot2)
  library(dplyr)
  library(tidyr)
})

# 参数解析
args <- commandArgs(trailingOnly = TRUE)
if (length(args) < 2) {
  cat("Usage: Rscript plot_kaks.R <kaks_results.csv> <output_prefix>\n")
  quit(status = 1)
}

input_file <- args[1]
output_prefix <- args[2]

# 读取数据
cat("Reading Ka/Ks results...\n")
kaks_data <- read.csv(input_file)

# 创建短基因对ID，完整基因对写入映射表，避免图中长标签重叠。
kaks_data$pair_simple <- paste0("P", 1:nrow(kaks_data))
kaks_data$pair_label <- paste0(kaks_data$gene1, " vs ", kaks_data$gene2)

# ==================== 1. Ka/Ks分布直方图 ====================
cat("Generating Ka/Ks distribution plot...\n")
p1 <- ggplot(kaks_data, aes(x = Ka.Ks)) +
  geom_histogram(bins = 30, fill = "steelblue", color = "white", alpha = 0.7) +
  geom_vline(xintercept = 1, linetype = "dashed", color = "red", size = 1) +
  geom_vline(xintercept = mean(kaks_data$Ka.Ks, na.rm = TRUE), 
             linetype = "dotted", color = "darkgreen", size = 1) +
  annotate("text", x = 1, y = Inf, label = "Neutral\n(Ka/Ks=1)", 
           color = "red", vjust = 1.5, hjust = -0.1, size = 3) +
  annotate("text", x = mean(kaks_data$Ka.Ks, na.rm = TRUE), y = Inf, 
           label = paste("Mean\n", round(mean(kaks_data$Ka.Ks, na.rm = TRUE), 2)), 
           color = "darkgreen", vjust = 1.5, hjust = -0.1, size = 3) +
  labs(
    title = "Ka/Ks Distribution",
    subtitle = paste("Mean Ka/Ks =", round(mean(kaks_data$Ka.Ks, na.rm = TRUE), 4)),
    x = "Ka/Ks Ratio",
    y = "Count",
    caption = "Red dashed line: Ka/Ks = 1 (neutral selection)\nGreen dotted line: Mean Ka/Ks"
  ) +
  theme_minimal(base_size = 12) +
  theme(
    plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
    plot.subtitle = element_text(hjust = 0.5, size = 11),
    panel.grid.minor = element_blank()
  )

# ==================== 2. Ka vs Ks散点图 ====================
cat("Generating Ka vs Ks scatter plot...\n")
p2 <- ggplot(kaks_data, aes(x = Ks, y = Ka)) +
  geom_point(aes(color = Ka.Ks), size = 4, alpha = 0.8) +
  geom_text(aes(label = pair_simple), size = 2.5, vjust = -1, hjust = 0.5) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", color = "red", size = 1) +
  annotate("text", x = max(kaks_data$Ks) * 0.9, y = max(kaks_data$Ks) * 0.9, 
           label = "Ka=Ks\n(neutral)", color = "red", size = 3) +
  scale_color_gradient2(
    low = "blue", mid = "green", high = "red",
    midpoint = 1, name = "Ka/Ks"
  ) +
  labs(
    title = "Ka vs Ks Scatter Plot",
    subtitle = paste("n =", nrow(kaks_data), "gene pairs"),
    x = "Ks (Synonymous substitutions)",
    y = "Ka (Non-synonymous substitutions)",
    caption = "Red dashed line: Ka = Ks (neutral selection)\nLabels indicate gene pair ID (P1, P2, ...)"
  ) +
  theme_minimal(base_size = 12) +
  theme(
    plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
    plot.subtitle = element_text(hjust = 0.5, size = 11),
    legend.position = "right"
  )

# ==================== 3. Ks分布图 ====================
cat("Generating Ks distribution plot...\n")
p3 <- ggplot(kaks_data, aes(x = Ks)) +
  geom_histogram(bins = 30, fill = "coral", color = "white", alpha = 0.7) +
  geom_density(aes(y = ..count.. * 0.1), color = "darkred", size = 1) +
  geom_vline(xintercept = mean(kaks_data$Ks, na.rm = TRUE), 
             linetype = "dotted", color = "darkblue", size = 1) +
  annotate("text", x = mean(kaks_data$Ks, na.rm = TRUE), y = Inf,
           label = paste("Mean Ks =", round(mean(kaks_data$Ks, na.rm = TRUE), 1)),
           color = "darkblue", vjust = 1.5, hjust = -0.1, size = 3) +
  labs(
    title = "Ks Distribution",
    subtitle = paste("Mean Ks =", round(mean(kaks_data$Ks, na.rm = TRUE), 4)),
    x = "Ks (Synonymous substitutions per site)",
    y = "Count"
  ) +
  theme_minimal(base_size = 12) +
  theme(
    plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
    plot.subtitle = element_text(hjust = 0.5, size = 11)
  )

# ==================== 4. 分歧时间分布图 ====================
if ("divergence_time_mya" %in% names(kaks_data)) {
  cat("Generating divergence time plot...\n")
  p4 <- ggplot(kaks_data, aes(x = divergence_time_mya)) +
    geom_histogram(bins = 30, fill = "purple", color = "white", alpha = 0.7) +
    geom_vline(xintercept = mean(kaks_data$divergence_time_mya, na.rm = TRUE), 
               linetype = "dotted", color = "darkgreen", size = 1) +
    annotate("text", x = mean(kaks_data$divergence_time_mya, na.rm = TRUE), y = Inf,
             label = paste("Mean =", round(mean(kaks_data$divergence_time_mya, na.rm = TRUE), 0), "Mya"),
             color = "darkgreen", vjust = 1.5, hjust = -0.1, size = 3) +
    labs(
      title = "Divergence Time Distribution",
      subtitle = paste("Mean =", round(mean(kaks_data$divergence_time_mya, na.rm = TRUE), 2), "Mya"),
      x = "Divergence Time (Million years ago)",
      y = "Count"
    ) +
    theme_minimal(base_size = 12) +
    theme(
      plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
      plot.subtitle = element_text(hjust = 0.5, size = 11)
    )
}

# ==================== 5. 每个基因对的Ka/Ks柱状图 ====================
cat("Generating individual Ka/Ks bar plot...\n")
p5 <- ggplot(kaks_data, aes(x = reorder(pair_simple, Ka.Ks), y = Ka.Ks, fill = Ka.Ks)) +
  geom_bar(stat = "identity", color = "white", width = 0.7) +
  geom_hline(yintercept = 1, linetype = "dashed", color = "red", size = 1) +
  geom_text(aes(label = round(Ka.Ks, 2)), vjust = -0.5, size = 3, fontface = "bold") +
  annotate("text", x = 1, y = 1.15, label = "Neutral (Ka/Ks=1)", 
           color = "red", hjust = 0, size = 3) +
  scale_fill_gradient2(
    low = "blue", mid = "green", high = "red",
    midpoint = 1, name = "Ka/Ks"
  ) +
  labs(
    title = "Ka/Ks Values for Each Gene Pair",
    subtitle = paste("Mean Ka/Ks =", round(mean(kaks_data$Ka.Ks, na.rm = TRUE), 4)),
    x = "Gene Pair ID (see kaks_gene_pairs.csv)",
    y = "Ka/Ks Ratio",
    caption = "Red dashed line: Ka/Ks = 1 (neutral selection)\nBars sorted by Ka/Ks value\nValues shown on top of bars"
  ) +
  theme_minimal(base_size = 12) +
  theme(
    plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
    plot.subtitle = element_text(hjust = 0.5, size = 11),
    axis.text.x = element_text(angle = 45, hjust = 1, size = 9),
    legend.position = "right"
  )

# ==================== 6. 每个基因对的分歧时间柱状图 ====================
if ("divergence_time_mya" %in% names(kaks_data)) {
  cat("Generating individual divergence time bar plot...\n")
  p6 <- ggplot(kaks_data, aes(x = reorder(pair_simple, divergence_time_mya), 
                              y = divergence_time_mya, fill = divergence_time_mya)) +
    geom_bar(stat = "identity", color = "white", width = 0.7) +
    geom_text(aes(label = paste0(round(divergence_time_mya, 0), "")), 
              vjust = -0.5, size = 3, fontface = "bold") +
    scale_fill_gradient(low = "lightblue", high = "darkblue", name = "Mya") +
    labs(
      title = "Divergence Time for Each Gene Pair",
      subtitle = paste("Mean =", round(mean(kaks_data$divergence_time_mya, na.rm = TRUE), 0), "Mya"),
      x = "Gene Pair ID (see kaks_gene_pairs.csv)",
      y = "Divergence Time (Million years ago)",
      caption = "Bars sorted by divergence time\nValues shown on top of bars (Mya)"
    ) +
    theme_minimal(base_size = 12) +
    theme(
      plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
      plot.subtitle = element_text(hjust = 0.5, size = 11),
      axis.text.x = element_text(angle = 45, hjust = 1, size = 9),
      legend.position = "right"
    )
}

# ==================== 7. Ka和Ks对比柱状图 ====================
cat("Generating Ka and Ks comparison bar plot...\n")
kaks_long <- kaks_data %>%
  select(pair_simple, Ka, Ks) %>%
  pivot_longer(cols = c(Ka, Ks), names_to = "Type", values_to = "Value")

p7 <- ggplot(kaks_long, aes(x = pair_simple, y = Value, fill = Type)) +
  geom_bar(stat = "identity", position = "dodge", color = "white", width = 0.8) +
  scale_fill_manual(values = c("Ka" = "coral", "Ks" = "steelblue"), name = "") +
  labs(
    title = "Ka and Ks Comparison",
    subtitle = "Non-synonymous vs Synonymous substitutions",
    x = "Gene Pair ID (see kaks_gene_pairs.csv)",
    y = "Substitution Rate",
    caption = "Ka (coral): Non-synonymous\nKs (blue): Synonymous"
  ) +
  theme_minimal(base_size = 12) +
  theme(
    plot.title = element_text(hjust = 0.5, face = "bold", size = 14),
    plot.subtitle = element_text(hjust = 0.5, size = 11),
    axis.text.x = element_text(angle = 45, hjust = 1, size = 9),
    legend.position = "top"
  )

# ==================== 8. 基因对信息表格图 ====================
cat("Generating gene pair information table...\n")
# 创建表格数据
table_data <- kaks_data %>%
  select(pair_simple, pair_label, gene1, gene2, Ka.Ks, divergence_time_mya) %>%
  mutate(
    Ka.Ks = round(Ka.Ks, 3),
    divergence_time_mya = round(divergence_time_mya, 0)
  )

# 保存表格为CSV
write.csv(table_data, paste0(output_prefix, ".kaks_gene_pairs.csv"), row.names = FALSE)

# ==================== 保存所有图表 ====================
cat("Saving PDF...\n")
pdf_file <- paste0(output_prefix, ".kaks_plots.pdf")
pdf(pdf_file, width = 12, height = 8)
print(p1)
print(p2)
print(p3)
if (exists("p4")) {
  print(p4)
}
print(p5)
if (exists("p6")) {
  print(p6)
}
print(p7)
dev.off()

# 保存PNG（每个图单独保存）
cat("Saving PNG files...\n")

# 图1: Ka/Ks分布
png(paste0(output_prefix, ".kaks_distribution.png"), width = 800, height = 600, res = 100)
print(p1)
dev.off()

# 图2: Ka vs Ks散点
png(paste0(output_prefix, ".kaks_scatter.png"), width = 800, height = 600, res = 100)
print(p2)
dev.off()

# 图3: Ks分布
png(paste0(output_prefix, ".ks_distribution.png"), width = 800, height = 600, res = 100)
print(p3)
dev.off()

# 图4: 分歧时间分布
if (exists("p4")) {
  png(paste0(output_prefix, ".divergence_distribution.png"), width = 800, height = 600, res = 100)
  print(p4)
  dev.off()
}

# 图5: 每个基因对的Ka/Ks柱状图
png(paste0(output_prefix, ".kaks_barplot.png"), width = 1000, height = 600, res = 100)
print(p5)
dev.off()

# 图6: 每个基因对的分歧时间柱状图
if (exists("p6")) {
  png(paste0(output_prefix, ".divergence_barplot.png"), width = 1000, height = 600, res = 100)
  print(p6)
  dev.off()
}

# 图7: Ka和Ks对比
png(paste0(output_prefix, ".ka_ks_comparison.png"), width = 1000, height = 600, res = 100)
print(p7)
dev.off()

cat("Plots saved successfully!\n")
cat("PDF:", pdf_file, "\n")
cat("PNG files:\n")
cat("  -", paste0(output_prefix, ".kaks_distribution.png"), "\n")
cat("  -", paste0(output_prefix, ".kaks_scatter.png"), "\n")
cat("  -", paste0(output_prefix, ".ks_distribution.png"), "\n")
if (exists("p4")) {
  cat("  -", paste0(output_prefix, ".divergence_distribution.png"), "\n")
}
cat("  -", paste0(output_prefix, ".kaks_barplot.png"), "\n")
if (exists("p6")) {
  cat("  -", paste0(output_prefix, ".divergence_barplot.png"), "\n")
}
cat("  -", paste0(output_prefix, ".ka_ks_comparison.png"), "\n")
cat("CSV table:", paste0(output_prefix, ".kaks_gene_pairs.csv"), "\n")