#  parse command line arguments
library(ComplexHeatmap)
library(circlize)
library(ggplot2)
library(dplyr)
library(tidyr)
library(stringr)
library(RColorBrewer)
library(ragg)
library(shiny)
library(InteractiveComplexHeatmap)

args <- commandArgs(trailingOnly = TRUE)

if (length(args) == 0) {
    cat("No arguments provided. Exiting script.\n -i : Input is rgi result generated by add_rgi_togenePA.py, -m metadata CSV file, -out output heatmap file, -meta_cols : metadata column numbers like 2,3,4\n")
  quit(save = "no", status = 1)
}

arg_list <- list()
for (i in seq(1, length(args), 2)) {
  arg_list[[args[i]]] <- args[i + 1]
}

input_csv <- arg_list[["-i"]]
metadata_file <- arg_list[["-m"]]
output_file <- arg_list[["-out"]]
meta_col_selected <- arg_list[["-mc"]]
meta_col_selected <- as.integer(meta_col_selected)
meta_cols <- NULL

output_html <- arg_list[["-out_html"]]

set_font <- function() {
  par(family = "Courier")
}

set_font()

if ("-meta_cols" %in% names(arg_list)) {
    meta_cols <- unlist(strsplit(arg_list[["-meta_cols"]], ","))
}

# prohibit too long legend text. 
wrap_after_two_semicolons <- function(s) {
  parts <- unlist(strsplit(s, split = ";", fixed = TRUE))
  new_text <- ""
  for (i in seq_along(parts)) {
    new_text <- paste(new_text, parts[i], sep = "")
    if (i < length(parts) && (i %% 2 == 0)) {
      new_text <- paste(new_text, ";", "\n", sep = "")
    } else if (i < length(parts)) {
      new_text <- paste(new_text, ";", sep = "")
    }
  }
  new_text
}

wrap_after_semicolon <- function(s) {
  parts <- unlist(strsplit(s, split = ";", fixed = TRUE))
  new_text <- paste(parts, collapse = ";\n")
  new_text
}

clean_column_names <- function(df) {
  names(df) <- gsub(" ", "_", names(df))
  return(df)
}

card_raw <- read.csv(input_csv, header = TRUE,check.names=FALSE)

card_raw <- clean_column_names(card_raw)

#input should be csv. 
metadata <- read.csv(metadata_file, header = TRUE, check.names = FALSE)

selected_colname = colnames(metadata)[meta_col_selected]

data <- read.csv(input_csv, header = TRUE, check.names = FALSE)
if (!is.null(meta_cols)) {
    meta_cols <- as.numeric(meta_cols)
    metadata <- metadata[, meta_cols, drop = FALSE]
}





#####################################
############# set this ##############
# set metadata rownames(the output should be grepped by genome_analysis_accession)
# genome_analysis_accession
# manually create based on scritps and get them . 
############# set this ##############
#####################################
rownames(metadata) <- metadata[,1] # should revise automated metadata  
# free below line 
# rownames(metadata) <- metadata[,colnaems(metadata[,])]


card_heatmap_data <- card_raw[, 2:(ncol(card_raw) - 3)]


selected_columns <- c("Completeness", "Contamination", "classification", colnames(metadata[selected_colname])) #meta1380 %>% select(ncol(.)))
meta_filtered_forvis <- metadata[, selected_columns]
meta_filtered_forvis <- meta_filtered_forvis %>%
  mutate(!!selected_colname := replace(.[[selected_colname]], .[[selected_colname]] == "", "NA"))

classification_split <- strsplit(as.character(meta_filtered_forvis$classification), ";")
classification_df <- do.call(rbind, classification_split)
colnames(classification_df) <- c("Domain", "Phylum", "Class", "Order", "Family", "Genus", "Species")
classification_df <- as.data.frame(classification_df)
meta_filtered_forvis <- cbind(meta_filtered_forvis[, -which(colnames(meta_filtered_forvis) == "classification")], classification_df)

