# =============================================================================
# utils_grs.R - Internal helpers for grs_ series
# =============================================================================


# Generate the plink2 driver R script executed on RAP via Swiss Army Knife.
# Faithfully mirrors scripts 06 (standard) and 13 (large), differing only
# in n_threads and plink_memory which are injected at generation time.
#
# CRITICAL: output_dir MUST be "/home/dnanexus/out/out/" - Swiss Army Knife
# auto-uploads everything in that path to the job --destination on RAP.
.bgen2pgen_script <- function(n_threads, plink_memory, maf = 0.01) {
  maf_label <- sprintf("%03d", round(maf * 100))
  c(
    "#!/usr/bin/env Rscript",
    "# Auto-generated by ukbflow::grs_bgen2pgen()",
    "args <- commandArgs(trailingOnly = TRUE)",
    'if (length(args) == 0L) stop("Usage: Rscript grs_bgen2pgen.R <chr>")',
    "chr <- as.integer(args[1L])",
    'if (is.na(chr) || chr < 1L || chr > 22L) stop("chr must be 1-22")',
    "",
    "suppressPackageStartupMessages(library(cli))",
    "",
    'cli_h1("BGEN -> PGEN: Chromosome {chr}")',
    "",
    "# --- Paths (fixed UKB RAP layout) ---",
    'geno_dir    <- "/mnt/project/Bulk/Imputation/UKB imputation from genotype"',
    'bgen_file   <- file.path(geno_dir, paste0("ukb22828_c", chr, "_b0_v3.bgen"))',
    'sample_file <- file.path(geno_dir, "ukb22828_c1_b0_v3.sample")',
    "",
    "# CRITICAL: Swiss Army Knife auto-uploads /home/dnanexus/out/out/ to --destination",
    'output_dir <- "/home/dnanexus/out/out"',
    paste0('out_prefix <- file.path(output_dir, paste0("chr", chr, "_maf', maf_label, '"))'),
    "dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)",
    "",
    paste0("n_threads    <- ", n_threads, "L"),
    paste0("plink_memory <- ", plink_memory, "L"),
    paste0("maf_filter   <- ", maf),
    "",
    "# --- Input checks ---",
    'if (!file.exists(bgen_file)) {',
    '  cli_alert_danger("BGEN not found: {bgen_file}")',
    '  quit(save = "no", status = 1L)',
    "}",
    'if (!file.exists(sample_file)) {',
    '  cli_alert_danger("Sample file not found: {sample_file}")',
    '  quit(save = "no", status = 1L)',
    "}",
    "",
    "# --- Skip if already done ---",
    'if (file.exists(paste0(out_prefix, ".pgen"))) {',
    '  cli_alert_warning("PGEN already exists, skipping.")',
    '  quit(save = "no", status = 0L)',
    "}",
    "",
    "# --- Install / locate plink2 ---",
    'plink_bin <- file.path(Sys.getenv("HOME"), "plink2")',
    'if (!file.exists(plink_bin)) {',
    '  cli_alert_info("Installing plink2 ...")',
    "  # Pinned version - matches original analysis environment on UKB RAP",
    '  url      <- "https://s3.amazonaws.com/plink2-assets/alpha5/plink2_linux_x86_64_20240105.zip"',
    '  dest_zip <- file.path(tempdir(), "plink2.zip")',
    '  download.file(url, dest_zip, quiet = FALSE, mode = "wb")',
    '  unzip(dest_zip, exdir = Sys.getenv("HOME"))',
    '  system(sprintf("chmod +x %s", shQuote(plink_bin)))',
    '  file.remove(dest_zip)',
    '  cli_alert_success("plink2 installed")',
    "} else {",
    '  cli_alert_success("plink2 found")',
    "}",
    "ver <- system(sprintf(\"%s --version\", shQuote(plink_bin)), intern = TRUE)",
    "cli_alert_info(\"Version: {ver[1L]}\")",
    "",
    "# --- plink2: BGEN ->  PGEN ---",
    "plink_cmd <- sprintf(",
    '  "%s --bgen %s ref-first --sample %s --maf %s --make-pgen --out %s --threads %d --memory %d",',
    "  shQuote(plink_bin), shQuote(bgen_file), shQuote(sample_file),",
    "  maf_filter, shQuote(out_prefix), n_threads, plink_memory",
    ")",
    'cli_alert_info("Running plink2 ...")',
    "t0     <- Sys.time()",
    "status <- system(plink_cmd)",
    'elapsed <- round(difftime(Sys.time(), t0, units = "mins"), 1L)',
    "",
    "# --- Check result ---",
    'pgen_file <- paste0(out_prefix, ".pgen")',
    "if (file.exists(pgen_file) && status == 0L) {",
    '  sz <- round(file.size(pgen_file) / 1024^3, 2)',
    '  cli_alert_success("Done in {elapsed} min -- {pgen_file} ({sz} GB)")',
    '  log_file <- paste0(out_prefix, ".log")',
    "  if (file.exists(log_file)) {",
    "    ll    <- readLines(log_file)",
    '    vline <- grep("variants remaining", ll, value = TRUE)',
    '    if (length(vline) > 0L) cli_alert_info("{vline[length(vline)]}")',
    "  }",
    '  quit(save = "no", status = 0L)',
    "} else {",
    '  cli_alert_danger("plink2 failed -- check {out_prefix}.log")',
    '  quit(save = "no", status = 1L)',
    "}"
  )
}


