# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                                                             #
#   02-meta-analysis.R                                                         #
#   Bayesian AD meta-analysis – continuous outcome (y)                        #
#   Stage 1: per-study Cohen's d from baseline-adjusted lm()                 #
#   Stage 2: random-effects pooling in JAGS                                   #
#   Prior:   mu ~ N(0, 1e-6);  tau ~ half-normal(0, 0.5)                     #
#                                                                             #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #


# 0. Packages ----------------------------------------------------------------

library(readxl)
library(dplyr)
library(purrr)
library(runjags)
library(coda)
library(HDInterval)
library(ggplot2)
library(stringr)
library(ggridges)


# 1. Helper: pooled within-group SD ------------------------------------------

pooled.sd = function(y, trt) {
  groups = split(y, trt)
  groups = Filter(function(g) length(g) > 1, groups)
  df = sapply(groups, function(g) length(g) - 1)
  ss = sapply(groups, function(g) var(g, na.rm = TRUE) * (length(g) - 1))
  sqrt(sum(ss) / sum(df))
}


# 2. Stage 1: per-study Cohen's d from baseline-adjusted lm() ---------------

path   = "_data/data.xlsx"
sheets = readxl::excel_sheets(path)

es.list = purrr::map_dfr(sheets, function(sh) {

  dat = readxl::read_excel(path, sheet = sh) %>%
    dplyr::mutate(treat.bin = as.numeric(as.numeric(factor(treat))>1))

  fit    = lm(y ~ treat.bin + baseline, data = dat)
  adj.md = coef(fit)[["treat.bin"]]
  se.md  = summary(fit)$coefficients["treat.bin", "Std. Error"]
  cc     = dat[complete.cases(dat[, c("y", "treat.bin", "baseline")]), ]
  sd.p   = pooled.sd(cc$y, cc$treat.bin)
  d      = adj.md / sd.p
  var.d  = (se.md / sd.p)^2

  data.frame(study = sh, n = nrow(cc), d = d, var.d = var.d, se.d = sqrt(var.d))
})

print(es.list)


# 3. Stage 2: Bayesian random-effects meta-analysis in JAGS -----------------
#
# Model:
#   d[j]       ~ N(theta[j],  var.d[j])    — known sampling variance
#   theta[j]   ~ N(mu, tau^2)              — random effects
#   mu         ~ N(0, 1e-6)               — flat prior on overall mean
#   tau        ~ half-normal(0, 0.5)       — weakly informative heterogeneity prior
#
# Note: JAGS dnorm is parameterised by precision (= 1/variance).
#       The per-study sampling variances are treated as fixed and known,
#       which is the standard two-stage Bayesian meta-analysis assumption.

k = nrow(es.list)
data.jags = list(
  k     = k,
  d     = es.list$d,
  prec  = 1 / es.list$var.d   # known per-study precisions
)

M = "
model {

  for (j in 1:k) {
    d[j]     ~ dnorm(theta[j], prec[j])       # likelihood: observed Cohen's d
    theta[j] ~ dnorm(mu, inv.tau.sq)          # random effect for study j
  }

  mu         ~ dnorm(0, 1.0E-6)               # flat prior on pooled effect
  tau        ~ dnorm(0, pow(0.5, -2)) T(0,)   # half-normal(0, 0.5) on heterogeneity
  inv.tau.sq <- pow(tau, -2)

  # Predictive distribution for a new (unobserved) study
  theta.new ~ dnorm(mu, inv.tau.sq)
}

#monitor# mu, theta, tau, theta.new
"

fit.jags = run.jags(
  model    = M,
  data     = data.jags,
  n.chains = 4,
  burnin   = 5000,
  sample   = 80000,
  thin     = 2,
  summarise = FALSE
)

# Convergence check
summary(fit.jags)
mcmc = coda::as.mcmc(fit.jags)

cat("\n── Gelman-Rubin R-hat ─────────────────────────────────────────────────\n")
print(coda::gelman.diag(fit.jags, multivariate = FALSE))
cat("\n── Effective Sample Size ──────────────────────────────────────────────\n")
print(round(coda::effectiveSize(fit.jags)))
plot(fit.jags)

# 4. Posterior summaries -----------------------------------------------------

# Study-level posteriors: theta[1] ... theta[k]
theta.cols = paste0("theta[", 1:k, "]")

study.post = purrr::map_dfr(seq_len(k), function(j) {
  x = mcmc[, theta.cols[j]]
  data.frame(
    study  = es.list$study[j],
    md     = median(x),
    lo     = HDInterval::hdi(x)[1],
    hi     = HDInterval::hdi(x)[2],
    sd.post = sd(x)
  )
})

