# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                                                             #
#   04-ipd-meta-analysis.R                                                    #
#   IPD-AD Model – Continuous Outcome (y)                                     #
#   Adapted from 03-analysis.R (Sutton et al., 2008 framework)                #
#                                                                             #
#   Data:    implist.rda  (multiply imputed data)                             #
#   Treat:   dichotomised as treat == 1 → 0 (control), treat > 1 → 1 (active) #
#   Outcome: y (continuous, used as-is)                                       #
#   Pooling: posterior pooling across all imputation sets                     #
#                                                                             #
#   AD data: 5 synthetic aggregate-data studies are simulated (Section B).    #
#            AD studies report Cohen's d + SE directly; no sd.p scaling is    #
#            applied to their coefficients in Part 3 of the JAGS model.       #
#                                                                             #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #


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

library(readxl)
library(dplyr)
library(purrr)
library(tidyr)
library(rjags)
library(runjags)
library(ggplot2)
library(ggridges)
library(forcats)
library(HDInterval)
library(stringr)


# 1. Helper functions -----------------------------------------------------

# Pooled within-group SD (Cohen's d denominator)
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) * (length(g) - 1))
  sqrt(sum(ss) / sum(df))
}

# Within-cluster centering scaled by overall SD
clusterScale = function(x, cluster) {
  mu.k = ave(x, cluster, FUN = mean)
  (x - mu.k) / sd(x)
}


# 2. Build implist ---------------------------------------------------------

path = "_data/_imputation/implist.rda"
load(path)

# Get number of studies an study.names
k.ipd = nlevels(implist[[1]]$id.study %>% as.factor)
study.names = levels(implist[[1]]$study %>% as.factor)


# 3. Simulate aggregate-data (AD) studies ----------------------------------
# In a real analysis, AD studies provide a published Cohen's d and its SE.
# Here we simulate 5 plausible AD studies consistent with a true pooled d ≈ −0.30.
# SE approximation for Cohen's d with n_arm participants per arm (balanced):
#   SE ≈ sqrt(2/n_arm + d²/(4*n_arm))

set.seed(2026)
k.agg = 5
ad.names = c("hart2009", "kellner2012", "novak2015", "reeves2017", "vogel2021")
n.arm.ad = sample(40:120, k.agg, replace = TRUE)
es.agg = rnorm(k.agg, mean = -0.1, sd = 0.20)
se.agg = sqrt(2 / n.arm.ad + es.agg^2 / (4 * n.arm.ad))

message("\nSimulated AD studies (SMD):")
print(data.frame(study = ad.names, n.arm = n.arm.ad,
  d = round(es.agg, 3), se = round(se.agg, 3)))


# 4. Main IPD-AD analysis ---------------------------------------------------

nsets = length(implist)
sets = seq_len(nsets)
fits = vector("list", nsets)

for (i in seq_len(nsets)) {

  dat = implist[[sets[i]]]
  sd.p = split(dat, dat$id.study) %>%
    purrr::map_dbl(function(x) pooled.sd(x$y, x$treat.bin))

  data.jags = 
    list(
      n.ipd = nrow(dat), k.ipd = k.ipd, study = dat$id.study,
      cov = clusterScale(dat$baseline, dat$id.study), 
      y = dat$y, treat = dat$treat.bin,
      sd.p = sd.p, k.agg = k.agg, 
      es = es.agg, se = se.agg)

  M = "
    model {
  
   ## ── Part 1: Likelihood for IPD studies ─────────────────────────────────
   ## (Sutton et al., 2008, Stat Med, doi:10.1002/sim.2916)
    for (i in 1:n.ipd) {
      y[i]    ~ dnorm(yhat[i], prec.sigma.e.y[study[i]])
      yhat[i] <- alpha[study[i]]
                + delta[study[i]] * treat[i]
                + beta0[study[i]] * cov[i]
    }
  
    for (i in 1:k.ipd) {
      alpha[i]          ~ dnorm(0, 1.0E-6)
      beta0[i]          ~ dnorm(0, 1.0E-6)
      sigma.e.y[i]      ~ dunif(0, 100)
      prec.sigma.e.y[i] <- inverse(pow(sigma.e.y[i], 2))
    }
  
   ## ── Part 2: Likelihood for AD studies ──────────────────────────────────
   ## Each published Cohen's d (es[i]) is a noisy draw from the study's
   ## true effect d[k.ipd + i], with precision determined by its SE.
    for (i in 1:k.agg) {
      es[i]             ~ dnorm(d[k.ipd + i], prec.sigma.agg[i])
      prec.sigma.agg[i] <- inverse(pow(se[i], 2))
    }
  
   ## ── Part 3: Random-effects pooling across IPD + AD studies ────────────
   ## IPD: raw treatment coefficient → Cohen's d via per-study pooled SD
    for (j in 1:k.ipd) {
      delta[j] <- d[j] * sd.p[j]
      d[j]     ~ dnorm(mu, inv.tau.sq)
    }
   ## AD: effect is already on Cohen's d scale — no sd.p scaling needed
    for (j in (k.ipd + 1):(k.ipd + k.agg)) {
      d[j] ~ dnorm(mu, inv.tau.sq)
    }
  
    mu         ~ dnorm(0.0, 1.0E-6)
    tau        ~ dnorm(0, pow(0.5, -2)) T(0,)  
    inv.tau.sq <- inverse(pow(tau, 2))
    }
  
  #monitor# d, mu, tau
  "

  message("IPD-AD continuous model | imputation set ", i, " / ", nsets)
  fits[[i]] = run.jags(M, data = data.jags, summarise = FALSE, n.chains = 4,
    burnin = 1000, sample = 50000)
}


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

