#!/usr/bin/env Rscript
# ---------------------------------------------------------------------------
# z93_distribution_analysis.R (2026-08-18)
# Country x clade x era analysis of R1a-Z93 for the distribution blog post.
# Inputs:  runs/ydna-z93-aadr-ancients-2026-07-24.csv   (212 ancient males)
#          scratch JSON of YFull v14.05.00 modern strata flags (path via arg)
# Outputs: runs/z93-distribution-analysis-2026-08-18.md + two CSVs
# Regime: counts and proportions only; no model claims. Fisher tests where
# expected counts allow; correspondence analysis on the era x clade table.
# ---------------------------------------------------------------------------
args <- commandArgs(trailingOnly = TRUE)
flags_json <- ifelse(length(args) >= 1, args[1], stop("need strata JSON path"))
suppressMessages({ library(jsonlite) })

ROOT <- "/home/cdr/projects/sohi-qpadm"
anc  <- read.csv(file.path(ROOT, "runs/ydna-z93-aadr-ancients-2026-07-24.csv"),
                 stringsAsFactors = FALSE)
stopifnot(nrow(anc) == 212)

era <- function(bp) {
  bp <- suppressWarnings(as.numeric(bp))
  cut(bp, breaks = c(-Inf, 1000, 1800, 2800, 3500, 4300, Inf),
      labels = c("6 Medieval+", "5 Migration-era 150-950CE",
                 "4 Scythian-Saka-Xiongnu 850BCE-150CE",
                 "3 LBA-Iron 1550-850BCE", "2 Sintashta-Andronovo 2350-1550BCE",
                 "1 pre-Sintashta >2350BCE"))
}
anc$era <- era(anc$date_bp)
anc$clade <- anc$relation

sink(file.path(ROOT, "runs/z93-distribution-analysis-2026-08-18.md"))
cat("# Z93 distribution analysis - 2026-08-18\n")
cat("Inputs: ydna-z93-aadr-ancients-2026-07-24.csv (n=212, AADR v66 x YFull v14.04)\n")
cat("        + YFull v14.05.00 modern strata flags (curl 2026-08-18)\n\n")

## 1. era x clade table -------------------------------------------------------
t1 <- table(anc$era, anc$clade)
cat("## 1. Ancient era x clade counts\n\n")
print(t1)
z2124_share <- prop.table(t1, 1)[, "Z94 Z2124 sibling"]
cat("\nZ2124-arm share by era:\n"); print(round(z2124_share, 3))

## 2. era x country -----------------------------------------------------------
t2 <- table(anc$era, anc$entity)
keep <- colSums(t2) >= 4
cat("\n## 2. Ancient era x country (countries with n>=4)\n\n")
print(t2[, keep])

## 3. clade x country ---------------------------------------------------------
t3 <- table(anc$clade, anc$entity)
cat("\n## 3. Ancient clade x country (countries with n>=4)\n\n")
print(t3[, colSums(t3) >= 4])

## 4. Fisher test: does clade composition differ between steppe core and east?
grp <- ifelse(anc$entity %in% c("Russia", "Ukraine", "Moldova", "Hungary",
                                "Serbia", "Bulgaria", "Romania"), "west-steppe",
       ifelse(anc$entity %in% c("Kazakhstan", "Kyrgyzstan", "China", "Mongolia",
                                "Uzbekistan", "Turkmenistan"), "east-steppe", NA))
sub <- anc[!is.na(grp) & anc$clade %in%
             c("Z93 basal / unresolved", "Z94 Z2124 sibling"), ]
t4 <- table(grp[!is.na(grp) & anc$clade %in%
                  c("Z93 basal / unresolved", "Z94 Z2124 sibling")], sub$clade)
cat("\n## 4. West-steppe vs east-steppe: basal Z93 vs Z2124 arm\n\n")
print(t4)
ft <- fisher.test(t4)
cat(sprintf("\nFisher exact: OR = %.2f, p = %.4g (basal-Z93 enrichment in the west)\n",
            ft$estimate, ft$p.value))