meta_filtered_forvis_ordered <- meta_filtered_forvis[match(colnames(card_heatmap_data), rownames(meta_filtered_forvis)),]

top_annotation <- HeatmapAnnotation(df = meta_filtered_forvis_ordered,
                                    annotation_name_gp = gpar(fontsize = 7))


basic_metadata <- meta_filtered_forvis_ordered[, !colnames(meta_filtered_forvis_ordered) %in% c(selected_colname,"Domain", "Phylum", "Class", "Order", "Family")]
user_metadata <- meta_filtered_forvis_ordered[, selected_colname, drop = FALSE]

top_annotation_boxplot <- HeatmapAnnotation(Average_value = anno_boxplot(as.matrix(card_heatmap_data), axis = TRUE,height = unit(1.0, "cm"),
                               axis_param = list(labels_rot = 0)),
  annotation_name_side = "left", annotation_name_rot = 0,
  annotation_name_gp = gpar(fontsize = 11, fontfamily = "Courier",fontface="bold")#,
  #height = unit(1.0, "cm")
)
top_annotation_basic <- HeatmapAnnotation(
  df = basic_metadata,
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 11, fontfamily = "Courier",fontface="bold"),
  annotation_legend_param = list(
  title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold"),
  labels_gp = gpar(fontsize = 11, fontfamily = "Courier")
  )
)
top_annotation_user <- HeatmapAnnotation(
  df = user_metadata,
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 11, fontfamily = "Courier", col = "red", fontface="bold"),
    annotation_legend_param = list(
    title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier")
  )
)





pdf_width <- (2.5*ncol(card_heatmap_data) +90 + 110 +100 +40) / 25.4
pdf_height <- (2.5*nrow(card_heatmap_data) + 45 +40 + 40 ) / 25.4

if (pdf_width >= 20.9) {
  pdf_width <- 20.9
}

if (pdf_height >= 15) {
  pdf_height <- 15.65
}

# row_annotation <- rowAnnotation(
#     df = arg_info,
#     annotation_name_gp = gpar(fontsize = 8)
# )

arg_info <- card_raw[, (ncol(card_raw) - 2):ncol(card_raw)]


rowsum_data <- data.frame(rowsum = rowSums(card_heatmap_data))



# ############
# reorder based on largest category
integrated_data <- cbind(card_heatmap_data, arg_info) %>%
  mutate(rowsum = rowSums(select(., 1:(ncol(card_heatmap_data)))))

integrated_data_sorted <- integrated_data %>%
  group_by(Resistance_Mechanism) %>%
  arrange(Resistance_Mechanism, desc(rowsum)) %>%
  mutate(row_index = row_number()) %>%
  ungroup() %>%
  arrange(Resistance_Mechanism, row_index)

row_order <- integrated_data_sorted$row_index

  if(n_distinct(integrated_data_sorted$CARD_Short_Name) > 90) {
  top_90_card <- integrated_data_sorted %>%
    group_by(CARD_Short_Name) %>%
    summarise(total_rowsum = sum(rowsum)) %>%
    arrange(desc(total_rowsum)) %>%
    slice_head(n = 90) %>%
    pull(CARD_Short_Name)
  
  integrated_data_sorted <- integrated_data_sorted %>%
    mutate(CARD_Short_Name_display = ifelse(CARD_Short_Name %in% top_90_card, CARD_Short_Name, "Other"))
} else {
  integrated_data_sorted$CARD_Short_Name_display <- integrated_data_sorted$CARD_Short_Name
}


# left annotation ,  represent average  number of genes. 





#generate split matrix 
integrated_matrix <- as.matrix(integrated_data_sorted[, 1:ncol(card_heatmap_data)])



split_data <- split(as.data.frame(t(integrated_matrix)), meta_filtered_forvis_ordered[, selected_colname])
split_matrix <- do.call(cbind, lapply(split_data, function(x) colMeans(x, na.rm = TRUE)))

