# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#                                                                             #
#   01-linear-regression.R                                                    #
#                                                                             #
#   Simple Bayesian regression: y ~ treatment                                 #
#   First study from example-data.xlsx                                        #
#   Engine: JAGS via runjags                                                  #
#   Model: y[i] ~ N(beta0 + beta1*treat[i], sigma²); vague priors.            #
#                                                                             #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #


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

library(readxl)
library(runjags)
library(coda)
library(ggplot2)
library(dplyr)
library(tidyr)
library(broom)
library(HDInterval)


# 1. Data ------------------------------------------------------------------

# Get first dataset (one study)
path = "_data/data.xlsx"
dat = readxl::read_excel(path, sheet = 1) %>%
  dplyr::mutate(treat.bin = as.integer(treat > 1)) %>%
  dplyr::filter(!is.na(y), !is.na(treat.bin))

# Give details about the study
cat(sprintf("Study: %s  |  n = %d  |  n.control = %d  |  n.treatment = %d\n",
  dat$study[1], nrow(dat),
  sum(dat$treat.bin == 0), sum(dat$treat.bin == 1)))

data.jags = list(n = nrow(dat), y = dat$y, treat = dat$treat.bin)


# 2. JAGS model --------------------------------------------------------------
# y[i] ~ N(beta0 + beta1 * treat[i],  sigma²)
# Priors:
#   beta0, beta1 ~ N(0, 0.001)  (flat / uninformative)
#   sigma        ~ Gamma(0.001,0.001)
# beta1 = adjusted mean difference (active – control)

M = "
model {

  for (i in 1:n) {
    y[i]    ~ dnorm(mu[i], tau)
    mu[i]  <- beta0 + beta1 * treat[i]
  }
  beta0 ~ dnorm(0, 0.001)      # intercept (control group mean)
  beta1 ~ dnorm(0, 0.001)      # treatment effect (active - control)
  tau  ~ dgamma(0.001,0.001)   # 1/residual SD
  sigma <- 1 / sqrt(tau)
}
#monitor# beta0, beta1, sigma
"


# 3. MCMC sampling ----------------------------------------------------------

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

# Create a summary
summary(fit)

# Compare with a frequentist model
lm(y ~ treat.bin, data = dat) %>% 
  broom::tidy(conf.int=TRUE)


# 4. Model diagnostics ------------------------------------------------------

cat("\n── Gelman-Rubin R-hat ─────────────────────────────────────────────────\n")
print(coda::gelman.diag(fit))

cat("\n── Effective Sample Size ──────────────────────────────────────────────\n")
print(coda::effectiveSize(fit))

# Trace + density plots 
plot(fit)


# 5. Posterior summary ------------------------------------------------------

# Pool all chains into a single matrix
mcmc = coda::as.mcmc(fit)

posterior.summary = function(x, name) {
  data.frame(parameter = name, mean = mean(x), median = median(x), sd = sd(x),
    hdi.lo = HDInterval::hdi(x)[1], hdi.hi = HDInterval::hdi(x)[2])
}

smry = dplyr::bind_rows(
  posterior.summary(mcmc[, "beta0"], "beta0 (control mean)"),
  posterior.summary(mcmc[, "beta1"], "beta1 (treatment effect)"),
  posterior.summary(mcmc[, "sigma"], "sigma (residual SD)")
)

cat("\n── Posterior Summary ──────────────────────────────────────────────────\n")
print(smry, digits = 3, row.names = FALSE)


# 6. Posterior of beta1 -----------------------------------------------------

beta1 = mcmc[, "beta1"]

# Compute quantities for annotation
beta1.med = median(beta1)
beta1.hdi = HDInterval::hdi(beta1)
beta1.mean = mean(beta1)