## 5. temporal trend of Z2124-arm dominance (logistic on date) ---------------
sub5 <- anc[!is.na(suppressWarnings(as.numeric(anc$date_bp))), ]
sub5$is2124 <- as.integer(sub5$clade == "Z94 Z2124 sibling")
sub5$kyr <- as.numeric(sub5$date_bp) / 1000
m <- glm(is2124 ~ kyr, family = binomial, data = sub5)
cat("\n## 5. Logistic trend: P(Z2124-arm | date)\n\n")
print(summary(m)$coefficients)
cat("Interpretation: sign of kyr coefficient <0 means the Z2124 arm's share of\n")
cat("sampled Z93 RISES toward the present.\n")

## 6. correspondence analysis on era x clade (base-R SVD) --------------------
P <- t1 / sum(t1); r <- rowSums(P); c <- colSums(P)
S <- diag(1/sqrt(r)) %*% (as.matrix(P) - r %o% c) %*% diag(1/sqrt(c))
sv <- svd(S)
rowc <- diag(1/sqrt(r)) %*% sv$u[, 1:2] %*% diag(sv$d[1:2])
colc <- diag(1/sqrt(c)) %*% sv$v[, 1:2] %*% diag(sv$d[1:2])
rownames(rowc) <- rownames(t1); rownames(colc) <- colnames(t1)
cat("\n## 6. Correspondence analysis (dim1 inertia share = ",
    round(sv$d[1]^2 / sum(sv$d^2), 3), ")\n\nEra scores:\n", sep = "")
print(round(rowc, 3)); cat("\nClade scores:\n"); print(round(colc, 3))

## 7. modern strata: country composition + diversity --------------------------
fl <- fromJSON(flags_json)
countries <- c("India","Pakistan","Bangladesh","Sri Lanka","Afghanistan","Iran",
  "Iraq","Turkey","Kyrgyzstan","Kazakhstan","Russia","China","Saudi Arabia",
  "Qatar","Kuwait","United Arab Emirates","Bahrain","Oman","Yemen","Syria",
  "Lebanon","Jordan","Palestine","Israel","Armenia","Azerbaijan","Georgia",
  "Ukraine","Belarus","Poland","Germany","England","Hungary","Mongolia","Nepal",
  "Uzbekistan","Tajikistan","Turkmenistan","Singapore","Albania","Moldova")
mat <- sapply(fl, function(s) {
  vapply(countries, function(cn) { x <- s[[cn]]; ifelse(is.null(x), 0, as.numeric(x)) },
         numeric(1))
})
rownames(mat) <- countries
mat <- mat[rowSums(mat) > 0, ]
write.csv(mat, file.path(ROOT, "runs/z93-modern-country-x-stratum-2026-08-18.csv"))
cat("\n## 7. Modern YFull country x stratum matrix (countries, n>0)\n\n")
print(mat[order(-rowSums(mat)), ])
sh <- apply(mat, 1, function(x) { p <- x[x>0]/sum(x); -sum(p*log(p)) })
cat("\nShannon diversity of stratum composition per country (top 12 by n):\n")
top <- names(sort(rowSums(mat), decreasing = TRUE))[1:12]
print(round(sh[top], 2))

## 8. ethnic/language tags per stratum (top 5, NOT countries) ----------------
cat("\n## 8. Leading ethnic/language tags per stratum (tester self-report)\n\n")
for (s in names(fl)) {
  v <- unlist(fl[[s]]); v <- v[!(names(v) %in% c(countries, "Ancient DNA"))]
  v <- sort(v, decreasing = TRUE)[1:5]
  cat(s, ": ", paste(names(v), v, sep = "=", collapse = ", "), "\n", sep = "")
}
sink()

anc_out <- anc[, c("genetic_id","date_bp","year_bce","entity","locality",
                   "matched_node","clade","era","publication")]
write.csv(anc_out, file.path(ROOT, "runs/z93-ancient-era-clade-2026-08-18.csv"),
          row.names = FALSE)
cat("done\n")