#!/usr/bin/env Rscript
# =============================================================================
# MCNV2 - Test Data Simulation
# All parameters are estimated DIRECTLY from real data.
#
# NOTE: This script requires the annotated inheritance file from MCNV2
# (output of compute_inheritance()) as input.
# The input file used to calibrate the simulation parameters cannot be
# shared due to data privacy constraints (SPARK cohort).
# The simulated output files (sim_cnvs.tsv, sim_pedigree.tsv) are
# available at: https://zenodo.org/records/19860597
#
# Usage:
#   Rscript simulate_mcnv2_testdata.R --real_data /path/to/cnvs_inheritance.tsv
# =============================================================================

suppressPackageStartupMessages({
  library(dplyr)
  library(readr)
})

# ── Command-line arguments ───────────────────────────────────────────────────
args <- commandArgs(trailingOnly = TRUE)
get_arg <- function(flag, default=NULL) {
  i <- which(args == flag)
  if (length(i) && i < length(args)) args[i+1] else default
}

REAL_DATA  <- get_arg("--real_data",  "/path/to/your/cnvs_inheritance.tsv")
OUT_CNVS   <- get_arg("--out_cnvs",   "sim_cnvs.tsv")
OUT_PED    <- get_arg("--out_ped",    "sim_pedigree.tsv")
SEED       <- as.integer(get_arg("--seed", "42"))
set.seed(SEED)  # Ensures full reproducibility

cat("=================================================================\n")
cat("MCNV2 - Simulation from real data distributions\n")
cat(sprintf("  Input file : %s\n", REAL_DATA))
cat("=================================================================\n\n")

# =============================================================================
# STEP 1: Load and prepare real data
# =============================================================================
cat("=== Step 1: Loading real data ===\n")

if (!file.exists(REAL_DATA)) stop("File not found: ", REAL_DATA)

df_raw <- read_tsv(REAL_DATA, show_col_types=FALSE)

# Detect family_status column (handles typo variant 'family_statue')
fs_col <- intersect(c("family_status","family_statue"), colnames(df_raw))[1]
if (is.na(fs_col)) stop("Column family_status not found")

# Detect algos_overlap column (handles multiple naming conventions)
ov_col <- intersect(c("algos_overlap","twoAlgs","TwoAlgos","two_algos","AlgosOverlap"),
                    colnames(df_raw))[1]
if (!is.na(ov_col) && ov_col != "algos_overlap") {
  df_raw <- df_raw %>% rename(algos_overlap = all_of(ov_col))
  cat(sprintf("  Overlap column: '%s' renamed to 'algos_overlap'\n", ov_col))
}
# Normalize to 0-1 scale if stored as 0-100
if (!is.na(ov_col) && max(df_raw$algos_overlap, na.rm=TRUE) > 1) {
  df_raw <- df_raw %>% mutate(algos_overlap = algos_overlap / 100)
  cat("  0-100 scale detected -> normalized to 0-1\n")
}
if (is.na(ov_col)) {
  cat("  WARN: algos_overlap column not found -> set to NA\n")
  df_raw$algos_overlap <- NA_real_
}

# Detect end position column (Stop or End)
stop_col <- intersect(c("Stop", "End"), colnames(df_raw))[1]
if (is.na(stop_col)) stop("Stop/End column not found")

# Keep children only, deduplicate CNVs
df <- df_raw %>%
  filter(.data[[fs_col]] == "child") %>%
  rename(Stop = all_of(stop_col)) %>%
  distinct(SampleID, Chr, Start, Stop, Type, .keep_all=TRUE) %>%
  mutate(
    Size          = as.numeric(Stop) - as.numeric(Start),
    transmitted   = as.character(transmitted_cnv) %in% c("True","TRUE","true","1"),
    NbreAlgos     = as.integer(NbreAlgos),
    algos_overlap = as.numeric(algos_overlap),
    Score         = as.numeric(Score),
    SNP           = as.numeric(SNP)
  ) %>%
  filter(!is.na(Size), Size > 0)

