# =========================
# 05 画 18 张 regional/locuscompare-like 图
# =========================

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

library(ggplot2)
library(patchwork)

make_one_locuscompare_plot <- function(dat, gene, outcome) {
  dat <- data.table::as.data.table(dat)

  # 必要列检查
  required <- c("position_pqtl", "p_pqtl", "pos_gwas", "p_gwas")
  missing_required <- setdiff(required, names(dat))
  if (length(missing_required) > 0) {
    stop("Missing columns for plotting: ", paste(missing_required, collapse = ", "))
  }

  dat[, x_pqtl := neglog10(p_pqtl)]
  dat[, y_gwas := neglog10(p_gwas)]
  dat[, ppos_mb := position_pqtl / 1e6]
  dat[, gpos_mb := pos_gwas / 1e6]

  # lead SNP：优先选择 pQTL 最小 P 值
  lead <- dat[which.min(p_pqtl)]
  lead_snp <- lead$SNP[1]
  lead_pos <- lead$position_pqtl[1] / 1e6

  p_main <- ggplot(dat, aes(x = x_pqtl, y = y_gwas)) +
    geom_point(size = 0.9, alpha = 0.75) +
    geom_point(
      data = dat[SNP == lead_snp],
      aes(x = x_pqtl, y = y_gwas),
      size = 1.8
    ) +
    labs(
      x = expression(-log[10](pQTL~italic(P))),
      y = expression(-log[10](GWAS~italic(P))),
      title = paste0(gene, " - ", outcome)
    ) +
    theme_bw(base_size = 9) +
    theme(
      plot.title = element_text(face = "bold", size = 11),
      panel.grid.minor = element_blank()
    )

  p_pqtl <- ggplot(dat, aes(x = ppos_mb, y = x_pqtl)) +
    geom_point(size = 0.8, alpha = 0.75) +
    geom_vline(xintercept = lead_pos, linetype = "dashed") +
    annotate("text", x = lead_pos, y = max(dat$x_pqtl, na.rm = TRUE), label = lead_snp,
             vjust = -0.2, size = 2.7) +
    labs(
      x = "Genomic position (Mb)",
      y = expression(-log[10](pQTL~italic(P))),
      title = "pQTL association signal"
    ) +
    theme_bw(base_size = 8) +
    theme(
      plot.title = element_text(size = 9),
      panel.grid.minor = element_blank()
    )

  p_gwas <- ggplot(dat, aes(x = gpos_mb, y = y_gwas)) +
    geom_point(size = 0.8, alpha = 0.75) +
    geom_vline(xintercept = lead_pos, linetype = "dashed") +
    labs(
      x = "Genomic position (Mb)",
      y = expression(-log[10](GWAS~italic(P))),
      title = "FinnGen GWAS association signal"
    ) +
    theme_bw(base_size = 8) +
    theme(
      plot.title = element_text(size = 9),
      panel.grid.minor = element_blank()
    )

  # 近似 locuscompare 版式：左主图，右侧上下两个 regional plot
  combined <- p_main + (p_pqtl / p_gwas) + plot_layout(widths = c(1, 1.05))
  list(plot = combined, lead_snp = lead_snp, snps_plotted = nrow(dat),
       pQTL_max_neglog10P = max(dat$x_pqtl, na.rm = TRUE),
       GWAS_max_neglog10P = max(dat$y_gwas, na.rm = TRUE))
}

plot_summary <- list()
panel_letters <- letters[1:18]
panel_i <- 1

for (gene in GENE_ORDER) {
  for (outcome in OUTCOME_ORDER) {
    f <- file.path(DIR_HARMONISED, paste0(gene, "_", outcome, "_harmonised.tsv.gz"))
    if (!file.exists(f)) {
      warning("Missing harmonised file: ", f)
      next
    }

    dat <- data.table::fread(f)
    dat <- dat[!is.na(p_pqtl) & !is.na(p_gwas) & p_pqtl > 0 & p_gwas > 0]

    if (nrow(dat) == 0) {
      warning("No valid data for plotting: ", gene, " ", outcome)
      next
    }

    res <- make_one_locuscompare_plot(dat, gene, outcome)

    png_file <- file.path(DIR_PLOTS_PNG, paste0(sprintf("%02d", panel_i), "_", gene, "_", outcome, "_regional_locuscompare.png"))
    pdf_file <- file.path(DIR_PLOTS_PDF, paste0(sprintf("%02d", panel_i), "_", gene, "_", outcome, "_regional_locuscompare.pdf"))

    ggsave(png_file, res$plot, width = PLOT_WIDTH_IN, height = PLOT_HEIGHT_IN, dpi = PLOT_DPI)
    ggsave(pdf_file, res$plot, width = PLOT_WIDTH_IN, height = PLOT_HEIGHT_IN)

    plot_summary[[length(plot_summary) + 1]] <- data.frame(
      Panel = paste0("(", panel_letters[panel_i], ")"),
      Gene = gene,
      Outcome = outcome,
      SNPs_plotted = res$snps_plotted,
      Lead_SNP_labeled = res$lead_snp,
      pQTL_max_neglog10P = res$pQTL_max_neglog10P,
      GWAS_max_neglog10P = res$GWAS_max_neglog10P,
      PNG_file = png_file,
      PDF_file = pdf_file,
      stringsAsFactors = FALSE
    )

    cat("Plot generated:", gene, outcome, "\n")
    panel_i <- panel_i + 1
  }
}

plot_summary <- data.table::rbindlist(plot_summary, fill = TRUE)
openxlsx::write.xlsx(
  plot_summary,
  file.path(DIR_TABLES, "original_loci_regional_plot_summary.xlsx"),
  overwrite = TRUE
)

# 合并成 3 列 × 6 行的大图
if (nrow(plot_summary) == 18) {
  plot_objects <- list()
  for (idx in seq_len(nrow(plot_summary))) {
    dat <- data.table::fread(file.path(
      DIR_HARMONISED,
      paste0(plot_summary$Gene[idx], "_", plot_summary$Outcome[idx], "_harmonised.tsv.gz")
    ))
    dat <- dat[!is.na(p_pqtl) & !is.na(p_gwas) & p_pqtl > 0 & p_gwas > 0]
    p <- make_one_locuscompare_plot(dat, plot_summary$Gene[idx], plot_summary$Outcome[idx])$plot
    p <- p + plot_annotation(tag_levels = NULL, tag_prefix = "", tag_suffix = "")
    plot_objects[[idx]] <- p
  }

  combined_all <- wrap_plots(plot_objects, ncol = 3)
  ggsave(file.path("05_regional_plots", "Supplementary_Figure_S1_18_panels.png"),
         combined_all, width = 18, height = 28, dpi = 300)
  ggsave(file.path("05_regional_plots", "Supplementary_Figure_S1_18_panels.pdf"),
         combined_all, width = 18, height = 28)
}

write_status("05_make_locuscompare_plots", TRUE, "Individual plots and plot summary generated.")