library(ComplexHeatmap)
library(circlize)
library(ggplot2)
library(dplyr)
library(tidyr)
library(stringr)
library(RColorBrewer)

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


args <- commandArgs(trailingOnly = TRUE)

if (length(args) == 0) {
    cat("No arguments provided. Exiting script.\n -i : Input is CARD csv 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"]]

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


vfdb_raw <- read.csv(input_csv, header = TRUE, check.names = FALSE)
metadata <- read.csv(metadata_file, header = TRUE, check.names = FALSE)
selected_colname = colnames(metadata)[meta_col_selected]
print(selected_colname)
if (ncol(vfdb_raw) - 7 != nrow(metadata)) {
    stop("Mismatch in the number of columns in heatmap data and the number of rows in metadata.")
}



# for visualization partitioning
# group_col_name <- colnames(metadata_file)[as.integer(meta_col_selected)]
heatmap_data <- vfdb_raw[, 2:(ncol(vfdb_raw) - 6)]


vf_info <- vfdb_raw[, (ncol(vfdb_raw) - 5):ncol(vfdb_raw)] # last 6 fields

heatmap_data_summarized <- heatmap_data %>% 
  mutate(VFID = vf_info$VF_ID) %>%
  group_by(VFID) %>%
  summarise_all(sum)

character_data <- vf_info %>%
  group_by(VF_ID) %>%
  summarise_all(~first(.))


select_numeric_columns <- function(df) {
  df %>% select(where(is.numeric))
}




heatmap_data_summarized <- bind_cols(character_data, heatmap_data_summarized)
heatmap_data_summarized <- heatmap_data_summarized %>%
  mutate(
    VFC = str_extract(VFC_ID, "VFC\\d+"),
    VF = str_extract(VF_ID, "VF\\d+")
  ) %>%
  { print(select(., VFC_ID, VFC, VF_ID, VF)); . } %>%
  mutate(
    rowsum = rowSums(select_numeric_columns(select(., 8:ncol(.))), na.rm = TRUE)
  
  ) %>%
  { print(select(., VFC_ID, VF_ID, rowsum)); . } %>%
  group_by(VFC) %>%
  arrange(VFC, desc(rowsum)) %>%
  mutate(
    row_index = row_number()
  ) %>%
  { print(select(., VFC, row_index)); . } %>%
  ungroup() %>%
  arrange(VFC, row_index)

print(head(heatmap_data_summarized))


# VF 수가 90개 이상일 때 legend 줄이기
if(n_distinct(as.character(heatmap_data_summarized$VF)) > 90) {
  top_90_vf <- heatmap_data_summarized %>%
    group_by(VF) %>%
    summarise(total_rowsum = sum(rowsum)) %>%
    top_n(90, total_rowsum) %>%
    pull(VF)
  heatmap_data_summarized <- heatmap_data_summarized %>%
    mutate(VF_display = ifelse(VF %in% top_90_vf, VF, "Other"))
} else {
  heatmap_data_summarized$VF_display <- heatmap_data_summarized$VF
}
print(head(heatmap_data_summarized))



heatmap_data_summarized$VFC <- factor(heatmap_data_summarized$VFC, levels = unique(heatmap_data_summarized$VFC))
heatmap_data_summarized$VF_display <- factor(heatmap_data_summarized$VF_display, levels = unique(heatmap_data_summarized$VF_display))

heatmap_data_summarized$VF_ID <- factor(heatmap_data_summarized$VF_ID)
heatmap_data_summarized$VFC_ID <- factor(heatmap_data_summarized$VFC_ID)

set.seed(33)
vf_names <- unique(heatmap_data_summarized$VF_ID)
vfc_categories <- unique(heatmap_data_summarized$VFC_ID)
vf_colors <- rainbow(length(vf_names))
vfc_colors <- rainbow(length(vfc_categories))
vf_color_map <- setNames(vf_colors, vf_names)
vfc_color_map <- setNames(vfc_colors, vfc_categories)
head(vf_color_map)
head(vfc_color_map)
# heatmap_data_summarized$VF_ID <- factor(heatmap_data_summarized$VF_ID, levels = vf_names)
# heatmap_data_summarized$VFC_ID <- factor(heatmap_data_summarized$VFC_ID, levels = vfc_categories)
heatmap_data_summarized$VF_ID <- factor(heatmap_data_summarized$VF_ID, levels = names(vf_color_map))
heatmap_data_summarized$VFC_ID <- factor(heatmap_data_summarized$VFC_ID, levels = names(vfc_color_map))

all(names(vf_color_map) %in% levels(heatmap_data_summarized$VF_ID)) # Should return TRUE
all(names(vfc_color_map) %in% levels(heatmap_data_summarized$VFC_ID)) # Should return TRUE

head(heatmap_data_summarized)
# metadata add 
rownames(metadata) <- metadata[,1]

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