map(fits, function(x){
  if (!is.null(x)) as.mcmc(x)
}) %>% do.call(rbind, .) -> mcmc

k.all = k.ipd + k.agg
all.names = c(study.names, ad.names, "Overall")
col.names = all.names %>%
  stringr::str_replace("([a-zA-Z])([0-9])", "\\1, \\2") %>%
  stringr::str_to_title()

n.draw = min(2e6, nrow(mcmc))
plot.dat = mcmc[sample(nrow(mcmc), n.draw), 1:(k.all + 1)] %>%
  { colnames(.) = col.names; . } %>%
  as_tibble() %>%
  pivot_longer(everything()) %>%
  arrange(name) %>%
  group_by(name) %>%
  mutate(median = median(value)) %>%
  ungroup() %>%
  mutate(median = ifelse(name == tail(col.names, 1), 100, median),
    name = fct_reorder(name, median, .desc = TRUE))

ad.display = ad.names %>%
  stringr::str_replace("([a-zA-Z])([0-9])", "\\1, \\2") %>%
  stringr::str_to_title()
overall.display = tail(col.names, 1)

plot.dat2 = plot.dat %>%
  group_by(name) %>%
  dplyr::summarise(lo = hdi(value)[1], md = median(value), hi = hdi(value)[2], sd = sd(value)) %>%
  dplyr::mutate(source = dplyr::case_when(
    name == overall.display ~ "overall", name %in% ad.display ~ "ad", TRUE ~ "ipd"))

tau = mcmc[, "tau"] %>%
  { c(lo = sprintf("%.2f", hdi(.)[1]), md = sprintf("%.2f", median(.)), hi = sprintf("%.2f", hdi(.)[2])) }

pt.col = plot.dat2 %>%
  arrange(desc(md)) %>%
  dplyr::mutate(col = dplyr::case_when(
    source == "overall" ~ "purple", source == "ad" ~ "turquoise", TRUE ~ "black")) %>%
  pull(col)

p = ggplot(
  plot.dat %>%
    dplyr::left_join(plot.dat2 %>% dplyr::select(name, source), by = "name") %>%
    dplyr::slice_sample(n = min(1e6, nrow(.))) %>%
    dplyr::arrange(name),
  aes(x = value, y = name, fill = source, color = source)) +
  geom_linerange(data = plot.dat2 %>% arrange(desc(md)), color = pt.col,
    position = position_nudge(y = -0.15), linewidth = 0.2, aes(x = md, xmin = lo, xmax = hi)) +
  geom_point(data = plot.dat2 %>% arrange(desc(md)), color = pt.col, shape = 15,
    position = position_nudge(y = -0.15), aes(x = md, size = sd^-1)) +
  scale_size(range = c(0.0001, 1)) +
  geom_vline(xintercept = 0, linetype = "dotted") +
  geom_density_ridges(rel_min_height = 0.01, alpha = 0.8, size = 0.4) +
  theme_ridges() + ylab("") +
  xlab(bquote("Cohen's" ~ italic(d) ~ "  " ~ tau == .(tau["md"]) ~ " [" ~ .(tau["lo"]) ~ ";" ~ .(tau["hi"]) ~ "]")) +
  scale_fill_manual(values = c("overall" = "purple", "ipd" = "lightgray", "ad" = "lightblue"),
    labels = c("overall" = "Overall", "ipd" = "IPD study", "ad" = "AD study")) +
  scale_color_manual(values = c("overall" = "purple", "ipd" = "darkgray", "ad" = "turquoise"),
    labels = c("overall" = "Overall", "ipd" = "IPD study", "ad" = "AD study")) +
  scale_x_continuous(breaks = c(-1, -0.5, 0, 0.5), limits = c(-1.5, 1.2)) +
  geom_text(data = plot.dat2 %>% arrange(desc(md)), color = "black", size = 2.8, hjust = 1,
    position = position_nudge(y = 0.15), fontface = c("bold", rep("plain", k.all)),
    aes(x = 0.85, label = sprintf("%.2f", md))) +
  geom_text(data = plot.dat2 %>% arrange(desc(md)), color = "black", size = 2.8, hjust = 0,
    position = position_nudge(y = 0.15), fontface = c("bold", rep("plain", k.all)),
    aes(x = 0.87, label = paste0("[", sprintf("%.2f", lo), "; ", sprintf("%.2f", hi), "]"))) +
  coord_cartesian(clip = "off") +
  theme(legend.position = "right", legend.title = element_blank(),
    panel.grid.major.y = element_line("lightgray", 0.1), axis.title.x = element_text(size = 8),
    axis.text.x = element_text(size = 8), axis.ticks.y = element_line("lightgray", 0.1),
    axis.text.y = element_text(size = 8, face = c("bold", rep("plain", k.all))))