print(paste("Dimensions of card_heatmap_data:", paste(dim(card_heatmap_data), collapse = " x ")))
print(paste("Dimensions of split_matrix:", paste(dim(split_matrix), collapse = " x ")))

distribution_colors <- scales::hue_pal()(ncol(split_matrix))
distribution_legend <- Legend(
    labels = colnames(split_matrix),
    legend_gp = gpar(fill = distribution_colors),
    title = "Resistance_averages",
    title_gp = gpar(fontsize = 12, fontfamily = "Courier",fontface="bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier")    
)
# left_annotation <- rowAnnotation(
#   Resistance_averages = anno_points(split_matrix, 
#                              pch = 1:ncol(split_matrix),
#                              size = unit(1, "mm"),
#                              title_gp = gpar(fontsize = 12, fontfamily = "Courier",fontface="bold"),

#                              gp = gpar(col = scales::hue_pal()(ncol(split_matrix))),
#                              width = unit(3, "cm")
#   ),  annotation_name_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface="bold")
# )

left_annotation <- rowAnnotation(
  Resistance_averages = anno_points(
    split_matrix,
    pch = 1:ncol(split_matrix),
    size = unit(1, "mm"),
    title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold"),
    gp = gpar(col = scales::hue_pal()(ncol(split_matrix))),
    width = unit(3, "cm")
  ),
  annotation_name_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold")
)



# grouped_data <- arg_info %>%
#   mutate(rowsum = rowSums(card_heatmap_data)) %>%
#   group_by(Resistance_Mechanism) %>%
#   arrange(Resistance_Mechanism, desc(rowsum)) %>%
#   mutate(row_index = row_number()) %>%
#   ungroup() %>%
#   arrange(Resistance_Mechanism, row_index)



# ############


# ############
# prohibit too long category  . 
#integrated_data_sorted
# if(n_distinct(integrated_data_sorted$CARD_Short_Name) > 90) {
#   top_90_card <- integrated_data %>%
#     group_by(CARD_Short_Name) %>%
#     summarise(total_rowsum = sum(rowsum)) %>%
#     arrange(desc(total_rowsum)) %>%
#     slice_head(n = 90) %>%
#     pull(CARD_Short_Name)


# if(n_distinct(grouped_data$CARD_Short_Name) > 90) {
#   top_90_card <- grouped_data %>%
#     group_by(CARD_Short_Name) %>%
#     summarise(total_rowsum = sum(rowsum)) %>%
#     arrange(desc(total_rowsum)) %>%
#     slice_head(n = 90) %>%
#     pull(CARD_Short_Name)
  
#   grouped_data <- grouped_data %>%
#     mutate(CARD_Short_Name_display = ifelse(CARD_Short_Name %in% top_90_card, CARD_Short_Name, "Other"))
# } else {
#   grouped_data$CARD_Short_Name_display <- grouped_data$CARD_Short_Name
# }
# prohibit too long category  . 
# ############


#card_heatmap_data <- card_heatmap_data[grouped_data$row_index, ]
#arg_info <- arg_info[grouped_data$row_index, ]



# prohibit too long legned text  # Durg_class information will be only shown in the hover text. 
#integrated_data$Resistance_Mechanism <- sapply(integrated_data$Resistance_Mechanism, wrap_after_semicolon)
integrated_data_sorted$Resistance_Mechanism <- sapply(integrated_data_sorted$Resistance_Mechanism, wrap_after_semicolon)

#arg_info$"Drug_Class" <- sapply(arg_info$"Drug_Class", wrap_after_two_semicolons)

head(integrated_data)


row_annotation <- HeatmapAnnotation(
  Resistance_Mechanism = integrated_data_sorted$Resistance_Mechanism,
  CARD_Short_Name = integrated_data_sorted$CARD_Short_Name_display,
  which = 'row',
  annotation_name_gp = gpar(fontsize = 11, fontfamily = "Courier", fontface = "bold"),
  annotation_legend_param = list(
    title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier")
  )
)

