#!/home/bayegy/pipelines/metagenome/miniconda2/envs/fcx_bac_assem/bin/Rscript
rm(list=ls())
suppressPackageStartupMessages(library(ggplot2,quietly = TRUE))
suppressPackageStartupMessages(library(plotly, quietly = TRUE))
suppressPackageStartupMessages(library(tidyr, quietly = TRUE))
suppressPackageStartupMessages(library(RColorBrewer, quietly = TRUE))
suppressPackageStartupMessages(library(ggsci, quietly = TRUE))
suppressPackageStartupMessages(library(ggrepel, quietly = TRUE))
# suppressPackageStartupMessages(library(paletteer, quietly = TRUE))
suppressPackageStartupMessages(library(reshape2, quietly = TRUE))
suppressPackageStartupMessages(library(dplyr, quietly = TRUE))
suppressPackageStartupMessages(library(argparse, quietly = TRUE))
suppressPackageStartupMessages(library(jsonlite))

# 创建一个解析器对象
# 定义参数
parser <- ArgumentParser()
parser$add_argument("-i", "--input_file", type = "character",
                    help = "Input file name")
parser$add_argument("-o", "--outdir", type = "character",
                    help = "Output directory")
parser$add_argument("-a", "--annotion_value", choices = c('top','bottom'),
                    default = "bottom" ,
                    help = "Value for annotion (top or bottom)")
parser$add_argument("-p", "--pad_size", type = "integer",
                  default = 10,
                  help = "Pad size for the plot")
parser$add_argument("-nt", "--node_thickness_size", type = "integer",
                    default = 10,
                    help = "Node thickness size")
parser$add_argument("-lc", "--line_color", type = "character",
                    default='black',
                     help = "Line color")
parser$add_argument("-lw", "--line_width", type = "double",
                    default=0.5,
                    help = "Line width")
parser$add_argument("-lkc", "--link_color", type = "character",
                    default='grey',
                    help = "Link color")
parser$add_argument("-pt", "--p_title", type = "character",
                    default='',
                    help = "Plot title")
parser$add_argument("-pts", "--p_title_size", type = "integer",
                    default='14',
                    help = "Plot title size")
parser$add_argument("-ptc", "--p_title_color", type = "character",
                    default='black',
                    help = "Plot title color")

parser$add_argument("-pa", "--p_arrangement", type = "character",
                    choices = c("snap", "perpendicular", "freeform", "fixed"),
                    default="perpendicular",
                    help = "Plot arrangement")
parser$add_argument("-sn", "--sort_n", type = "integer",
                  default=0,
                  help = "Sort by specific number of species")
parser$add_argument("-rs", "--is_remove_species_class", type = "logical",
                    default=TRUE,
                    help = "Remove species class")
parser$add_argument("-ac", "--is_add_class", type = "logical",
                    default=TRUE,
                    help = "Add species class label")
parser$add_argument("-ans", "--annotion_font_size", type = "integer",
                    default='14',
                    help = "Plot annotion_font_size size")
parser$add_argument("-mp", "--is_mapping", type = "logical",
                    default=FALSE,
                    help = "Add mapping")
parser$add_argument("-cl", "--columns_length", type = "integer",
                    default=6,
                    help = "columns length")
parser$add_argument("-nlt", "--node_label_type", type = "character",
                    choices = c("name", "name_percent", "name_abundance", "name_abundance_percent"),
                    default="name",
                    help = "Node label content")
parser$add_argument("-ps", "--percent_scope", type = "character",
                    choices = c("level", "global"),
                    default="level",
                    help = "Percent scope")
# 解析参数
args <- parser$parse_args()

# 根据解析的参数设置变量
annotion_value <- args$annotion_value
pad_size <- args$pad_size
node_thickness_size <- args$node_thickness_size
line_color <- args$line_color
line_width <- args$line_width
link_color <- args$link_color
p_title <- args$p_title
p_title_size <- args$p_title_size
p_title_color <- args$p_title_color
input_file <- args$input_file
p_arrangement <- args$p_arrangement
sort_n <- args$sort_n
is_remove_species_class <- args$is_remove_species_class
is_add_class <- args$is_add_class
annotion_font_size<-args$annotion_font_size
outdir <- args$outdir
is_mapping<-args$is_mapping
columns_length<-args$columns_length
node_label_type<-args$node_label_type
percent_scope<-args$percent_scope
percent_digits<-4