# Generate the plink2 scoring driver R script executed on RAP via Swiss Army Knife.
# n_threads and plink_memory are injected at generation time (same pattern as
# .bgen2pgen_script).  The script accepts three command-line args:
#   <weight_file>   basename of the weight file (available at /mnt/project/<name>)
#   <score_name>    GRS column name written in the output CSV (e.g. GRS_gwas_cscc)
#   <pgen_dir>      directory of PGEN files on RAP
#
# Output: /home/dnanexus/out/out/<score_name>_scores.csv
# (Swiss Army Knife auto-uploads that path to the job --destination on RAP.)
.grs_score_script <- function(n_threads, plink_memory, maf = 0.01) {
  maf_label <- sprintf("%03d", round(maf * 100))
  c(
    "#!/usr/bin/env Rscript",
    "# Auto-generated by ukbflow::grs_score()",
    "args <- commandArgs(trailingOnly = TRUE)",
    'if (length(args) < 3L) stop("Usage: Rscript grs_score_driver.R <weight_file> <score_name> <pgen_dir>")',
    "weight_file <- args[1L]",
    "score_name  <- args[2L]",
    "pgen_dir    <- args[3L]",
    "",
    "suppressPackageStartupMessages({",
    "  library(data.table)",
    "  library(cli)",
    "})",
    "",
    'cli_h1("GRS scoring: {score_name}")',
    "",
    "# --- Paths ---",
    'weight_path <- file.path("/mnt/project", weight_file)',
    'output_dir  <- "/home/dnanexus/out/out"',
    "dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)",
    "",
    paste0("n_threads    <- ", n_threads, "L"),
    paste0("plink_memory <- ", plink_memory, "L"),
    "chromosomes  <- 1:22",
    "",
    "# --- Check weight file ---",
    'if (!file.exists(weight_path)) {',
    '  cli_alert_danger("Weight file not found: {weight_path}")',
    '  quit(save = "no", status = 1L)',
    "}",
    "",
    "# --- Load + standardise weights (take first 3 cols: ID, A1, WEIGHT) ---",
    "w <- fread(weight_path, data.table = TRUE)",
    'if (ncol(w) < 3L) stop("Weight file must have >= 3 columns: ID, A1, WEIGHT")',
    "w <- w[, 1:3]",
    'setnames(w, c("ID", "A1", "WEIGHT"))',
    'cli_alert_success("Loaded {nrow(w)} SNPs")',
    "",
    "# Write standardised weight and variant-list temp files for plink2",
    'wf_tmp <- file.path(tempdir(), "weights_std.txt")',
    'vl_tmp <- file.path(tempdir(), "variants.txt")',
    'fwrite(w, wf_tmp, sep = "\\t")',
    "fwrite(w[, .(ID)], vl_tmp, col.names = FALSE)",
    "",
    "# --- Install / locate plink2 (pinned to alpha5/20240105) ---",
    'plink_bin <- file.path(Sys.getenv("HOME"), "plink2")',
    'if (!file.exists(plink_bin)) {',
    '  cli_alert_info("Installing plink2 ...")',
    '  url      <- "https://s3.amazonaws.com/plink2-assets/alpha5/plink2_linux_x86_64_20240105.zip"',
    '  dest_zip <- file.path(tempdir(), "plink2.zip")',
    '  download.file(url, dest_zip, quiet = FALSE, mode = "wb")',
    '  unzip(dest_zip, exdir = Sys.getenv("HOME"))',
    '  system(sprintf("chmod +x %s", shQuote(plink_bin)))',
    '  file.remove(dest_zip)',
    '  cli_alert_success("plink2 installed")',
    "} else {",
    '  cli_alert_success("plink2 found")',
    "}",
    'ver <- system(sprintf("%s --version", shQuote(plink_bin)), intern = TRUE)',
    'cli_alert_info("plink2 version: {ver[1L]}")',
    "",
    "# --- Score per chromosome ---",
    "chr_results <- list()",
    "failed_chrs <- integer(0L)",
    "",
    "for (chr in chromosomes) {",
    paste0('  pgen_prefix <- file.path(pgen_dir, paste0("chr", chr, "_maf', maf_label, '"))'),
    '  out_prefix  <- file.path(tempdir(), paste0(score_name, "_chr", chr))',
    "",
    "  plink_cmd <- sprintf(",
    '    "%s --pfile %s --extract %s --score %s 1 2 3 header cols=+scoresums --out %s --threads %d --memory %d",',
    "    shQuote(plink_bin), shQuote(pgen_prefix), shQuote(vl_tmp),",
    "    shQuote(wf_tmp), shQuote(out_prefix), n_threads, plink_memory",
    "  )",
    "",
    "  status     <- system(plink_cmd, ignore.stdout = TRUE, ignore.stderr = TRUE)",
    '  score_file <- paste0(out_prefix, ".sscore")',
    "",
    "  if (status == 0L && file.exists(score_file)) {",
    "    chr_dt <- fread(score_file, data.table = TRUE)",
    '    setnames(chr_dt, gsub("^#", "", names(chr_dt)))',
    '    if (nrow(chr_dt) > 0L && "IID" %in% names(chr_dt) && "SCORE1_SUM" %in% names(chr_dt)) {',
    '      chr_results[[as.character(chr)]] <- chr_dt[, .(IID, SCORE1_SUM)]',
    '      cli_alert_success("chr{chr}: {format(nrow(chr_dt), big.mark=\',\')} individuals")',
    "    } else {",
    "      failed_chrs <- c(failed_chrs, chr)",
    '      cli_alert_warning("chr{chr}: empty result (no matching SNPs?)")',
    "    }",
    "  } else {",
    "    failed_chrs <- c(failed_chrs, chr)",
    '    cli_alert_warning("chr{chr}: plink2 failed (exit status {status})")',
    "  }",
    "}",
    "",
    'if (length(chr_results) == 0L) stop("All chromosomes failed for ", score_name)',
    'cli_alert_info("{length(chr_results)}/22 chromosomes succeeded; {length(failed_chrs)} failed")',
    "",
    "# --- Merge chromosomes: sum SCORE1_SUM across all successful chromosomes ---",
    "# nxt[merged, ...] is a right-join on merged; NA score_next (chr absent) -> 0",
    "merged <- chr_results[[1L]]",
    'setnames(merged, "SCORE1_SUM", "score_sum")',
    "",
    "if (length(chr_results) > 1L) {",
    "  for (k in 2:length(chr_results)) {",
    "    nxt <- chr_results[[k]]",
    '    setnames(nxt, "SCORE1_SUM", "score_next")',
    '    merged <- nxt[merged, on = "IID"]',
    "    merged[is.na(score_next), score_next := 0]",
    "    merged[, score_sum := score_sum + score_next]",
    "    merged[, score_next := NULL]",
    "  }",
    "}",
    "",
    'setnames(merged, "score_sum", score_name)',
    "",
    "# --- Summary ---",
    "v <- merged[[score_name]]",
    'cli_alert_success("N={format(nrow(merged), big.mark=\',\')}, mean={round(mean(v), 4)}, sd={round(sd(v), 4)}, range=[{round(min(v), 4)}, {round(max(v), 4)}]")',
    "",
    "# --- Save to Swiss Army Knife auto-upload path ---",
    'out_file <- file.path(output_dir, paste0(score_name, "_scores.csv"))',
    "fwrite(merged, out_file)",
    'cli_alert_success("Saved: {out_file}")',
    "",
    'quit(save = "no", status = 0L)'
  )
}