row_annotation_hover <- HeatmapAnnotation(
  Drug_Class = integrated_data_sorted$Drug_Class,
  show_annotation_name = TRUE,
  show_legend = FALSE,
  which = 'row'
)


# 호버 텍스트를 위한 추가 annotation (Drug_Class 포함)



message("which")
head(row_annotation)
#annotation_width = unit(4, "cm")

length(unique(as.vector(card_heatmap_data)))
#this is for PA 

# breaks <- c(0, min(vfdb_matrix[vfdb_matrix>0]), max(vfdb_matrix))
# print(breaks)
# #doncha --> solved  last 5 field were attached 
# colors <- c("lightgrey", "white", "red")


process_card_heatmap_data <- function(card_heatmap_data) {
  if (all(card_heatmap_data %in% c(0, 1))) {
    # Presence/Absence case
    return(list(
      color_function = colorRamp2(c(0, 1), c("lightgrey", "red")),
      legend_labels = c("Absent", "Present"),
      at = c(0, 1)
    ))
  } else {
    # Gene count case
    breaks <- c(0, min(card_heatmap_data[card_heatmap_data > 0]), max(card_heatmap_data))
    colors <- c("lightgrey", "white", "red")
    return(list(
      color_function = colorRamp2(breaks, colors),
      legend_labels = as.character(breaks),
      at = breaks
    ))
  }
}

print(card_heatmap_data[1:3,1:3])
print(card_heatmap_data[(nrow(card_heatmap_data)-3):nrow(card_heatmap_data),(ncol(card_heatmap_data)-3):ncol(card_heatmap_data)])

result <- process_card_heatmap_data(card_heatmap_data)
my_color_function <- result$color_function
legend_labels <- result$legend_labels
legend_at <- result$at


ht1 <- Heatmap(
    as.matrix(integrated_data_sorted[, 1:ncol(card_heatmap_data)]),
    #as.matrix(card_heatmap_data),
    name = "Antimicrobial resistance\n gene number",
    col = my_color_function,

    top_annotation = c(top_annotation_boxplot, top_annotation_basic, top_annotation_user),
    right_annotation = c(row_annotation, row_annotation_hover),
    #row_order = order(as.integer(factor(grouped_data$Resistance_Mechanism)), grouped_data$row_index),
    row_order = order(as.integer(factor(integrated_data_sorted$Resistance_Mechanism)), integrated_data_sorted$row_index),
    left_annotation=left_annotation,
    heatmap_legend_param = list(
    legend_direction = "horizontal",
    title_gp = gpar(fontsize = 12, fontfamily = "Courier",fontface="bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier"),
    labels = legend_labels,
    at = legend_at
    ), 
    show_row_names = FALSE,
    show_column_names = FALSE,
    show_column_dend = FALSE,
    show_row_dend = FALSE,
    column_split  = meta_filtered_forvis_ordered[,selected_colname],
    row_names_gp = gpar(fontsize = 11, fontfamily = "Courier"),
    column_names_gp = gpar(fontsize = 11, fontfamily = "Courier"),
    column_title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface = "bold", col = "blue")
    
)

ht1 <- draw(ht1,
padding = unit(c(2, 2, 2, 2), "cm"),
annotation_legend_list = list(distribution_legend),annotation_legend_side = "left")

#draw(lgd, x = unit(1, "cm"), y = unit(1, "cm"), just = c("left", "bottom"))

output_file <- "heatmap_CARD.pdf"
pdf(file = output_file, width = pdf_width, height = pdf_height)
#pdf(file = output_file, width = 20, height = 15)
draw(ht1)
dev.off()







htShiny(ht1, output_ui_float = TRUE, save = output_html)

#write.csv(card_heatmap_data, file = output_table)