cat(sprintf("  Unique CNVs (children): %d\n", nrow(df)))

# Extract pedigree structure from real data
ped_col <- intersect(c("TrioKey","trio_key"), colnames(df_raw))[1]
real_trios <- df_raw %>%
  filter(.data[[fs_col]] == "child") %>%
  distinct(SampleID, .data[[ped_col]])
N_TRIOS <- nrow(real_trios)
cat(sprintf("  Real trios: %d\n\n", N_TRIOS))

# =============================================================================
# STEP 2: Estimate simulation parameters from real data
# All parameters are derived empirically — no arbitrary values used.
# =============================================================================
cat("=== Step 2: Parameter estimation ===\n")

# ── 2a. CNVs per child — Negative Binomial distribution ───────────────────
# The negative binomial models overdispersed count data (variance > mean)
cnvs_per_child <- df %>% count(SampleID)
nb_mu   <- mean(cnvs_per_child$n)           # mean CNVs per child
nb_var  <- var(cnvs_per_child$n)             # variance
nb_size <- nb_mu^2 / (nb_var - nb_mu)        # dispersion parameter
cat(sprintf("  CNVs/child   : mu=%.1f, size=%.2f\n", nb_mu, nb_size))

# ── 2b. CNV type proportions (DEL/DUP) ────────────────────────────────────
type_probs <- prop.table(table(df$Type))
cat(sprintf("  Types        : DEL=%.1f%%, DUP=%.1f%%\n",
            type_probs["DEL"]*100, type_probs["DUP"]*100))

# ── 2c. CNV size by type — log-normal parameters ──────────────────────────
size_params <- df %>%
  filter(!is.na(Size), Size > 0) %>%
  group_by(Type) %>%
  summarise(log10_mean=mean(log10(Size)), log10_sd=sd(log10(Size)), .groups="drop")
cat("  Size by type:\n")
for (i in seq_len(nrow(size_params)))
  cat(sprintf("    %s : log10_mean=%.3f, log10_sd=%.3f\n",
              size_params$Type[i], size_params$log10_mean[i], size_params$log10_sd[i]))

# ── 2d. Chromosome assignment distribution ────────────────────────────────
# Chromosomes are sampled directly from observed values (preserves exact frequencies)
chr_real <- df$Chr
cat(sprintf("  Chromosomes  : %d distinct values\n", n_distinct(chr_real)))

# ── 2e. Number of supporting algorithms ───────────────────────────────────
algo_probs <- prop.table(table(df$NbreAlgos))
cat(sprintf("  NbreAlgos    : 1=%.1f%%, 2=%.1f%%\n",
            algo_probs["1"]*100, algo_probs["2"]*100))

# ── 2f. Inheritance probability — stratified lookup table ─────────────────
# P(inherited | Type x size_bin x NbreAlgos) — full stratification
# This preserves the observed correlation between inheritance and CNV properties
size_breaks <- c(0, 30000, 50000, 100000, 200000, 500000, 1e6, Inf)
size_labels <- c("1-30kb","30-50kb","50-100kb","100-200kb","200-500kb","500-1M",">1M")

p_inh_strat <- df %>%
  mutate(size_bin = cut(Size, breaks=size_breaks, labels=size_labels, right=FALSE),
         key = paste0(Type,"_",size_bin,"_a",NbreAlgos)) %>%
  group_by(key) %>%
  summarise(p=mean(transmitted, na.rm=TRUE), n=n(), .groups="drop") %>%
  filter(n >= 5)   # exclude strata with insufficient data
p_inh_lookup <- setNames(p_inh_strat$p, p_inh_strat$key)

# Progressive fallbacks: NbreAlgos only, then global rate
p_inh_by_algos <- df %>%
  group_by(NbreAlgos) %>%
  summarise(p=mean(transmitted, na.rm=TRUE), .groups="drop")
