# =========================
# 03 根据区域 SNP 查询 pQTL 数据
# =========================

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

library(ieugwasr)

loci <- read_loci_file()
unique_loci <- unique(loci[, .(Gene, Protein_ID, Protein, chr, region_start, region_end)])

jwt_ok <- nchar(Sys.getenv("OPENGWAS_JWT")) > 0
cat("OpenGWAS JWT exists:", jwt_ok, "\n")
if (!jwt_ok) {
  stop("OPENGWAS_JWT is empty. Please set token in .Renviron before querying pQTL.")
}

query_pqtl_for_snps <- function(snps, protein_id, chunk_size = PQTL_QUERY_CHUNK_SIZE) {
  snps <- unique(snps[!is.na(snps) & snps != ""])
  if (length(snps) == 0) return(data.table::data.table())

  chunks <- split(snps, ceiling(seq_along(snps) / chunk_size))
  res_list <- list()

  for (k in seq_along(chunks)) {
    cat("Querying pQTL:", protein_id, "chunk", k, "of", length(chunks), "\n")

    tmp <- tryCatch(
      {
        data.table::as.data.table(
          ieugwasr::associations(
            variants = chunks[[k]],
            id = protein_id,
            proxies = 0,
            opengwas_jwt = Sys.getenv("OPENGWAS_JWT")
          )
        )
      },
      error = function(e) {
        warning("pQTL query failed for ", protein_id, " chunk ", k, ": ", conditionMessage(e))
        data.table::data.table()
      }
    )

    if (nrow(tmp) > 0) res_list[[length(res_list) + 1]] <- tmp
    Sys.sleep(SLEEP_SECONDS_BETWEEN_QUERIES)
  }

  data.table::rbindlist(res_list, fill = TRUE)
}

summary_list <- list()

for (i in seq_len(nrow(unique_loci))) {
  gene <- unique_loci$Gene[i]
  protein_id <- unique_loci$Protein_ID[i]
  chr <- unique_loci$chr[i]
  region_start <- unique_loci$region_start[i]
  region_end <- unique_loci$region_end[i]

  cat("Processing pQTL:", gene, protein_id, "\n")

  # 使用三个 outcome 的 GWAS 区域 SNP 并集作为查询 SNP
  gwas_files <- file.path(DIR_GWAS_REGIONS, paste0(gene, "_", OUTCOME_ORDER, "_GWAS_original_region.tsv.gz"))
  gwas_files <- gwas_files[file.exists(gwas_files)]

  if (length(gwas_files) == 0) {
    warning("No GWAS region files for gene: ", gene)
    next
  }

  gwas_union <- data.table::rbindlist(lapply(gwas_files, data.table::fread), fill = TRUE)
  snps <- unique(gwas_union$SNP)

  pqtl_raw <- query_pqtl_for_snps(snps, protein_id)
  if (nrow(pqtl_raw) > 0) {
    # 尽量标准化后再限制区域
    pqtl_std <- standardise_pqtl(pqtl_raw)
    pqtl_std <- pqtl_std[
      chr_pqtl == chr &
        position_pqtl >= region_start &
        position_pqtl <= region_end
    ]
    pqtl_std[, Gene := gene]
    pqtl_std[, Protein_ID := protein_id]
  } else {
    pqtl_std <- data.table::data.table()
  }

  out_file <- file.path(DIR_PQTL_REGIONS, paste0(gene, "_pQTL_original_region.tsv.gz"))
  data.table::fwrite(pqtl_std, out_file, sep = "\t")

  summary_list[[length(summary_list) + 1]] <- data.frame(
    Gene = gene,
    Protein_ID = protein_id,
    chr = chr,
    region_start = region_start,
    region_end = region_end,
    SNPs_queried = length(snps),
    pQTL_SNPs_returned = nrow(pqtl_std),
    Output_file = out_file,
    stringsAsFactors = FALSE
  )
}

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

write_status("03_extract_pqtl_regions", TRUE, "pQTL regional files and summary generated.")