###以上均是参数默认值
set_annotion_y_size <- function(annotion) {
  # 检查annotion参数的值，并根据条件设置annotion_y_size
  if (annotion == "top") {
    annotion_y_size <- 1
  } else if (annotion == "bottom") {
    annotion_y_size <- -0.1
  } else {
    # 如果annotion不是"top"或"bottom"，可以设置为NA或者返回错误信息
    annotion_y_size <--0.05
    # 或者可以警告用户并返回一个错误信息
    warning("The value of annotion must be either 'top' or 'bottom'.")
  }

  # 返回annotion_y_size的值
  return(annotion_y_size)
}
annotion_y_size=set_annotion_y_size(annotion_value)

palTableau45<-function(n) {
  color_data<-colorRampPalette((pal_igv( "default")(45)))(n)
#   if (n<=45){
#     color_data<-paletteer_d("ggsci::default_igv", n = 45)
#   }else{
#     color_data<-colorRampPalette((pal_igv( "default")(45)))(n)
#   }
  }

sankey_otu_cols= c('kingdom', 'phylum', 'class', 'order', 'family', 'genus','strain')
sankey_otu_cols=sankey_otu_cols[0:columns_length]
species_str_pattern <- "(k|p|c|o|f|g|s)__"
species_str_pattern2<-"(\\[\\])"
# 定义替换函数
replace_string <- function(column, pattern,total_df) {
  total_df[[column]] <- gsub(pattern, "", total_df[[column]], ignore.case = TRUE)
  return(total_df[[column]])
}

##生成
genSankey <- function(df) {
  columns_list <- colnames(df)
  value_cols <- columns_list[length(columns_list)]
  cat_cols <- columns_list[-length(columns_list)]
  for (i in 1:(length(cat_cols)-1)) {
    if (i == 1) {
      sourceTargetDf <- df[, c(cat_cols[i], cat_cols[i+1], value_cols)]
      colnames(sourceTargetDf) <- c('source', 'target', 'count')
    } else {
      tempDf <- df[, c(cat_cols[i], cat_cols[i+1], value_cols)]
      colnames(tempDf) <- c('source', 'target', 'count')
      sourceTargetDf <- rbind(sourceTargetDf, tempDf)
    }
    sourceTargetDf <- aggregate(count ~ source + target, data = sourceTargetDf, sum)
  }

  return(sourceTargetDf)
}
sort_df <- function(df, n) {
  last_column <- colnames(df)[ncol(df)]
  df <- df[order(df[, last_column], decreasing = TRUE), , drop = FALSE]
  if (nrow(df) >= n) {
    df <- df[1:n,]
  } else {
    df<-df
    # 处理列数不足的情况
    # 可以根据需求选择合适的处理方式，比如选择所有列或者填充缺失的列
  }
  return(df)
}
# input_file<-'ss2s.txt'
data <- read.delim(input_file, stringsAsFactors = FALSE, check.names = FALSE)
df<-data

if(is_mapping){
  total_df <-df
  total_df <- total_df[total_df$kingdom != "unclassified", ]

}else{
  df <- separate(df, colnames(df)[length(colnames(df))], into = sankey_otu_cols, sep = ";", extra = "drop")
  abundance_data <- rowSums(data[,2:(ncol(data)-1)])
  subset_df <- df[,sankey_otu_cols]
  total_df <- merge(subset_df, data.frame(abundance = abundance_data), by = 'row.names', all = TRUE)
  total_df <- total_df[,-1]
}
###是否去除”g_“


#pattern <- "^[0-9]+$)"
#is_valid <- grepl(pattern, sort_n)


df_rows=nrow(total_df)
if (sort_n==0){
  total_df<-sort_df(total_df,df_rows)
}else{
  if(sort_n>=df_rows){
    total_df<-sort_df(total_df,df_rows)
  }else{total_df<-sort_df(total_df,sort_n)}
}

if(is_remove_species_class){
  total_df[sankey_otu_cols] <- lapply(sankey_otu_cols, replace_string, species_str_pattern,total_df)
}
abundance_sankeydf<-genSankey(total_df)
node_list2<- reshape2::melt(total_df, id = 'abundance') %>%
  mutate(color = palTableau45(n()))

node_list2 <- node_list2[!duplicated(node_list2$value), ]
abundance_sankeydf$IDsource <- match(abundance_sankeydf$source, node_list2$value) - 1
abundance_sankeydf$IDtarget <- match(abundance_sankeydf$target, node_list2$value) - 1
abundance_sankeydf<-abundance_sankeydf%>%
  mutate(color = palTableau45(n()))