p_inh_fb <- setNames(p_inh_by_algos$p, p_inh_by_algos$NbreAlgos)

global_inh <- mean(df$transmitted, na.rm=TRUE)
cat("  P(inherited | Type x size_bin x NbreAlgos) — strata with n>=5:\n")
for (nm in names(p_inh_lookup))
  cat(sprintf("    %s -> %.3f\n", nm, p_inh_lookup[nm]))
cat(sprintf("  Global rate  : %.3f\n", global_inh))

# ── 2g. Score and SNP — empirical resampling tables ───────────────────────
# Resampling is stratified by size bin x NbreAlgos x inheritance status
# This preserves observed correlations between quality metrics and CNV properties
df_keyed <- df %>%
  filter(!is.na(Score), !is.na(SNP)) %>%
  mutate(
    size_bin = cut(Size, breaks=size_breaks, labels=size_labels, right=FALSE),
    key = paste0(size_bin, "_a", NbreAlgos, "_", ifelse(transmitted, "inh", "dn"))
  )

score_table <- df_keyed %>%
  group_by(key) %>%
  summarise(vals = list(Score), .groups="drop")
score_lookup <- setNames(score_table$vals, score_table$key)

snp_table <- df_keyed %>%
  group_by(key) %>%
  summarise(vals = list(SNP), .groups="drop")
snp_lookup <- setNames(snp_table$vals, snp_table$key)

# Fallback tables by NbreAlgos x inheritance status (if size bin is missing)
score_fallback <- df %>%
  filter(!is.na(Score)) %>%
  mutate(key = paste0("a", NbreAlgos, "_", ifelse(transmitted, "inh", "dn"))) %>%
  group_by(key) %>% summarise(vals=list(Score), .groups="drop")
score_fb <- setNames(score_fallback$vals, score_fallback$key)

snp_fallback <- df %>%
  filter(!is.na(SNP)) %>%
  mutate(key = paste0("a", NbreAlgos, "_", ifelse(transmitted, "inh", "dn"))) %>%
  group_by(key) %>% summarise(vals=list(SNP), .groups="drop")
snp_fb <- setNames(snp_fallback$vals, snp_fallback$key)

cat("  Score/SNP    : empirical resampling by size bin x NbreAlgos x inheritance\n")
cat(sprintf("  Filled bins  : Score=%d, SNP=%d\n",
            length(score_lookup), length(snp_lookup)))

# Quality check: verify Score correlation with inheritance status in real data
score_check <- df %>%
  filter(!is.na(Score)) %>%
  group_by(NbreAlgos, transmitted) %>%
  summarise(med_score=median(Score), n=n(), .groups="drop")
cat("  Score median | NbreAlgos x inheritance:\n")
for (i in seq_len(nrow(score_check)))
  cat(sprintf("    NbreAlgos=%d, inherited=%s : median=%.0f (n=%d)\n",
              score_check$NbreAlgos[i], score_check$transmitted[i],
              score_check$med_score[i], score_check$n[i]))

# ── 2h. Algorithm overlap — Beta distribution ─────────────────────────────
# For NbreAlgos=2 only; parameters estimated separately for inherited and de novo
# Beta distribution models proportions bounded in (0, 1)
fit_beta <- function(vals) {
  vals <- vals[!is.na(vals) & vals > 0 & vals < 1]
  m <- mean(vals); v <- min(var(vals), m*(1-m)*0.99)
  k <- m*(1-m)/v - 1
  list(shape1=m*k, shape2=(1-m)*k)
}

ov_data <- df %>% filter(NbreAlgos == 2, !is.na(algos_overlap))
beta_inh    <- fit_beta(ov_data$algos_overlap[ov_data$transmitted])
beta_denovo <- fit_beta(ov_data$algos_overlap[!ov_data$transmitted])
cat(sprintf("  Beta overlap (inherited) : shape1=%.2f, shape2=%.2f -> median~%.2f\n",
            beta_inh$shape1, beta_inh$shape2,
            qbeta(0.5, beta_inh$shape1, beta_inh$shape2)))