p.beta1 = ggplot(data.frame(beta1 = beta1), aes(x = beta1)) +
  stat_function(fun = approxfun(density(beta1), rule = 2), geom = "area",
    xlim = c(beta1.hdi[1], beta1.hdi[2]), fill = "forestgreen", alpha = 0.25) +
  geom_density(color = "forestgreen", linewidth = 0.8, fill = NA) +
  geom_vline(xintercept = beta1.med, color = "forestgreen", linetype = "solid", linewidth = 0.7) +
  geom_vline(xintercept = 0, color = "black", linetype = "dotted", linewidth = 0.5) +
  geom_vline(xintercept = beta1.hdi[1], color = "forestgreen", linetype = "dashed", linewidth = 0.4) +
  geom_vline(xintercept = beta1.hdi[2], color = "forestgreen", linetype = "dashed", linewidth = 0.4) +
  annotate("text", x = beta1.med, y = Inf, label = sprintf("median = %.2f", beta1.med),
    vjust = 2, hjust = -0.1, size = 3, color = "forestgreen") +
  annotate("text", x = mean(beta1.hdi), y = 0,
    label = sprintf("95%% HDI [%.2f, %.2f]", beta1.hdi[1], beta1.hdi[2]),
    vjust = -0.6, size = 3, color = "forestgreen") +
  labs(x = expression(beta[1] ~ "(treatment effect: active – control)"),
    y = "Posterior density", title = expression("Posterior distribution of " ~ beta[1])) +
  theme_classic(base_size = 10) +
  theme(plot.title = element_text(face = "bold"))
p.beta1


# 7. Posterior tail probabilities ------------------------------------------
# These answer questions of the form: "what is the probability that the
# treatment effect exceeds / falls below some threshold?"
# Here we use three illustrative thresholds:
#   0   = any effect at all (direction probability)
#   3   = a minimally important difference (example; adjust to your outcome)
#  -3   = harm threshold (symmetric example)

thresh.mid = 3     # example MID; replace with a clinically meaningful value

cat("\n── Posterior Tail Probabilities for beta1 ─────────────────────────────\n")
cat(sprintf("  P(beta1 > 0)           = %.3f   [treatment better than control]\n", mean(beta1 > 0)))
cat(sprintf("  P(beta1 < 0)           = %.3f   [treatment worse than control]\n", mean(beta1 < 0)))
cat(sprintf("  P(beta1 > %g)          = %.3f   [exceeds example MID of %g]\n",
  thresh.mid, mean(beta1 > thresh.mid), thresh.mid))
cat(sprintf("  P(beta1 < -%g)         = %.3f   [harm: worse by more than %g]\n",
  thresh.mid, mean(beta1 < -thresh.mid), thresh.mid))
cat(sprintf("  P(|beta1| < %g)        = %.3f   [practical equivalence zone]\n",
  thresh.mid, mean(abs(beta1) < thresh.mid)))

region.df = data.frame(beta1 = beta1)
p.tails = ggplot(region.df, aes(x = beta1)) +
  stat_function(fun = approxfun(density(beta1), rule = 2), geom = "area",
    xlim = c(thresh.mid, max(beta1)), fill = "orange", alpha = 0.4) +
  stat_function(fun = approxfun(density(beta1), rule = 2), geom = "area",
    xlim = c(min(beta1), -thresh.mid), fill = "darkblue", alpha = 0.4) +
  stat_function(fun = approxfun(density(beta1), rule = 2), geom = "area",
    xlim = c(-thresh.mid, thresh.mid), fill = "gray70", alpha = 0.25) +
  geom_density(color = "darkgrey", linewidth = 0.8, fill = NA) +
  geom_vline(xintercept = 0, linetype = "dotted", color = "black", linewidth = 0.5) +
  geom_vline(xintercept = thresh.mid, linetype = "dashed", color = "orange", linewidth = 0.5) +
  geom_vline(xintercept = -thresh.mid, linetype = "dashed", color = "darkblue", linewidth = 0.5) +
  annotate("text", x = thresh.mid + 0.3, y = Inf,
    label = sprintf("P(>±%g)\n%.3f", thresh.mid, mean(beta1 > thresh.mid)),
    vjust = 1.5, hjust = 0, size = 2.8, color = "orange") +
  annotate("text", x = -thresh.mid - 0.3, y = Inf,
    label = sprintf("P(<−%g)\n%.3f", thresh.mid, mean(beta1 < -thresh.mid)),
    vjust = 1.5, hjust = 1, size = 2.8, color = "darkblue") +
  annotate("text", x = 0, y = Inf, label = sprintf("ROPE\n%.3f", mean(abs(beta1) < thresh.mid)),
    vjust = 1.5, hjust = 0.5, size = 2.8, color = "gray40") +
  labs(x = expression(beta[1] ~ "(treatment effect)"), y = "Posterior density",
    title = "Posterior tail probabilities",
    subtitle = sprintf("Threshold = ±%.0f  (replace with clinically meaningful MID)", thresh.mid)) +
  theme_classic(base_size = 10) +
  theme(plot.title = element_text(face = "bold"),
    plot.subtitle = element_text(size = 8, color = "gray40"))
p.tails