#!/usr/bin/env Rscript


#### use it as :
#### Rscript circultree.R --tree /media/bayegy/disk1/bioin3/cp_database/data/bin_tree.tre --heat /media/bayegy/disk1/bioin3/cp_database/data/heat.txt --abundance /media/bayegy/disk1/bioin3/cp_database/data/bin_abundance_taxonomy.xls --mapping-file /media/bayegy/disk1/bioin3/cp_database/data/分组信息表.tsv --type /media/bayegy/disk1/bioin3/cp_database/data/All.CARD.Drug.Class.txt --category "Category1" --class "p__" --top 4 --piecol "#FF7900,#7EB500" --heatcol "#EEEEEE,#0072CC,#DC143C" --typecol "#EEEEEE,#7EB500,#DC143C" --classcol "Paired" --scale TRUE --out-dir /media/bayegy/disk1/bioin3/cp_database

library(ggtree)
library(ggplot2)
library(dplyr)
library(ape)
library(argparse)
library(reshape2)
library(scatterpie)
library(ggtreeExtra)
library(treeio)
library(ggnewscale)
library(ggpp)
library("stringr")
library(RColorBrewer)
library(plyr)
library(patchwork)
library(grid)
library(ggrepel)
library(svglite)
library(cowplot)
library(ggtext)


parser <- ArgumentParser()
parser$add_argument("--tree")
parser$add_argument("--heat")
parser$add_argument("--abundance")
parser$add_argument("--mapping-file")
parser$add_argument("--type")
parser$add_argument("--category")
parser$add_argument("--class", default = "p__") ### 可选：k__,p__,等
parser$add_argument("--top", default = 4)
parser$add_argument("--bin-top", default = 20, required = FALSE)
parser$add_argument("--piecol")
parser$add_argument("--heatcol", default = "#EEEEEE,#0072CC,#DC143C", required = FALSE)
parser$add_argument("--typecol", default = "#EEEEEE,#7EB500,#DC143C", required = FALSE)
parser$add_argument("--classcol", default = "Paired", required = FALSE) # 可选：Set3，Paired等12色RColorBrewer调色盘
parser$add_argument("--scale", default = FALSE) # 是否绘制比例尺到进化树上
parser$add_argument("--out-dir", default = "./")

### 调整环形图位置,字号的参数：

parser$add_argument("--cir1", default = -0.5, required = FALSE) # 进化树分类标签位置，内圈
parser$add_argument("--cir2", default = 1, required = FALSE) # 饼图位置
parser$add_argument("--cir3", default = 0.08, required = FALSE) # 热图位置1
parser$add_argument("--cir4", default = 0.03, required = FALSE) # 热图位置2，见p3绘制过程
parser$add_argument("--cir5", default = 0.08, required = FALSE) # type number 热图位置
parser$add_argument("--cir6", default = 0.15, required = FALSE) # 分类瓦片位置
parser$add_argument("--cir7", default = 0.2, required = FALSE) # 进化树分类标签位置，外圈

parser$add_argument("--heatwidth1", default = 0.1, required = FALSE) # 热图宽度
parser$add_argument("--heatwidth2", default = 0.25, required = FALSE) # type number 热图宽度
parser$add_argument("--heatwidth3", default = 0.05, required = FALSE) # 分类瓦片宽度


parser$add_argument("--tree-size", default = 0.1, required = FALSE) # 进化树线段size
parser$add_argument("--size1", default = 2, required = FALSE) # 进化树label标签字体大小
parser$add_argument("--size2", default = 2, required = FALSE) # 进化树分类标签字体大小，内圈
parser$add_argument("--size3", default = 2, required = FALSE) # 热图字体大小
parser$add_argument("--size4", default = 2, required = FALSE) # type number 热图字体大小
parser$add_argument("--size5", default = 3, required = FALSE) ## 进化树分类标签字体大小，外圈
parser$add_argument("--size6", default = 8, required = FALSE) ## 图例字体大小
parser$add_argument("--size7", default = 0.5, required = FALSE) ## 图例边框宽度


parser$add_argument("--rect", default = 0.5, required = FALSE) ## 图例方块大小
parser$add_argument("--width", default = 0.8, required = FALSE) ## 图例下划线长度

args <- parser$parse_args()