#####为了将最后一列移到右侧，将最后一列转为空白。
add_new_row <- function(node_list2) {
  new_values <- rep('', ncol(node_list2))
  new_row <- data.frame(t(new_values))
  names(new_row) <- colnames(node_list2)
  new_row$color <- '#FFFFFF'
  node_list2 <- rbind(node_list2, new_row)
  return(node_list2)
}

# 定义一个R函数，用于添加新行并移动最后一列到右侧
add_final_row_and_move_last_column <- function(abundance_sankeydf, total_df) {
  # 筛选出满足条件的行
  filtered_df <- abundance_sankeydf[abundance_sankeydf$target %in% total_df$genus,]
  # 选择需要的列并重命名
  selected_df <- filtered_df %>% select(target, IDtarget)
  # 重命名列
  names(selected_df) <- c("source", "IDsource")
  # 添加新列并设置值
  selected_df$target<-''
  selected_df$count <- 1
  selected_df$color <- "#FFFFFF"
  # 计算新的IDtarget值
  new_IDtarget <- max(abundance_sankeydf$IDtarget) + 1
  selected_df$IDtarget <- new_IDtarget

  # 创建一个新的数据框，将原始数据框和新行合并
  abundance_sankeydf2 <- rbind(abundance_sankeydf, selected_df)

  # 将最后一列移动到最右侧
  abundance_sankeydf2 <- abundance_sankeydf2[, rev(names(abundance_sankeydf2))]

  # 返回修改后的数据框
  return(abundance_sankeydf2)
}

# 假设abundance_sankeydf和total_df是已经存在的数据框
# 调用函数
abundance_sankeydf_for_label <- abundance_sankeydf
if (!is_mapping) {
  abundance_sankeydf <- add_final_row_and_move_last_column(abundance_sankeydf, total_df)
  node_list2<-add_new_row(node_list2)
}

value_col_name <- colnames(total_df)[ncol(total_df)]
cat_cols_for_total <- colnames(total_df)[-ncol(total_df)]
valid_rows <- rep(TRUE, nrow(total_df))
if (length(cat_cols_for_total) >= 2) {
  valid_rows <- !is.na(total_df[[cat_cols_for_total[1]]]) &
    total_df[[cat_cols_for_total[1]]] != "" &
    !is.na(total_df[[cat_cols_for_total[2]]]) &
    total_df[[cat_cols_for_total[2]]] != ""
}
total_abundance <- sum(total_df[[value_col_name]][valid_rows], na.rm = TRUE)

node_long_df <- reshape2::melt(
  total_df[valid_rows, c(cat_cols_for_total, value_col_name), drop = FALSE],
  id.vars = value_col_name
)
names(node_long_df) <- c("abundance", "level", "node")
node_long_df <- node_long_df[!is.na(node_long_df$node) & node_long_df$node != "", ]
node_stat_df <- aggregate(abundance ~ level + node, data = node_long_df, sum, na.rm = TRUE)
level_total_df <- aggregate(abundance ~ level, data = node_stat_df, sum, na.rm = TRUE)
names(level_total_df) <- c("level", "level_total")
node_stat_df <- merge(node_stat_df, level_total_df, by = "level", all.x = TRUE)
node_stat_df$percent <- 0
if (percent_scope == "global") {
  node_stat_df$percent <- ifelse(total_abundance > 0, node_stat_df$abundance / total_abundance * 100, 0)
} else {
  node_stat_df$percent <- ifelse(node_stat_df$level_total > 0, node_stat_df$abundance / node_stat_df$level_total * 100, 0)
}
node_stat_df$percent_str <- paste0(formatC(node_stat_df$percent, format = "f", digits = percent_digits), "%")
node_stat_df$abundance_str <- format(signif(node_stat_df$abundance, 4), scientific = FALSE, trim = TRUE)
node_stat_df$label <- node_stat_df$node
if (node_label_type == "name_percent") {
  node_stat_df$label <- paste0(node_stat_df$node, "<br>", node_stat_df$percent_str)
} else if (node_label_type == "name_abundance") {
  node_stat_df$label <- paste0(node_stat_df$node, "<br>", node_stat_df$abundance_str)
} else if (node_label_type == "name_abundance_percent") {
  node_stat_df$label <- paste0(node_stat_df$node, "<br>", node_stat_df$abundance_str, " (", node_stat_df$percent_str, ")")
}

