# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                                                             #
#   05-ipd-meta-analysis-binary.R                                             #
#   IPD-AD Model – Binary Outcome (y dichotomised, logistic likelihood)       #
#   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 > median(y) → 1, otherwise 0  (median cut)                     #
#   Pooling: posterior pooling across all imputation sets                     #
#                                                                             #
#   AD data: 5 synthetic aggregate-data studies are simulated (Section B).    #
#            AD studies report log-OR + SE. No sd.p scaling is needed since   #
#            the log-OR is already on a natural standardised scale.           #
#                                                                             #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #


# 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 -----------------------------------------------------

clusterScale = function(x, cluster) {
  mu.k = ave(x, cluster, FUN = mean)
  (x - mu.k) / sd(x)
}


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

load("_data/_imputation/implist.rda")
implist %>% map(function(x){
  split(x, ~id.study) %>% 
    map_dfr(function(y){
      y$y.bin = ifelse(y$y > median(y$y), 1, 0); y
    })
}) -> implist
k.ipd = nlevels(implist[[1]]$study %>% as.factor())
study.names = levels(implist[[1]]$study %>% as.factor())


# 3. Simulate aggregate-data (AD) studies ----------------------------------
# AD studies for a binary outcome supply a log-OR and its SE.
# SE approximation for log-OR from a 2×2 table with n_arm per arm
# and event probability p ≈ 0.40 in the control group, OR ≈ exp(log_or_agg):
#   SE ≈ sqrt(1/(n_arm*p*(1-p)) + 1/(n_arm*p_trt*(1-p_trt)))
# Simplified: SE ≈ sqrt(4/n_arm) for balanced ~50% event rates (conservative).

set.seed(2024)
k.agg = 5
ad.names = c("hart2009", "kellner2012", "novak2015", "reeves2017", "vogel2021")

n.arm.ad = sample(40:120, k.agg, replace = TRUE)
log.or.agg = rnorm(k.agg, mean = -0.50, sd = 0.35)
se.agg = sqrt(4 / n.arm.ad)

message("\nSimulated AD studies (log-OR scale):")
print(data.frame(study = ad.names, n.arm = n.arm.ad, log.or = round(log.or.agg, 3),
  OR = round(exp(log.or.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]]]
  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.bin, treat = dat$treat.bin,
    k.agg = k.agg, es = log.or.agg, se = se.agg)

  M = "
  model {

    # Part 1: Likelihood for IPD studies (Sutton et al., 2008, Stat Med)
    for (i in 1:n.ipd) {
      y[i]         ~ dbern(p[i])
      logit(p[i]) <- alpha[study[i]]
                    + delta[study[i]] * treat[i]
                    + beta0[study[i]] * cov[i]
    }

    # Weakly informative priors on logistic intercepts/slopes
    for (i in 1:k.ipd) {
      alpha[i] ~ dnorm(0, 1.0E-3)
      beta0[i] ~ dnorm(0, 1.0E-3)
    }

    # Part 2: Likelihood for AD studies — each es[i] is noisy draw from delta[k.ipd + i]
    for (i in 1:k.agg) {
      es[i]             ~ dnorm(delta[k.ipd + i], prec.sigma.agg[i])
      prec.sigma.agg[i] <- inverse(pow(se[i], 2))
    }

    # Part 3: Random-effects pooling — delta[j] is log-OR; all studies share mu, tau
    for (j in 1:(k.ipd + k.agg)) {
      delta[j] ~ dnorm(mu, inv.tau.sq)
    }

    mu         ~ dnorm(0.0, 1.0E-3)
    tau        ~ dnorm(0, pow(0.5, -2))T(0,)
    inv.tau.sq <- inverse(pow(tau, 2))
  }

 #monitor# delta, mu, tau
 "

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


# 5. Summary of the results -------------------------------------------------

mcmc = fits %>%
  purrr::map(~coda::as.mcmc(.)) %>%
  do.call(rbind, .)

mu.samp = mcmc[, "mu"]
cat(sprintf("\n── Pooled estimate (logistic IPD-AD) ──\n"))
cat(sprintf("  mu (log-OR) median : %+.3f\n",  median(mu.samp)))
cat(sprintf("  OR median          : %.3f\n",   exp(median(mu.samp))))
cat(sprintf("  95%% HDI (log-OR)  : [%.3f, %.3f]\n",
            hdi(mu.samp)[1], hdi(mu.samp)[2]))
cat(sprintf("  95%% HDI (OR)      : [%.3f, %.3f]\n",
            exp(hdi(mu.samp)[1]), exp(hdi(mu.samp)[2])))
cat(sprintf("  P(mu < 0)          : %.3f\n\n", mean(mu.samp < 0)))


# 6. Moderator analysis – baseline as treatment moderator -------------------
# gamma[j]: study-specific treat x baseline interaction on the log-OR scale.
# Estimable only from IPD; AD studies cannot contribute participant-level
# covariate data but continue to anchor mu and tau.

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

for (i in seq_len(nsets)) {

  dat = implist[[sets[i]]]
  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.bin, treat = dat$treat.bin,
    k.agg = k.agg, es = log.or.agg, se = se.agg
  )

  M.mod <- "
    model {

    ## ── Part 1: IPD likelihood with treat x baseline interaction ───────────
    for (i in 1:n.ipd) {
      y[i]         ~ dbern(p[i])
      logit(p[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-3)
      beta0[i] ~ dnorm(0, 1.0E-3)
    }

    ## ── Part 2: AD likelihood ───────────────────────────────────────────────
    ## You can hash this part out if there is no AD
     for (i in 1:k.agg) {
       es[i]             ~ dnorm(delta[k.ipd + i], prec.sigma.agg[i])
       prec.sigma.agg[i] <- inverse(pow(se[i], 2))
    }

    ## ── Part 3: Random-effects pooling ─────────────────────────────────────
    ## Hash this part out if you do not have any AD
     for (j in 1:(k.ipd + k.agg)) {
       delta[j] ~ dnorm(mu, inv.tau.sq)
    }
    ## Moderation: IPD studies only
     for (j in 1:k.ipd) {
       gamma[j] ~ dnorm(mu.gamma, inv.tau.gamma.sq)
    }

      mu         ~ dnorm(0.0, 1.0E-3)
      tau        ~ dt(0, pow(0.5, -2), 1) T(0,)
      inv.tau.sq <- inverse(pow(tau, 2))

      mu.gamma         ~ dnorm(0.0, 1.0E-3)
      tau.gamma        ~ dt(0, pow(0.5, -2), 1) T(0,)
      inv.tau.gamma.sq <- inverse(pow(tau.gamma, 2))
    }
    #monitor# delta, mu, tau, gamma, mu.gamma, tau.gamma
 "

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

mcmc.mod = fits.mod %>%
  purrr::map(~coda::as.mcmc(.)) %>%
  do.call(rbind, .)

mu.gamma.samp = mcmc.mod[, "mu.gamma"]

cat("\n── Moderator: Baseline x Treatment Interaction (logistic IPD-AD) ──\n")
cat(sprintf("  mu.gamma  median  : %+.3f  (log-OR scale)\n",  median(mu.gamma.samp)))
cat(sprintf("  exp(mu.gamma)     : %.3f   (OR per SD of within-cluster baseline)\n",
            exp(median(mu.gamma.samp))))
cat(sprintf("  95%% HDI (log-OR) : [%.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; AD studies inform mu via shared tau.\n\n")