tree <- as.character(args$tree)
heat <- as.character(args$heat)
abundance <- as.character(args$abundance)
mapping_file <- as.character(args$mapping_file)
type <- as.character(args$type)
category <- as.character(args$category)
class <- as.character(args$class)
top <- as.numeric(args$top)
bin_top <- as.numeric(args$bin_top)
piecol <- as.character(args$piecol)
heatcol <- as.character(args$heatcol)
typecol <- as.character(args$typecol)
classcol <- as.character(args$classcol)
scale <- as.logical(args$scale)
out_dir <- as.character(args$out_dir)

cir1 <- as.numeric(args$cir1)
cir2 <- as.numeric(args$cir2)
cir3 <- as.numeric(args$cir3)
cir4 <- as.numeric(args$cir4)
cir5 <- as.numeric(args$cir5)
cir6 <- as.numeric(args$cir6)
cir7 <- as.numeric(args$cir7)


heatwidth1 <- as.numeric(args$heatwidth1)
heatwidth2 <- as.numeric(args$heatwidth2)
heatwidth3 <- as.numeric(args$heatwidth3)

tree_size <- as.numeric(args$tree_size)
size1 <- as.numeric(args$size1)
size2 <- as.numeric(args$size2)
size3 <- as.numeric(args$size3)
size4 <- as.numeric(args$size4)
size5 <- as.numeric(args$size5)
size6 <- as.numeric(args$size6)
size7 <- as.numeric(args$size7)

width <- as.numeric(args$width)
rect <- as.numeric(args$rect)


tree <- read.tree(tree)
heat <- read.csv(heat, sep = "\t", check.names = FALSE)
abundance <- read.csv(abundance, sep = "\t", check.names = FALSE)
if (length(abundance[, 1]) < bin_top) {
  bin_top <- length(abundance[, 1])
}

# 识别数值型列（自动排除非数值列）
numeric_cols <- sapply(abundance, is.numeric)

# 计算行总和（仅数值列）
abundance$Total <- rowSums(abundance[, numeric_cols, drop = FALSE], na.rm = TRUE)

# 按总和排序并取前10
abundance <- abundance[order(-abundance$Total), ] %>% # 降序排序
  head(bin_top) %>% # 取前n行丰度MAG ID
  subset(select = -Total) # 移除临时列


bin_top_name <- abundance[, 1]

# 获取树中所有tip标签
tree_tips <- tree$tip.label

# 找出bin_top_name中存在于树的tip标签
valid_names <- bin_top_name[bin_top_name %in% tree_tips]


err_fun <- function(valid_names) {
  if (length(valid_names) == 0) {
    stop(paste("丰度前", bin_top, "个bin均未包含在进化树中，请重新选择阈值"))
  }
}

result <- tryCatch(
  err_fun(valid_names),
  error = function(e) {
    message("报错: ", e$message)
    return(NULL)
  }
)


# 检查是否有被过滤的名称
if (length(valid_names) < length(bin_top_name)) {
  removed_names <- setdiff(bin_top_name, tree_tips)
  warning(paste(
    "以下名称不存在于树中，已被过滤:",
    paste(removed_names, collapse = ", ")
  ))
}

# 只保留有效的tip
tree <- keep.tip(tree, valid_names)


heat <- subset(heat, heat[, 1] %in% valid_names)
colnames(heat)[1] <- "MAG ID"
abundance <- subset(abundance, abundance[, 1] %in% valid_names)


mapping_file <- read.csv(mapping_file, sep = "\t", check.names = FALSE)
remove_empty_rows <- function(data, column_name) {
  data[!is.na(data[[column_name]]) &
    trimws(data[[column_name]]) != "", ]
}
mapping_file <- remove_empty_rows(mapping_file, category)

type <- read.csv(type, sep = "\t", check.names = FALSE)
type <- type %>% select(1, where(is.numeric))
type <- type %>%
  select(1, any_of(valid_names))


piecol <- strsplit(piecol, split = ",")[[1]]
heatcol <- strsplit(heatcol, split = ",")[[1]]
typecol <- strsplit(typecol, split = ",")[[1]]

### 排序后提取前n
type_top <- type[order(rowSums(type[, -c(1)]), decreasing = TRUE), ][1:top, ]


# 将表达量矩阵转换为长格式
abundance_long <- melt(abundance, id.vars = colnames(abundance)[1], variable.name = "#SampleID", value.name = "Abundance")
abundance_long <- merge(abundance_long, mapping_file, by = "#SampleID")
heat_long <- melt(heat, id.vars = colnames(heat)[1], variable.name = "type", value.name = "value")
type_long <- melt(type_top, id.vars = colnames(type)[1], variable.name = "ID", value.name = "value_type")