node_list2$key <- paste0(node_list2$variable, "|||", node_list2$value)
node_stat_df$key <- paste0(node_stat_df$level, "|||", node_stat_df$node)
node_list2$label <- node_stat_df$label[match(node_list2$key, node_stat_df$key)]
node_list2$label[is.na(node_list2$label)] <- node_list2$value[is.na(node_list2$label)]
node_list2$label[node_list2$value == ""] <- ""

node_list_txt<-file.path(outdir, paste0( "node_list2.txt"))
write.table(node_list2,file=node_list_txt,sep='\t', row.names = FALSE,quote = FALSE)
abundance_sankeydf_file<-file.path(outdir, paste0( "abundance_sankeydf2.txt"))
write.table(abundance_sankeydf,sep='\t',file=abundance_sankeydf_file,row.names = FALSE,quote = FALSE)
abundance_sankeydf$color <- colorspace::adjust_transparency(abundance_sankeydf$color, alpha = 0.5)

# 将特定的列转换为一个列表（list）
sankey_link_data <- list(
  source = abundance_sankeydf$IDsource,
  target = abundance_sankeydf$IDtarget,
  value = abundance_sankeydf$count,
  color = abundance_sankeydf$color
)
sankey_node_data<-list(
    label = node_list2$label,
    color=node_list2$color,
    pad = pad_size, thickness = node_thickness_size,
    line = list(color = line_color, width = 0)
  )
# 将列表转换为JSON格式的字符串
sankey_link_json_data <- toJSON(sankey_link_data , pretty = TRUE, auto_unbox = TRUE)
sankey_link_data_file<-file.path(outdir, paste0( "sankey_link_data.json"))
# 如果你想将JSON数据保存到文件
write(sankey_link_json_data, sankey_link_data_file, sep = "\n")
sankey_node_json_data <- toJSON(sankey_node_data , pretty = TRUE, auto_unbox = TRUE) # nolint
sankey_node_data_file<-file.path(outdir, paste0( "sankey_node_data.json"))
# 如果你想将JSON数据保存到文件
write(sankey_node_json_data, sankey_node_data_file, sep = "\n")
if (is_add_class){
  formatted_cols = unlist(lapply(sankey_otu_cols, function(col) paste("<b>", col, "<b>")))
  formatted_cols<- c(formatted_cols,"","")
  annotion_list=list(text =formatted_cols,
    x = seq(0, 1, length.out = 7),
    y = annotion_y_size, #below the bottom.  Use >1 for above the top.
    xanchor = "center", #center labels on steps
    showarrow = FALSE,
    font = list(
      size = annotion_font_size
    ))
    annotion_json_data <- toJSON(annotion_list, pretty = TRUE, auto_unbox = TRUE) # nolint
    annotion_json_data_file<-file.path(outdir, paste0( "sankey_annotion.json"))
    # 如果你想将JSON数据保存到文件
    write(annotion_json_data, annotion_json_data_file, sep = "\n")
}
# ploty_p<- plot_ly(
#   type = 'sankey', orientation = 'h',
#   valueformat = ".0f",
#   valuesuffix = "",
#   arrangement=p_arrangement,###snap""perpendicular""freeform""fixed"
#   node = list(
#     label = node_list2$value,
#     color=node_list2$color,
#     pad = pad_size, thickness = node_thickness_size,
#     line = list(color = line_color, width = 0)
#   ),
#   link = list(
#     source =abundance_sankeydf2$IDsource,
#     target = abundance_sankeydf2$IDtarget,
#     value = abundance_sankeydf2$count,
#     color = colorspace::adjust_transparency(abundance_sankeydf2$color, alpha = 0.5)
#   )
# )

# p <-ploty_p %>% layout(
#   title = p_title,
#   font = list(
#     size = p_title_size,color=p_title_color),
#   xaxis = list(showgrid = F, zeroline = F),
#   yaxis = list(showgrid = F, zeroline = F)

# )

# if (is_add_class){
#   formatted_cols = unlist(lapply(sankey_otu_cols, function(col) paste("<b>", col, "<b>")))
#   formatted_cols<- c(formatted_cols,"","")
#   p<-p %>%
#   add_annotations(
#     text =formatted_cols,
#     x = seq(0, 1, length.out = 7),
#     y = annotion_y_size, #below the bottom.  Use >1 for above the top.
#     xanchor = "center", #center labels on steps
#     showarrow = FALSE,
#     font = list(
#       size = annotion_font_size
#     )
#   ) # add step labels
# }


# html_output = file.path(outdir, paste0( "sankey.html"))
# htmlwidgets::saveWidget(p, file =html_output)
# svg_output=file.path(outdir, paste0( "sankey.svg"))
# save_image(p, file=svg_output)