# Pooled effect (mu)
mu.samp = mcmc[, "mu"]
overall.post = data.frame(
  study   = "Overall",
  md      = median(mu.samp),
  lo      = HDInterval::hdi(mu.samp)[1],
  hi      = HDInterval::hdi(mu.samp)[2],
  sd.post = sd(mu.samp)
)

# Heterogeneity (tau)
tau.samp = mcmc[, "tau"]
tau.md   = median(tau.samp)
tau.hdi  = HDInterval::hdi(tau.samp)

# Prediction interval from posterior predictive theta.new
theta.new.samp = mcmc[, "theta.new"]
pi.lo = quantile(theta.new.samp, 0.025)
pi.hi = quantile(theta.new.samp, 0.975)

cat("\n── Pooled effect (mu) ─────────────────────────────────────────────────\n")
cat(sprintf("  Median      : %+.3f\n",   overall.post$md))
cat(sprintf("  95%% HDI    : [%.3f, %.3f]\n", overall.post$lo, overall.post$hi))
cat(sprintf("  P(mu < 0)   : %.3f\n",   mean(mu.samp < 0)))
cat("\n── Heterogeneity (tau) ────────────────────────────────────────────────\n")
cat(sprintf("  Median      : %.3f\n",    tau.md))
cat(sprintf("  95%% HDI    : [%.3f, %.3f]\n", tau.hdi[1], tau.hdi[2]))
cat("\n── Predictive interval (theta.new) ────────────────────────────────────\n")
cat(sprintf("  95%% PI     : [%.3f, %.3f]\n", pi.lo, pi.hi))


# 5. Forest plot -------------------------------------------------------------

# Study posteriors: theta[1] … theta[k]
study.long = purrr::map_dfr(seq_len(k), function(j) {
  data.frame(
    label  = es.list$study[j],
    draw   = as.numeric(mcmc[, paste0("theta[", j, "]")]),
    type   = "study"
  )
})

# Pooled effect
mu.long = data.frame(
  label = "Overall (μ)",
  draw  = as.numeric(mcmc[, "mu"]),
  type  = "overall")

# Ordered: studies sorted by posterior median, Overall always on top
study.order = study.long %>%
  dplyr::group_by(label) %>%
  dplyr::summarise(med = median(draw), .groups = "drop") %>%
  dplyr::arrange(med) %>%
  dplyr::pull(label)

all.long = dplyr::bind_rows(study.long, mu.long) %>%
  dplyr::mutate(
    label = dplyr::case_when(
      label == "Overall (μ)" ~ label,
      TRUE ~ label %>%
        stringr::str_replace("([a-zA-Z])([0-9])", "\\1 \\2") %>%
        stringr::str_to_title()
    ),
    label = factor(label,
                   levels = c("Overall (μ)",
                              stringr::str_replace(study.order,
                                "([a-zA-Z])([0-9])", "\\1 \\2") %>%
                                stringr::str_to_title())))
anno = all.long %>%
  dplyr::group_by(label, type) %>%
  dplyr::summarise(
    med  = median(draw),
    lo   = HDInterval::hdi(draw)[1],
    hi   = HDInterval::hdi(draw)[2],
    prob.neg = mean(draw < 0),
    .groups = "drop") %>%
  dplyr::mutate(
    anno.txt = sprintf("%.2f [%.2f, %.2f]", med, lo, hi))

# x-axis limits: cover all 95% HDIs with a bit of padding
x.lo = min(anno$lo) - 0.15
x.hi = max(anno$hi) + 0.15