cat(sprintf("  Beta overlap (de novo)   : shape1=%.2f, shape2=%.2f -> median~%.2f\n",
            beta_denovo$shape1, beta_denovo$shape2,
            qbeta(0.5, beta_denovo$shape1, beta_denovo$shape2)))

# ── 2i. Problematic region (PR) positions — resampling from real positions ─
# A fraction of CNVs are placed at real PR coordinates to preserve
# the observed overlap rate with problematic genomic regions
pr_col <- intersect(c("cnv_problematic_region_overlap","prb_region_overlap","pr_overlap"),
                    colnames(df_raw))[1]
if (!is.na(pr_col)) {
  pr_positions <- df_raw %>%
    filter(.data[[fs_col]] == "child") %>%
    distinct(SampleID, Chr, Start, .data[[stop_col]], Type, .keep_all=TRUE) %>%
    filter(as.numeric(.data[[pr_col]]) > 50) %>%   # CNVs with >50% PR overlap
    select(Chr, Start, Stop=all_of(stop_col), Type) %>%
    mutate(Start=as.integer(Start), Stop=as.integer(Stop))
  cat(sprintf("  Real PR CNVs : %d positions (overlap > 50%%)\n", nrow(pr_positions)))
  PROP_IN_PR <- min(nrow(pr_positions) / nrow(df %>% distinct(SampleID, Chr, Start, Stop, Type)), 0.20)
  cat(sprintf("  PR proportion: %.3f\n", PROP_IN_PR))
} else {
  pr_positions <- NULL
  PROP_IN_PR   <- 0
  cat("  WARN: PR overlap column not found -> no PR simulation\n")
}

# ── 2j. GRCh38 chromosome lengths ─────────────────────────────────────────
# Used to ensure simulated CNV positions fall within valid chromosome boundaries
chr_lengths <- c(
  "chr1"=248956422,"chr2"=242193529,"chr3"=198295559,"chr4"=190214555,
  "chr5"=181538259,"chr6"=170805979,"chr7"=159345973,"chr8"=145138636,
  "chr9"=138394717,"chr10"=133797422,"chr11"=135086622,"chr12"=133275309,
  "chr13"=114364328,"chr14"=107043718,"chr15"=101991189,"chr16"=90338345,
  "chr17"=83257441,"chr18"=80373285,"chr19"=58617616,"chr20"=64444167,
  "chr21"=46709983,"chr22"=50818468
)

cat("\n")

# =============================================================================
# STEP 3: Simulate CNVs and pedigree
# =============================================================================
cat("=== Step 3: Simulation ===\n")

# ── Helper functions ───────────────────────────────────────────────────────

# Size resampling: sample directly from observed sizes (preserves exact distribution)
real_sizes <- df %>%
  filter(!is.na(Size), Size > 0) %>%
  select(Type, Size)

sim_size <- function(type) {
  pool <- real_sizes$Size[real_sizes$Type == type]
  if (length(pool) == 0) pool <- real_sizes$Size
  sample(pool, 1)
}