print(meta_filtered_forvis)

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)
head(meta_filtered_forvis)

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



# heatmap visualiztion add annotation
top_annotation <- HeatmapAnnotation(df = meta_filtered_forvis_ordered,
                                    annotation_name_gp = gpar(fontsize = 7)) 
# row_annotation <- HeatmapAnnotation(
#   VF_Name = anno_text(heatmap_data_summarized$VF_ID),
#   VFcategory = anno_text(heatmap_data_summarized$VFC_ID),
#   which='row',annotation_name_gp = gpar(fontsize = 7))


print(levels(heatmap_data_summarized$VF_ID))
print(levels(heatmap_data_summarized$VFC_ID))
head(heatmap_data_summarized)

row_annotation <-HeatmapAnnotation(
   VF_Name = heatmap_data_summarized$VF_display,
   VFcategory = heatmap_data_summarized$VFC,  
   #VF_ID = heatmap_data_summarized$VF_ID,
  #VFC_ID = heatmap_data_summarized$VFC_ID,
   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(
   VF_ID = heatmap_data_summarized$VF_ID,
  VFC_ID = heatmap_data_summarized$VFC_ID,
   show_annotation_name = FALSE, 
   show_legend = FALSE,
   which = 'row',
   height = unit(0, "mm")  
)






print(row_annotation)

#print(summary(heatmap_data_summarized[,8:ncol(heatmap_data_summarized)]))

#vfdb_matrix<- heatmap_data_summarized[,8:(ncol(heatmap_data_summarized) -5)]

vfdb_matrix <- as.matrix(heatmap_data_summarized[,8:(ncol(heatmap_data_summarized)-5)])
print("Matrix dimensions:")
print(dim(vfdb_matrix))


#asdasdasd
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")
vfdb_matrix <- as.matrix(heatmap_data_summarized[,8:(ncol(heatmap_data_summarized)-5)])
                                           

print(vfdb_matrix[(nrow(vfdb_matrix) - 5): nrow(vfdb_matrix), (ncol(vfdb_matrix) - 5): ncol(vfdb_matrix)])

print(vfdb_matrix[1:3,1:3])


print("Whatthe?2")
if (length(unique(as.vector(vfdb_matrix))) == 1) {

  my_color_function <- "red"
} else {
  my_color_function <- colorRamp2(breaks,colors)
}
print("Whatthe?3")

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]
print("Whatthe?4")

print(vfdb_matrix[(nrow(vfdb_matrix) - 5):nrow(vfdb_matrix), (ncol(vfdb_matrix) - 5):ncol(vfdb_matrix)])

#print(vfdb_matrix[1:3,1:12])

top_annotation_boxplot <- HeatmapAnnotation(Average_value = anno_boxplot(as.matrix(vfdb_matrix), 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")
)



print("Whatthe?5")
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")
  )
)
print("Whatthe?6")
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")
  )
)

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

distribution_colors <- scales::hue_pal()(ncol(split_matrix))
distribution_legend <- Legend(
    labels = colnames(split_matrix),
    legend_gp = gpar(fill = distribution_colors),
    title = "Virulence_factor_averages",
    title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface="bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier")    
)

left_annotation <- rowAnnotation(
  Virulence_factor_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")
)


heatmap <- Heatmap(
  as.matrix(heatmap_data_summarized[,8:(ncol(heatmap_data_summarized)-5)]),
  name = "Virulence Factors",
  col = my_color_function,
  top_annotation = c(top_annotation_boxplot, top_annotation_basic, top_annotation_user),
  right_annotation = c(row_annotation,row_annotation_hover),
  #right_annotation = row_annotation,
  left_annotation = left_annotation,
  show_row_names = FALSE,
  show_column_names = FALSE,
  show_column_dend = FALSE,
  row_order = order(as.integer(factor(heatmap_data_summarized$VFC)), heatmap_data_summarized$row_index),  # VFC,VF arrange
  column_split  = meta_filtered_forvis_ordered[,selected_colname],
  heatmap_legend_param = list(
    legend_direction = "horizontal",
    title_gp = gpar(fontsize = 12, fontfamily = "Courier", fontface="bold"),
    labels_gp = gpar(fontsize = 11, fontfamily = "Courier")
  ),
  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")  
)
# Save heatmap to file
output_file <- "heatmap_VFDB.pdf"

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

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

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




library(ragg)
library(ComplexHeatmap)
library(circlize)
library(InteractiveComplexHeatmap)

heatmap <- draw(heatmap,
    padding = unit(c(2, 2, 2, 2), "cm"),
    annotation_legend_list = list(distribution_legend),
    annotation_legend_side = "left"
)
pdf(file = output_file, width = pdf_width, height = pdf_height)


#pdf(file = output_file, width = 20, height = 15)
draw(heatmap)
dev.off()


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