# Create the plot
col.study   = "#4A90D9"   # steel blue for study densities
col.overall = "#E05A2B"   # burnt orange for pooled effect
col.zero    = "#555555"   # reference line
p = ggplot(all.long, aes(x = draw, y = label, fill = type)) +

  # Reference line at zero
  geom_vline(xintercept = 0, linetype = "dashed",
             colour = col.zero, linewidth = 0.45, alpha = 0.7) +
  # Density ridges
  ggridges::geom_density_ridges(
    aes(height = after_stat(density)),
    stat          = "density",
    trim          = TRUE,
    scale         = 0.88,        # overlap control: <1 = no overlap
    rel_min_height = 0.005,
    alpha         = 0.55,
    colour        = "white",
    linewidth     = 0.35,
    bandwidth     = "SJ"         # Sheather-Jones bandwidth
  ) +
  # Median tick
  geom_segment(
    data = anno,
    aes(x = med, xend = med,
        y    = as.numeric(label),
        yend = as.numeric(label) + 0.72),
    colour    = "white",
    linewidth = 0.9,
    inherit.aes = FALSE
  ) +
  # 95 % HDI bar below each density
  geom_segment(
    data = anno,
    aes(x = lo, xend = hi,
        y    = as.numeric(label),
        yend = as.numeric(label),
        colour = type),
    linewidth   = 1.4,
    lineend     = "round",
    inherit.aes = FALSE
  ) +
  # Annotation: "median [lo, hi]" to the right
  geom_text(
    data = anno,
    aes(x = x.hi + 0.01, y = as.numeric(label) + 0.35,
        label = anno.txt, colour = type),
    hjust       = 0,
    size        = 2.9,
    inherit.aes = FALSE
  ) +
  # Scales & colours
  scale_fill_manual(
    values = c(study = col.study, overall = col.overall),
    guide  = "none"
  ) +
  scale_colour_manual(
    values = c(study = col.study, overall = col.overall),
    guide  = "none"
  ) +
  scale_x_continuous(
    name   = "Cohen's *d* (posterior)",
    limits = c(x.lo, x.hi + 0.45),   # right padding for annotations
    expand = c(0, 0)
  ) + scale_y_discrete(name = NULL) +
  # Subtitle carries tau summary
  labs(
    title    = "Bayesian Random-Effects Meta-Analysis",
    subtitle = sprintf(
      "Posterior densities (θⱼ) and pooled effect (μ) · τ median = %.2f  95%% HDI [%.2f, %.2f]",
      tau.md, tau.hdi[1], tau.hdi[2]),
    caption  = "Vertical tick = posterior median · bar = 95% HDI · shaded = full posterior density"
  ) +
  theme_minimal(base_size = 11) +
  theme(
    base_family        = "Arial",
    plot.title         = element_text(face = "bold", size = 13),
    plot.subtitle      = element_text(size = 8.5, colour = "grey40"),
    plot.caption       = element_text(size = 7.5,  colour = "grey50", hjust = 0),
    axis.text.y        = element_text(size = 9.5),
    axis.title.x       = ggtext::element_markdown(),  # renders italic in title
    panel.grid.major.y = element_blank(),
    panel.grid.minor   = element_blank(),
    panel.grid.major.x = element_line(colour = "grey90", linewidth = 0.3),
    plot.margin        = margin(12, 60, 10, 8)   # extra right margin for labels
  )



# 6. Posterior of mu and tau -------------------------------------------------

# mu
p.mu = ggplot(data.frame(mu = as.numeric(mu.samp)), aes(x = mu)) +
  geom_density(fill = "#6b58a6", alpha = 0.3, color = "#6b58a6", linewidth = 0.7) +
  geom_vline(xintercept = median(mu.samp), color = "#6b58a6", linewidth = 0.7) +
  geom_vline(xintercept = 0, linetype = "dotted", color = "black") +
  annotate("text", x = median(mu.samp), y = Inf,
           label = sprintf("median = %.2f\n95%% HDI [%.2f, %.2f]",
                           median(mu.samp), HDInterval::hdi(mu.samp)[1],
                           HDInterval::hdi(mu.samp)[2]),
           vjust = 1.5, hjust = -0.05, size = 2.8, color = "#6b58a6") +
  labs(x = expression(mu ~ "(pooled Cohen's " * italic(d) * ")"),
       y = "Posterior density",
       title = expression("Posterior of " ~ mu)) +
  theme_classic(base_size = 9) +
  theme(plot.title = element_text(face = "bold"))

# tau
p.tau = ggplot(data.frame(tau = as.numeric(tau.samp)), aes(x = tau)) +
  geom_density(fill = "#e07b39", alpha = 0.3, color = "#e07b39", linewidth = 0.7) +
  geom_vline(xintercept = tau.md, color = "#e07b39", linewidth = 0.7) +
  annotate("text", x = tau.md, y = Inf,
           label = sprintf("median = %.2f\n95%% HDI [%.2f, %.2f]",
                           tau.md, tau.hdi[1], tau.hdi[2]),
           vjust = 1.5, hjust = -0.05, size = 2.8, color = "#e07b39") +
  labs(x = expression(tau ~ "(heterogeneity SD)"),
       y = "Posterior density",
       title = expression("Posterior of " ~ tau)) +
  theme_classic(base_size = 9) +
  theme(plot.title = element_text(face = "bold"))

p.mu
p.tau


# 7. Posterior tail probabilities -------------------------------------------

cat("\n── Posterior tail probabilities (mu) ──────────────────────────────────\n")
cat(sprintf("  P(mu < 0)      : %.3f   [any benefit]\n",       mean(mu.samp < 0)))
cat(sprintf("  P(mu < -0.2)   : %.3f   [small effect, d > 0.2]\n", mean(mu.samp < -0.2)))
cat(sprintf("  P(mu < -0.5)   : %.3f   [medium effect, d > 0.5]\n", mean(mu.samp < -0.5)))
cat(sprintf("  P(mu > 0)      : %.3f   [harm]\n\n",             mean(mu.samp > 0)))