# Generate one CNV row with all quality metrics
sim_row <- function(sample_id, chr, start, size, type, nbre_algos, is_inherited) {
  # Build lookup key: size bin x NbreAlgos x inheritance status
  sbin     <- as.character(cut(size, breaks=size_breaks, labels=size_labels, right=FALSE))
  inh_str  <- ifelse(is_inherited, "inh", "dn")
  key_full <- paste0(sbin, "_a", nbre_algos, "_", inh_str)
  key_fb   <- paste0("a", nbre_algos, "_", inh_str)

  # Progressive fallback: full key -> NbreAlgos key -> global
  resample <- function(lookup, fb, key_full, key_fb) {
    if (!is.null(lookup[[key_full]]) && length(lookup[[key_full]]) > 0)
      return(sample(lookup[[key_full]], 1))
    if (!is.null(fb[[key_fb]]) && length(fb[[key_fb]]) > 0)
      return(sample(fb[[key_fb]], 1))
    return(sample(fb[[1]], 1))
  }

  score <- resample(score_lookup, score_fb, key_full, key_fb)
  snp   <- resample(snp_lookup,   snp_fb,   key_full, key_fb)

  # Algorithm overlap: 0 if single algorithm, Beta-distributed if two algorithms
  if (nbre_algos == 1) {
    ov <- 0
  } else if (is_inherited) {
    ov <- round(rbeta(1, beta_inh$shape1, beta_inh$shape2), 4)
  } else {
    ov <- round(rbeta(1, beta_denovo$shape1, beta_denovo$shape2), 4)
  }

  data.frame(Chr=chr, Start=as.integer(start), Stop=as.integer(start+size-1L),
             Type=type, SampleID=sample_id,
             Score=score, SNP=snp, NbreAlgos=nbre_algos, algos_overlap=ov,
             stringsAsFactors=FALSE)
}

# ── Simulated pedigree ─────────────────────────────────────────────────────
# N_TRIOS complete trios with synthetic IDs (no real sample IDs)
pedigree <- data.frame(
  ChildID  = sprintf("SIM_C_%04d", seq_len(N_TRIOS)),
  FatherID = sprintf("SIM_F_%04d", seq_len(N_TRIOS)),
  MotherID = sprintf("SIM_M_%04d", seq_len(N_TRIOS)),
  stringsAsFactors=FALSE
)

# ── Simulate CNVs for all individuals ─────────────────────────────────────
all_rows <- list()

for (i in seq_len(N_TRIOS)) {
  child_id  <- pedigree$ChildID[i]
  father_id <- pedigree$FatherID[i]
  mother_id <- pedigree$MotherID[i]

  # Draw number of CNVs for this child from negative binomial
  n_cnvs <- max(rnbinom(1, mu=nb_mu, size=nb_size), 1L)

  child_rows  <- vector("list", n_cnvs)
  parent_rows <- list()

  for (j in seq_len(n_cnvs)) {
    # Simulate CNV properties
    type      <- sample(names(type_probs), 1, prob=as.numeric(type_probs))
    size      <- sim_size(type)
    chr       <- sample(chr_real, 1)  # sampled from observed chromosome distribution
    chr_key   <- if (startsWith(chr,"chr")) chr else paste0("chr",chr)
    chr_len   <- chr_lengths[chr_key]
    if (is.na(chr_len)) chr_len <- 100000000L
    max_start <- chr_len - size - 1000L
    start     <- if (max_start > 1000L) sample(1000L:max_start, 1) else 1000L

    # Place a fraction of CNVs at real PR coordinates (preserves observed PR overlap rate)
    use_pr <- !is.null(pr_positions) && nrow(pr_positions) > 0 &&
              runif(1) < PROP_IN_PR
    in_pr <- FALSE
    if (use_pr) {
      pr_row <- pr_positions[sample(nrow(pr_positions), 1), ]
      chr    <- pr_row$Chr
      start  <- pr_row$Start
      size   <- pr_row$Stop - pr_row$Start
      type   <- pr_row$Type
      in_pr  <- TRUE
    }

    nbre_algos <- sample(as.integer(names(algo_probs)), 1, prob=as.numeric(algo_probs))

    # Lookup inheritance probability with progressive fallback
    sbin_j    <- as.character(cut(size, breaks=size_breaks, labels=size_labels, right=FALSE))
    key_strat <- paste0(type, "_", sbin_j, "_a", nbre_algos)
    key_fb    <- as.character(nbre_algos)
    p_inh_this <- if (!is.null(p_inh_lookup[key_strat]) && !is.na(p_inh_lookup[key_strat])) {
      p_inh_lookup[key_strat]
    } else if (!is.na(p_inh_fb[key_fb])) {
      p_inh_fb[key_fb]
    } else {
      global_inh
    }

    # CNVs in PR regions are mostly false positives -> set as non-inherited
    is_inherited <- if (in_pr) FALSE else runif(1) < p_inh_this

    child_rows[[j]] <- sim_row(child_id, chr, start, size, type, nbre_algos, is_inherited)

    # For inherited CNVs: generate matching parental CNV with slight position jitter
    # Jitter ensures >=95% reciprocal overlap (realistic caller variability)
    if (is_inherited) {
      parent_id   <- sample(c(father_id, mother_id), 1)
      jitter      <- runif(1, 0, 0.05)
      p_start     <- max(1000L, round(start - size*jitter))
      p_size      <- round(size * (1 + 2*jitter))
      p_algos     <- sample(as.integer(names(algo_probs)), 1, prob=as.numeric(algo_probs))
      parent_rows <- c(parent_rows,
                       list(sim_row(parent_id, chr, p_start, p_size, type, p_algos, TRUE)))
    }
  }

  # Generate non-transmitted parental CNVs (independent of child CNVs)
  for (p_id in c(father_id, mother_id)) {
    n_extra <- max(rnbinom(1, mu=nb_mu*0.3, size=nb_size), 0L)
    for (k in seq_len(n_extra)) {
      type    <- sample(names(type_probs), 1, prob=as.numeric(type_probs))
      size    <- sim_size(type)
      chr     <- sample(chr_real, 1)
      chr_key <- if (startsWith(chr,"chr")) chr else paste0("chr",chr)
      chr_len <- chr_lengths[chr_key]
      if (is.na(chr_len)) chr_len <- 100000000L
      max_s   <- chr_len - size - 1000L
      st      <- if (max_s > 1000L) sample(1000L:max_s, 1) else 1000L
      na      <- sample(as.integer(names(algo_probs)), 1, prob=as.numeric(algo_probs))
      parent_rows <- c(parent_rows,
                       list(sim_row(p_id, chr, st, size, type, na, FALSE)))
    }
  }

  all_rows <- c(all_rows, child_rows, parent_rows)
}