# 6. Moderator analysis – baseline as treatment moderator -------------------
# The treat × baseline interaction (gamma) is estimable only from IPD.

fits.mod = vector("list", nsets)

for (i in seq_len(nsets)) {

  dat = implist[[sets[i]]]
  sd.p = split(dat, dat$id.study) %>%
    purrr::map_dbl(function(x) pooled.sd(x$y, x$treat.bin))

  data.jags.mod = list(n.ipd = nrow(dat), k.ipd = k.ipd, study = dat$id.study,
    cov = clusterScale(dat$baseline, dat$id.study), y = dat$y, treat = dat$treat.bin,
    sd.p = sd.p)

  M.mod <- "
  model {

    ## ── Part 1: IPD likelihood with treat x baseline interaction ───────────
    for (i in 1:n.ipd) {
      y[i]    ~ dnorm(yhat[i], prec.sigma.e.y[study[i]])
      yhat[i] <- alpha[study[i]]
                + delta[study[i]] * treat[i]
                + beta0[study[i]] * cov[i]
                + gamma[study[i]] * treat[i] * cov[i]
    }
    for (i in 1:k.ipd) {
      alpha[i]          ~ dnorm(0, 1.0E-6)
      beta0[i]          ~ dnorm(0, 1.0E-6)
      delta[i]          ~ dnorm(0, 1.0E-6)
      sigma.e.y[i]      ~ dunif(0, 100)
      prec.sigma.e.y[i] <- inverse(pow(sigma.e.y[i], 2))
    }
  
    ## ── Part 2: Random-effects pooling ─────────────────────────────────────
    for (j in 1:k.ipd) {
      ## We pool the interaction coefficients across trials
      gamma[j] ~ dnorm(mu.gamma, inv.tau.gamma.sq)
    }
  
    mu               ~ dnorm(0.0, 1.0E-6)
    mu.gamma         ~ dnorm(0.0, 1.0E-6)
    tau.gamma        ~ dt(0, pow(0.3, -2),1) T(0,)   # Half-Cauchy prior
    inv.tau.gamma.sq <- inverse(pow(tau.gamma, 2))
  }

  #monitor# d, mu, tau, gamma, mu.gamma, tau.gamma
  "

  message("Moderator model | imputation set ", i, " / ", nsets)
  fits.mod[[i]] = run.jags(M.mod, data = data.jags.mod, summarise = FALSE, n.chains = 4,
    burnin = 1000, sample = 10000)
}

map(fits.mod, function(x){
  if (!is.null(x)) as.mcmc(x)
}) %>% do.call(rbind, .) -> mcmc.mod

mu.gamma.samp = mcmc.mod[, "mu.gamma"]
cat("\n── Moderator: Baseline x Treatment Interaction (continuous IPD) ──\n")
cat(sprintf("  mu.gamma  median : %+.3f\n",  median(mu.gamma.samp)))
cat(sprintf("  95%% HDI          : [%.3f, %.3f]\n",
            hdi(mu.gamma.samp)[1], hdi(mu.gamma.samp)[2]))
cat(sprintf("  P(mu.gamma < 0)  : %.3f\n",   mean(mu.gamma.samp < 0)))
cat("  Note: gamma estimated from IPD only\n\n")