#!/usr/bin/env Rscript
#
# grouped plot for multiple samples with SNPs in one row, INDELs in another row, and a single common legend


library(ggplot2)
library(tools)
library(gridExtra)  # For grid.arrange
library(grid)

args = commandArgs(trailingOnly=TRUE)

if(length(args) < 2) {
    stop("Usage: rocplot_grouped.R output_name HG001:file1:name1,file2:name2... HG003:file1:name1,file2:name2...")
}

# Define pipeline colors - explicitly match the pipeline names as they appear in input
pipeline_colors <- c(
  "BWA-MEM" = "#E53935",  # Red for BWA-MEM
  "Giraffe(HPRC_pangenome)" = "#1E88E5",  # Blue for HPRC Pangenome
  "Giraffe(1kGP_pangenome)" = "#64B5F6",  # Lighter blue for 1kGP Pangenome
  "Giraffe(diploid)" = "#FB8C00",  # Bright Orange 
  "Giraffe(diploid_reported)" = "#FFB74D",  # Lighter Orange
  "Giraffe(Imputefirst_c20)" = "#673AB7",  # Deep purple for Giraffe c20
  "Giraffe(Imputefirst_c5)" = "#4527A0",  # Dark purple for Giraffe c5
  "LevioSAM2(Imputefirst_c5)" = "#AA00FF",  # Bright purple/pink for Leviosam2 c5
  "LevioSAM2(Imputefirst_c20)" = "#E040FB",  # Pink/magenta for Leviosam2 c20
  "Giraffe(benchmark)" = "#2E7D32",  # Dark Green
  "LevioSAM2(benchmark)" = "#43A047"  # Light Green
)

# Read single hap.py / xcmp dataset
read_single = function(x) {
    nx = strsplit(x, "\\:")[[1]]

    if(length(nx) == 1) {
        name = basename(file_path_sans_ext(x))
    } else {
        x = nx[1]
        name = nx[2]
    }
    cat(sprintf("Reading %s as %s\n", x, name))

    all_results = list()

    # Read SNP and INDEL data files
    all_results$roc_data_snp = read.csv(paste(x, "roc.Locations.SNP", "csv", "gz", sep="."))
    all_results$roc_data_indel = read.csv(paste(x, "roc.Locations.INDEL", "csv", "gz", sep="."))

    # Mark ALL points separately
    all_results$roc_data_snp_all = head(subset(all_results$roc_data_snp, QQ == min(all_results$roc_data_snp["QQ"])), n=1)
    all_results$roc_data_snp_all$Filter = "ALL"
    
    all_results$roc_data_indel_all = head(subset(all_results$roc_data_indel, QQ == min(all_results$roc_data_indel["QQ"])), n=1)
    all_results$roc_data_indel_all$Filter = "ALL"
    
    # Add ROC label to the rest of the data
    all_results$roc_data_snp$Filter = "ROC"
    all_results$roc_data_indel$Filter = "ROC"

    # Combine all results
    result = do.call(rbind, all_results)
    row.names(result) = NULL
    result$filename = x
    result$name = name
    result$igroup = paste(result$name, result$Type)
    return(result)
}

# Plot individual type (SNP or INDEL) for a specific sample
plot_type = function(data, type, xlimits, sample_id) {
    type_data = subset(data, Type == type)
    
    # Use ROC-style coordinates
    xaxis = "FPR"
    yaxis = "TPR"
    type_data$FPR = type_data$QUERY.FP / (type_data$QUERY.TOTAL - type_data$QUERY.UNK)
    type_data$TPR = type_data$TRUTH.TP / (type_data$TRUTH.TP + type_data$TRUTH.FN)
    cc = complete.cases(type_data[, c(xaxis, yaxis)])
    type_data = type_data[cc, ]
    
    yl_min = max(0.98, min(type_data[[yaxis]]) - 0.005)
    yl_max = 1.0
    
    # Print unique pipeline names for debugging
    cat(sprintf("Pipeline names in %s dataset for %s:\n", type, sample_id))
    print(unique(type_data$name))
    
    plt = ggplot(type_data, aes_string(x=xaxis, y=yaxis, color="name"))
    
    # ROC lines
    plt = plt +
        geom_line(data = subset(type_data, Filter == "ROC"),
                  mapping=aes(group=igroup),
                  size=1.5,
                  linetype=1)
    
    # ALL points
    plt = plt +
        geom_point(data = subset(type_data, Filter == "ALL"),
                   mapping=aes(group=igroup),
                   size=8)
    
    # Set styling
    plt = plt +
        scale_color_manual(values = pipeline_colors, 
                          breaks = unique(type_data$name)) +
        xlim(xlimits) +
        ylim(c(yl_min, yl_max)) +
        ggtitle(paste(sample_id, "-", type)) +
        theme_bw(base_size=16) +
        theme(plot.title = element_text(hjust = 0.5))
    
    return(plt)
}