cnvs <- bind_rows(all_rows)
cat(sprintf("  CNVs generated: %d (%d individuals, %d trios)\n",
            nrow(cnvs), n_distinct(cnvs$SampleID), N_TRIOS))

# =============================================================================
# STEP 4: Validation
# Compare simulated distributions to real data to verify calibration quality
# =============================================================================
cat("\n=== Step 4: Validation ===\n")
child_cnvs <- cnvs %>% filter(SampleID %in% pedigree$ChildID)
cpp <- child_cnvs %>% count(SampleID)
tp  <- prop.table(table(child_cnvs$Type))
sz  <- child_cnvs$Stop - child_cnvs$Start

cat(sprintf("CNVs/child   — mean: %.1f (real: %.1f) | sd: %.1f (real: %.1f)\n",
            mean(cpp$n), nb_mu, sd(cpp$n), sqrt(nb_mu^2/nb_size + nb_mu)))
cat(sprintf("Types        — DEL: %.1f%% (real: %.1f%%) | DUP: %.1f%% (real: %.1f%%)\n",
            tp["DEL"]*100, type_probs["DEL"]*100,
            tp["DUP"]*100, type_probs["DUP"]*100))
cat(sprintf("Median size  — %.0f bp (real: %.0f bp)\n",
            median(sz), median(df$Size)))
cat(sprintf("NbreAlgos=1  — %.1f%% (real: %.1f%%)\n",
            mean(child_cnvs$NbreAlgos==1)*100, algo_probs["1"]*100))

# =============================================================================
# STEP 5: Save output files
# =============================================================================
write_tsv(cnvs,     OUT_CNVS)
write_tsv(pedigree, OUT_PED)
cat(sprintf("\n-> %s (%d rows)\n",  OUT_CNVS, nrow(cnvs)))
cat(sprintf("-> %s (%d trios)\n",  OUT_PED,  nrow(pedigree)))