# 丰度分组
abundance_category <- aggregate(as.numeric(abundance_long$Abundance), by = list(label = abundance_long$`#OTU ID`, category = abundance_long[, category]), sum)

#### 合并分类数据到树状图
#### 根据选择类型分割得到分类信息

p1 <- ggtree(tree)
class_data <- data.frame(ID = abundance$`#OTU ID`, class = str_split_fixed(str_split_fixed(abundance$`Consensus Lineage`, class, 2)[, 2], ";", 2)[, 1])
p1_class_data <- merge(p1$data, class_data, by.x = "label", by.y = "ID")
p1_class_data <- p1_class_data[order(p1_class_data$y), ]


data1 <- aggregate(p1_class_data$node, by = list(type = p1_class_data$class), FUN = first)
data2 <- aggregate(p1_class_data$parent, by = list(type = p1_class_data$class), FUN = last)
p1_class_data$jiaodu <- p1_class_data$angle
#### class中只出现一次的分类放到树状图上显示
data3 <- aggregate(p1_class_data$y, by = list(type = p1_class_data$class), FUN = median)
data3$x <- as.integer(data3$x)
data4 <- aggregate(p1_class_data$y, by = list(type = p1_class_data$class), FUN = length)

data_out <- data3[data3$type %in% data4[data4$x != 1, ]$type, ]
data_in <- data3[data3$type %in% data4[data4$x == 1, ]$type, ]

result_list <- list()
for (label in unique(abundance_category$label)) {
  # 筛选当前 label 的数据
  subset_data <- abundance_category %>% filter(label == !!label)
  # 计算 C 和 T1 的总和
  total <- sum(subset_data$x)
  # 计算比例
  subset_data <- subset_data %>%
    mutate(proportion = x / total)
  # 如果总和为 0，跳过绘图
  if (sum(subset_data$x) == 0) {
    next
  }
  subset_data
  p <- ggplot(subset_data, aes(x = category, y = proportion, fill = category)) +
    geom_bar(stat = "identity", width = 1) +
    theme_minimal() +
    coord_polar() +
    theme_void() +
    theme(legend.position = "none") +
    scale_fill_manual(values = piecol)

  result_list[[label]] <- p
}


# 将结果列表转换为数据框
result_df <- tibble::tibble(
  label = names(result_list),
  pies = result_list
)

label_in <- p1_class_data[p1_class_data$y %in% data_in$x, c("label", "class", "jiaodu")]
label_out <- p1_class_data[p1_class_data$y %in% data_out$x, c("label", "class", "jiaodu")]

p1 <- ggtree(tree, layout = "circular", size = tree_size) + geom_tiplab(align = T, size = size1)
if (length(label_in$label) > 0) {
  p1 <- p1 + new_scale_fill() +
    geom_fruit(
      data = label_in,
      geom = geom_text,
      mapping = aes(
        x = 0,
        y = label,
        label = class
      ),
      stat = "identity",
      offset = cir1,
      size = size2,
      fontface = "bold",
      direction = "x"
    )
}

if (scale) {
  p1 <- p1 + geom_treescale(x = 0, y = 0.5, label = "Tree scale:", width = 1)
}


p2 <- p1 + new_scale_fill() +
  geom_fruit(
    data = result_df,
    geom = geom_plot,
    mapping = aes(y = label, label = pies),
    vp.width = 0.02, vp.height = 0.02, hjust = 0.5, vjust = 0.5,
    offset = cir2
  )

p3 <- p2
for (i in 1:length(unique(heat_long$type))) {
  if (i == 1) {
    p3 <- p3 +
      new_scale_fill() +
      geom_fruit(
        data = heat_long[heat_long$type == unique(heat_long$type)[i], ],
        geom = geom_tile,
        mapping = aes(y = `MAG ID`, x = type, fill = value, width = 0.05),
        stat = "identity",
        offset = cir3,
        pwidth = heatwidth1, # 调整此值以控制热图宽度
        axis.params = list(
          axis = "x",
          line.color = "white",
          text.size = size3,
          nbreak = 2,
          text.angle = -90,
          vjust = 0,
          hjust = 0
        )
      ) + guides(fill = "none") + scale_fill_gradientn(
        name = "Abundance",
        colours = heatcol
      )
  } else {
    p3 <- p3 +
      new_scale_fill() +
      geom_fruit(
        data = heat_long[heat_long$type == unique(heat_long$type)[i], ],
        geom = geom_tile,
        mapping = aes(y = `MAG ID`, x = type, fill = value, width = 0.05),
        stat = "identity",
        offset = cir4,
        pwidth = heatwidth1, # 调整此值以控制热图宽度
        axis.params = list(
          axis = "x",
          line.color = "white",
          text.size = size3,
          nbreak = 2,
          text.angle = -90,
          vjust = 0,
          hjust = 0
        )
      ) + guides(fill = "none") + scale_fill_gradientn(
        name = "Abundance",
        colours = heatcol
      )
  }
}

