#!/usr/bin/env Rscript
#
# Modified version of rocplot.R
# Changes:
# - Kept ROC lines but made them solid instead of dotted
# - Combined SNP and INDEL into a single plot
# - Removed PASS and SEL data processing
# - Removed precision-recall plot functionality
# - Added custom color palette for pipelines
# - Moved legend to bottom of plot
#

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

args = commandArgs(trailingOnly=TRUE)

if(length(args) < 2) {
    stop("Usage: rocplot.Rscript output_name prefix_1:name_1 ... ")
}

# 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)
plot_type = function(data, type, display_label, xlimits) {
    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:\n", display_label))
    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(display_label) +
        theme_bw(base_size=18) +
        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)
}

# Read all data files
data = do.call(rbind, lapply(args[2:length(args)], read_single))

# For proper plot visibility, x-limits vary for SNP and INDEL

# For SNP: ticks every 0.0005 units
# Filter by "SNP" (actual value in data) but display as "SNVs"
snp_plot = plot_type(data, "SNP", "SNVs", c(0, 0.0025)) + 
  scale_x_continuous(limits = c(0, 0.0025), 
                     breaks = seq(0, 0.0025, by = 0.0005))
# For INDEL: ticks every 0.001 units
# Filter by "INDEL" (actual value in data) but display as "indel"
indel_plot = plot_type(data, "INDEL", "indel", c(0.0010, 0.0040)) + 
  scale_x_continuous(limits = c(0.0010, 0.0040), 
                     breaks = seq(0.0010, 0.0040, by = 0.0005))

# Get legend from one of the plots with proper styling
legend = get_legend(snp_plot + 
                   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")))

# Remove legends from individual plots
snp_plot = snp_plot + theme(legend.position = "none")
indel_plot = indel_plot + theme(legend.position = "none")

# Arrange plots with gridExtra
plots = grid.arrange(
    arrangeGrob(snp_plot, indel_plot, ncol = 2),
    legend,
    heights = c(5, 1)
)

if (dev.cur() > 1) dev.off()

# Save the combined plot
ggsave(paste(args[1], "pdf", sep="."), 
       plots, width = 10, height = 5, dpi = 300)

# Print message to confirm completion
# cat(sprintf("Plot saved as %s\n", paste(args[1], "combined", "pdf", sep=".")))