# Extract legend from a plot
get_legend = function(p) {
    tmp = ggplot_gtable(ggplot_build(p))
    leg = which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
    legend = tmp$grobs[[leg]]
    return(legend)
}

# Parse the sample data arguments
output_name = args[1]
sample_data = list()

for (i in 2:length(args)) {
    parts = strsplit(args[i], ":")[[1]]
    sample_id = parts[1]
    
    # Remaining parts are the file:name pairs
    file_pairs = paste(parts[2:length(parts)], collapse=":")
    file_list = strsplit(file_pairs, ",")[[1]]
    
    sample_data[[sample_id]] = file_list
}

# Create a plot for each sample and type
all_plots = list()
all_data = list()

for (sample_id in names(sample_data)) {
    # Read all data files for this sample
    sample_files = sample_data[[sample_id]]
    data = do.call(rbind, lapply(sample_files, read_single))
    all_data[[sample_id]] = data
    
    # Create SNP and INDEL plots
    all_plots[[paste0(sample_id, "_SNP")]] = plot_type(
        data, 
        "SNP", 
        c(0, 0.0025), 
        sample_id
    ) + 
    scale_x_continuous(
        limits = c(0, 0.0025), 
        breaks = seq(0, 0.0025, by = 0.0005)
    ) + 
    theme(legend.position = "none")
    
    all_plots[[paste0(sample_id, "_INDEL")]] = plot_type(
        data, 
        "INDEL", 
        c(0.0010, 0.0040), 
        sample_id
    ) + 
    scale_x_continuous(
        limits = c(0.0010, 0.0040), 
        breaks = seq(0.0010, 0.0040, by = 0.0005)
    ) + 
    theme(legend.position = "none")
}


# Extract a common legend from the first sample's SNP plot
first_sample = names(sample_data)[1]
common_legend = get_legend(
    plot_type(all_data[[first_sample]], "SNP", c(0, 0.0025), first_sample) + 
    theme(
        legend.position = "bottom",
        legend.direction = "horizontal",
        legend.box = "horizontal",
        legend.title = element_blank(),
        legend.text = element_text(size = 10),
        legend.key.size = unit(1, "lines"),
        legend.spacing.x = unit(0.5, "cm"),
        legend.margin = margin(t = 10, r = 0, b = 0, l = 0, unit = "pt")
    )
)


# Arrange SNP plots in first row
snp_plots = list()
for (sample_id in names(sample_data)) {
    snp_plots = c(snp_plots, list(all_plots[[paste0(sample_id, "_SNP")]]))
}

# Arrange INDEL plots in second row
indel_plots = list()
for (sample_id in names(sample_data)) {
    indel_plots = c(indel_plots, list(all_plots[[paste0(sample_id, "_INDEL")]]))
}

# Combine the plots with the legend
combined_plots = grid.arrange(
    arrangeGrob(grobs = snp_plots, ncol = length(snp_plots)),
    arrangeGrob(grobs = indel_plots, ncol = length(indel_plots)),
    common_legend,
    heights = c(4, 4, 1)
)

# Add your note and draw the plot
grid.draw(combined_plots)
grid.text("Note: 1kGP_pangenome is analyzed in HG001 and HG002 samples only.", 
          x = 0.5, y = 0.02, just = "center", 
          gp = gpar(fontsize = 10, fontface = "italic"))


# Save the combined plot
if (dev.cur() > 1) dev.off()
ggsave(
    paste(output_name,"pdf", sep="."), 
    #paste(output_name, "grouped", "pdf", sep="."), 
    combined_plots, 
    width = min(length(snp_plots) * 5, 20),  # Limit width to 20 inches
    height = 10, 
    dpi = 300,
    device = "pdf"
)

# Print message to confirm completion
#cat(sprintf("Grouped plot saved as %s\n", paste(output_name, "grouped", "pdf", sep=".")))