### 添加ARG分组热图

p4 <- p3 +
  new_scale_fill() +
  geom_fruit(
    data = type_long,
    geom = geom_tile,
    mapping = aes(y = ID, x = type_long[, 1], fill = value_type, width = 0.1),
    axis.params = list(
      axis = "x",
      line.color = "white",
      text.size = size4,
      nbreak = 2,
      text.angle = -90,
      vjust = 0,
      hjust = 0
    ),
    offset = cir5,
    pwidth = heatwidth2, # 调整此值以控制热图宽度
  ) +
  guides(fill = "none") + scale_fill_gradientn(
    name = "Type number",
    colours = typecol
  )

### 外圈颜色，较少时固定颜色，太多时随机颜色。

mycol <- brewer.pal(n = 12, name = classcol)



if (length(data2$type) > length(mycol)) {
  mycol <- rainbow(length(data2$type))
}
p5 <- p4 + new_scale_fill() + geom_fruit(
  data = p1_class_data[, c("label", "class")],
  geom = geom_tile,
  mapping = aes(y = label, fill = class),
  stat = "identity",
  offset = cir6,
  width = heatwidth3,
  show.legend = FALSE
)


p6 <- p5 + new_scale_fill() +
  geom_fruit(
    data = p1_class_data[p1_class_data$y %in% data_out$x, c("label", "class", "jiaodu")],
    geom = geom_text_repel,
    mapping = aes(
      x = 0,
      y = label,
      label = class,
      angle = jiaodu - 90
    ),
    stat = "identity",
    offset = cir7,
    size = size5,
    fontface = "bold",
    direction = "x"
  )


### 添加图例

# 空白画布
plot_no <- ggplot() +
  theme_void()
plot_no <- plot_no + inset_element(p6, left = 0.1, bottom = 0, right = 1, top = 1)

#### 循环添加图例
for (i in 0:length(unique(heat_long$type)) + 1) {
  if (i < length(unique(heat_long$type)) + 1) {
    legend_grob <- grobTree(
      rectGrob(gp = gpar(fill = NA, col = NA)),
      # 绘制表示最小值白色矩形
      rectGrob(x = unit(0.1, "npc"), y = unit(0.7, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = heatcol[1], lwd = size7)),
      # 绘制表示中位数蓝色矩形
      rectGrob(x = unit(0.1, "npc"), y = unit(0.45, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = heatcol[2], lwd = size7)),
      # 绘制表示最大值红色矩形
      rectGrob(x = unit(0.1, "npc"), y = unit(0.2, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = heatcol[3], lwd = size7)),
      # 表示最小值的文本
      textGrob(as.character(min(heat_long[heat_long$type == unique(heat_long$type)[i], c("value")])), x = unit(0.5, "npc"), y = unit(0.7, "npc"), gp = gpar(fontsize = size6)),
      # 表示中位数的文本
      textGrob(as.character(median(heat_long[heat_long$type == unique(heat_long$type)[i], c("value")])), x = unit(0.5, "npc"), y = unit(0.45, "npc"), gp = gpar(fontsize = size6)),
      # 表示最大值的文本
      textGrob(as.character(max(heat_long[heat_long$type == unique(heat_long$type)[i], c("value")])), x = unit(0.5, "npc"), y = unit(0.2, "npc"), gp = gpar(fontsize = size6)),
      # 带有下划线的图例标题
      textGrob(paste(unique(heat_long$type)[i], " number", sep = ""), x = unit(0, "npc"), y = unit(1.0, "npc"), gp = gpar(fontface = "bold", fontsize = size6), just = "left"),
      linesGrob(x = unit(c(0, width), "npc"), y = unit(0.9, "npc"), gp = gpar(lwd = 1.5))
    )
    plot_no <- plot_no + inset_element(legend_grob, left = 0, bottom = 1.0 - (0.15 * i), right = 0.15, top = 1.15 - (0.15 * i))
  }

  if (i == length(unique(heat_long$type)) + 1) {
    legend_grob <- grobTree(
      rectGrob(gp = gpar(fill = NA, col = NA)),
      rectGrob(x = unit(0.1, "npc"), y = unit(0.7, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = typecol[1], lwd = size7)),
      rectGrob(x = unit(0.1, "npc"), y = unit(0.45, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = typecol[2], lwd = size7)),
      rectGrob(x = unit(0.1, "npc"), y = unit(0.2, "npc"), width = unit(rect, "cm"), height = unit(rect, "cm"), gp = gpar(fill = typecol[3], lwd = size7)),
      textGrob(as.character(min(type_long$value_type)), x = unit(0.5, "npc"), y = unit(0.7, "npc"), gp = gpar(fontsize = size6)),
      textGrob(as.character(median(type_long$value_type)), x = unit(0.5, "npc"), y = unit(0.45, "npc"), gp = gpar(fontsize = size6)),
      textGrob(as.character(max(type_long$value_type)), x = unit(0.5, "npc"), y = unit(0.2, "npc"), gp = gpar(fontsize = size6)),
      # 带有下划线的图例标题
      textGrob("Type numper", x = unit(0, "npc"), y = unit(1.0, "npc"), gp = gpar(fontface = "bold", fontsize = size6), just = "left"),
      linesGrob(x = unit(c(0, width), "npc"), y = unit(0.9, "npc"), gp = gpar(lwd = 1.5))
    )
    plot_no1 <- plot_no + inset_element(legend_grob, left = 0, bottom = 1.0 - (0.15 * i), right = 0.15, top = 1.15 - (0.15 * i))
  }
}


