# =========================
# 06 生成审稿人相关表格
# =========================

source("R/config.R")
source("R/utils.R")

library(dplyr)
library(readxl)
library(openxlsx)

# ---------- 1. Figure panel key ----------
panel_df <- expand.grid(
  Gene = GENE_ORDER,
  Outcome = OUTCOME_ORDER,
  stringsAsFactors = FALSE
)
panel_df <- panel_df[order(match(panel_df$Gene, GENE_ORDER), match(panel_df$Outcome, OUTCOME_ORDER)), ]
panel_df$Panel <- paste0("(", letters[seq_len(nrow(panel_df))], ")")
panel_df <- panel_df[, c("Panel", "Gene", "Outcome")]

openxlsx::write.xlsx(
  panel_df,
  file.path(DIR_TABLES, "Figure_panel_key.xlsx"),
  overwrite = TRUE
)

# ---------- 2. Supplementary Table S13 ----------
loci <- read_loci_file()

gwas_summary_file <- file.path(DIR_TABLES, "original_loci_GWAS_region_extraction_summary.xlsx")
pqtl_summary_file <- file.path(DIR_TABLES, "original_loci_pQTL_region_extraction_summary.xlsx")
harm_summary_file <- file.path(DIR_TABLES, "original_loci_harmonise_summary_final.xlsx")

gwas_summary <- if (file.exists(gwas_summary_file)) openxlsx::read.xlsx(gwas_summary_file) else data.frame()
pqtl_summary <- if (file.exists(pqtl_summary_file)) openxlsx::read.xlsx(pqtl_summary_file) else data.frame()
harm_summary <- if (file.exists(harm_summary_file)) openxlsx::read.xlsx(harm_summary_file) else data.frame()

regions <- as.data.frame(loci) %>%
  mutate(
    Gene = as.character(Gene),
    Protein_ID = as.character(Protein_ID),
    Outcome = as.character(Outcome)
  ) %>%
  select(any_of(c("Gene", "Protein_ID", "Protein", "Outcome", "Outcome_ID", "chr", "region_start", "region_end")))

if (nrow(gwas_summary) > 0) {
  regions <- regions %>%
    left_join(
      gwas_summary %>% select(any_of(c("Gene", "Protein_ID", "Outcome", "GWAS_SNPs_in_region"))),
      by = c("Gene", "Protein_ID", "Outcome")
    )
}

if (nrow(pqtl_summary) > 0) {
  regions <- regions %>%
    left_join(
      pqtl_summary %>% select(any_of(c("Gene", "Protein_ID", "SNPs_queried", "pQTL_SNPs_returned"))),
      by = c("Gene", "Protein_ID")
    )
}

if (nrow(harm_summary) > 0) {
  regions <- regions %>%
    left_join(
      harm_summary %>% select(any_of(c("Gene", "Protein_ID", "Outcome", "Overlap_SNPs"))),
      by = c("Gene", "Protein_ID", "Outcome")
    )
}

regions <- regions %>%
  mutate(
    Gene = factor(Gene, levels = GENE_ORDER),
    Outcome = factor(Outcome, levels = OUTCOME_ORDER)
  ) %>%
  arrange(Gene, Outcome) %>%
  mutate(Panel = paste0("(", letters[seq_len(n())], ")")) %>%
  select(Panel, everything()) %>%
  mutate(Gene = as.character(Gene), Outcome = as.character(Outcome))

wb <- createWorkbook()
addWorksheet(wb, "Supplementary Table S13")
writeData(wb, "Supplementary Table S13",
          "Supplementary Table S13. Genomic regions and SNP numbers used for regional association and locuscompare plots.",
          startRow = 1, startCol = 1)
writeData(wb, "Supplementary Table S13",
          "This table summarizes the genomic windows used to generate regional association and locuscompare plots for the six colocalized protein loci. GWAS_SNPs_in_region indicates the number of FinnGen GWAS variants extracted within each genomic window. pQTL_SNPs_returned indicates the number of protein-associated variants available within the same region. Overlap_SNPs indicates the number of variants shared between the pQTL and GWAS datasets and used for visualization.",
          startRow = 2, startCol = 1)
writeData(wb, "Supplementary Table S13", regions, startRow = 4, startCol = 1)

header_style <- createStyle(textDecoration = "bold", fgFill = "#D9EAF7", border = "Bottom")
addStyle(wb, "Supplementary Table S13", header_style, rows = 4, cols = seq_len(ncol(regions)), gridExpand = TRUE)
setColWidths(wb, "Supplementary Table S13", cols = seq_len(ncol(regions)), widths = "auto")

saveWorkbook(
  wb,
  file.path(DIR_TABLES, "Supplementary_Table_S13_regions_SNP_counts.xlsx"),
  overwrite = TRUE
)

# ---------- 3. Table 1 colocalization results ----------
# 如果 04_coloc_results 里存在 DVT/PTE/VTE 的 coloc 原始结果 Excel，则尝试自动整理。
# 如果列名不标准，脚本会生成一个空模板，方便手动粘贴。

pick_first_col <- function(df, patterns) {
  nms <- names(df)
  for (p in patterns) {
    hit <- grep(p, nms, ignore.case = TRUE, value = TRUE)
    if (length(hit) > 0) return(hit[1])
  }
  NA_character_
}

read_coloc_result <- function(file, outcome) {
  df <- tryCatch(readxl::read_excel(file), error = function(e) NULL)
  if (is.null(df)) return(data.frame())
  df <- as.data.frame(df)
  names(df) <- gsub("\\s+", "_", names(df))
  names(df) <- gsub("\\.", "_", names(df))

  gene_col <- pick_first_col(df, c("^Gene$", "gene", "symbol"))
  prot_col <- pick_first_col(df, c("Protein_ID", "protein", "^id$"))
  nsnps_col <- pick_first_col(df, c("nsnps", "n_snps", "snp"))
  h0_col <- pick_first_col(df, c("PP_H0", "PPH0", "H0"))
  h1_col <- pick_first_col(df, c("PP_H1", "PPH1", "H1"))
  h2_col <- pick_first_col(df, c("PP_H2", "PPH2", "H2"))
  h3_col <- pick_first_col(df, c("PP_H3", "PPH3", "H3"))
  h4_col <- pick_first_col(df, c("PP_H4", "PPH4", "H4"))

  needed <- c(gene_col, h0_col, h1_col, h2_col, h3_col, h4_col)
  if (any(is.na(needed))) return(data.frame())

  out <- data.frame(
    Gene = as.character(df[[gene_col]]),
    Protein_ID = if (!is.na(prot_col)) as.character(df[[prot_col]]) else NA_character_,
    Outcome = outcome,
    nsnps = if (!is.na(nsnps_col)) suppressWarnings(as.numeric(df[[nsnps_col]])) else NA_real_,
    PP.H0 = suppressWarnings(as.numeric(df[[h0_col]])),
    PP.H1 = suppressWarnings(as.numeric(df[[h1_col]])),
    PP.H2 = suppressWarnings(as.numeric(df[[h2_col]])),
    PP.H3 = suppressWarnings(as.numeric(df[[h3_col]])),
    PP.H4 = suppressWarnings(as.numeric(df[[h4_col]])),
    Source_file = basename(file),
    stringsAsFactors = FALSE
  )
  out <- out[!is.na(out$Gene) & out$Gene != "", ]
  out
}

coloc_files <- list.files("04_coloc_results", pattern = "\\.xlsx$", full.names = TRUE)

coloc_list <- list()
for (f in coloc_files) {
  bn <- basename(f)
  outcome <- if (grepl("DVT", bn, ignore.case = TRUE)) {
    "DVT"
  } else if (grepl("PTE|PULMEMB|PE", bn, ignore.case = TRUE)) {
    "PE"
  } else if (grepl("VTE", bn, ignore.case = TRUE)) {
    "VTE"
  } else {
    NA_character_
  }
  if (!is.na(outcome)) {
    coloc_list[[length(coloc_list) + 1]] <- read_coloc_result(f, outcome)
  }
}

coloc_df <- dplyr::bind_rows(coloc_list)

if (nrow(coloc_df) == 0) {
  coloc_df <- data.frame(
    Gene = character(),
    Protein_ID = character(),
    Outcome = character(),
    nsnps = numeric(),
    PP.H0 = numeric(),
    PP.H1 = numeric(),
    PP.H2 = numeric(),
    PP.H3 = numeric(),
    PP.H4 = numeric(),
    Significant_colocalization = character(),
    stringsAsFactors = FALSE
  )
} else {
  coloc_df <- coloc_df %>%
    filter(Gene %in% GENE_ORDER) %>%
    mutate(
      Gene = factor(Gene, levels = GENE_ORDER),
      Outcome = factor(Outcome, levels = OUTCOME_ORDER),
      Significant_colocalization = ifelse(PP.H4 >= 0.80, "Yes", "No")
    ) %>%
    arrange(Gene, Outcome) %>%
    mutate(Gene = as.character(Gene), Outcome = as.character(Outcome))
}

wb2 <- createWorkbook()
addWorksheet(wb2, "Table 1")
writeData(wb2, "Table 1",
          "Table 1. Colocalization results of six significant proteins associated with VTE-related outcomes.",
          startRow = 1, startCol = 1)
writeData(wb2, "Table 1",
          "PP.H0–PP.H4 represent posterior probabilities for the five coloc hypotheses. PP.H4 ≥ 0.80 was considered evidence supporting colocalization between protein levels and VTE-related outcomes. VTE, venous thromboembolism; DVT, deep vein thrombosis; PE, pulmonary embolism.",
          startRow = 2, startCol = 1)
writeData(wb2, "Table 1", coloc_df, startRow = 4, startCol = 1)
if (ncol(coloc_df) > 0) {
  addStyle(wb2, "Table 1", header_style, rows = 4, cols = seq_len(ncol(coloc_df)), gridExpand = TRUE)
  setColWidths(wb2, "Table 1", cols = seq_len(ncol(coloc_df)), widths = "auto")
}
saveWorkbook(
  wb2,
  file.path(DIR_TABLES, "Table_1_colocalization_results.xlsx"),
  overwrite = TRUE
)

write_status("06_make_tables", TRUE, "Reviewer-facing tables generated.")