##### 单独绘制一饼图图例

create_legend_grob <- function(labels, colors,
                               rect_size = rect, # 方块大小
                               text_size = size6, # 字体大小
                               spacing = -0.25) { # 行间距

  n <- length(labels)
  grobs <- list()

  # 动态计算布局参数
  y_pos <- rev(seq(0.7, by = spacing, length.out = n)) # 从上到下排列

  for (i in 1:n) {
    # 颜色方块
    if (i == 1) {
      title_grob <- textGrob(category, x = unit(0, "npc"), y = unit(1.0, "npc"), gp = gpar(fontface = "bold", fontsize = size6), just = "left")
      lines_grob <- linesGrob(x = unit(c(0, width), "npc"), y = unit(0.9, "npc"), gp = gpar(lwd = 1.5))
      grobs <- c(grobs, list(title_grob, lines_grob))
    }
    rect_grob <- rectGrob(
      x = unit(0.1, "npc"),
      y = unit(y_pos[i], "npc"),
      width = unit(rect_size, "cm"),
      height = unit(rect_size, "cm"),
      gp = gpar(fill = colors[i], col = "black", lwd = size7)
    )

    # 文本标签
    text_grob <- textGrob(
      labels[i],
      x = unit(0.5, "npc"), # 文字右对齐
      y = unit(y_pos[i], "npc"),
      just = "left",
      gp = gpar(fontsize = text_size)
    )

    grobs <- c(grobs, list(rect_grob, text_grob))
  }

  # 组合所有元素
  grobTree(
    rectGrob(gp = gpar(fill = "white", col = NA)), # 背景（可选）
    do.call(gList, grobs)
  )
}

legend_grob <- create_legend_grob(
  labels = unique(mapping_file[, category]),
  colors = piecol
)


##### 添加饼图图例：
plot_result <- plot_no1 + inset_element(legend_grob, left = 0.15, bottom = 0.85, right = 0.3, top = 1.0)



plot_svg <- paste(out_dir, "tree_heatmap.svg", sep = "/")
svglite(plot_svg)
plot_result
dev.off()

# write.csv(heat,file=paste(out_dir,"sum_counts.xls",sep="/"),sep="\t")
# write.csv(type,file=paste(out_dir,"CARD.txt",sep="/"),sep="\t")
# write.csv(abundance,file=paste(out_dir,"bin_abundance_taxonomy.xls",sep="/"),sep="\t")
# write.tree(tree,file=paste(out_dir,"bin_tree.tre",sep="/"),append=FALSE,digits=2,tree.names=FALSE)