# Bayesian model averaging and Figures 4--5 -------------------------------
#
# Standard BMS output lacks the cross-coefficient posterior covariances needed
# for interaction marginal effects.  The vendored additive extension records
# posterior second moments during MCMC and is therefore the primary estimator.

suppressPackageStartupMessages({library(BMS); library(dplyr); library(ggplot2)})
source(file.path(ROOT, "R", "bms_mcmc_extensions.R"), local = TRUE)

baseline <- readRDS(file.path(OUTPUTS, "baseline_ols.rds"))
df <- baseline$data
bma_data <- df %>% select(size, all_of(controls), rule_of_law, trust_index, trust_rule_of_law)

# Full manuscript settings are defaults.  REPLICATION_QUICK_BMA=1 is an
# explicitly non-publication smoke-test switch used only for development.
quick <- identical(Sys.getenv("REPLICATION_QUICK_BMA"), "1")
burn <- if (quick) 100L else 100000L
iterations <- if (quick) 500L else 200000L
stored_models <- if (quick) 100L else 5000L
set.seed(20260607)
bma_model <- bms_mcmc_cov(
  X.data = as.data.frame(bma_data), burn = burn, iter = iterations,
  g = "BRIC", mprior = "random", nmodel = stored_models, mcmc = "bd",
  user.int = FALSE, fixed.reg = institutional_terms
)

coef_raw <- as.data.frame(coef(bma_model, exact = FALSE, order.by.pip = FALSE,
                               include.constant = TRUE, std.coefs = FALSE))
coef_raw$term <- rownames(coef_raw); rownames(coef_raw) <- NULL
names(coef_raw) <- make.names(names(coef_raw))
moments <- bms_mcmc_posterior_moments(bma_model, terms = institutional_terms)
validation <- bms_mcmc_validate_against_bms(bma_model, terms = institutional_terms)
write.csv(validation, file.path(OUTPUTS, "bma_native_validation.csv"), row.names = FALSE)
saveRDS(bma_model, file.path(OUTPUTS, "bma_model.rds"))

z90 <- qnorm(.95)
table4 <- data.frame(
  term = institutional_terms,
  posterior_mean = unname(moments$posterior_mean[institutional_terms]),
  posterior_sd = sqrt(diag(moments$posterior_cov[institutional_terms, institutional_terms]))
) %>% mutate(low_90 = posterior_mean - z90 * posterior_sd,
             high_90 = posterior_mean + z90 * posterior_sd)
write.csv(table4, file.path(TABLES, "Table_4.csv"), row.names = FALSE)

# Posterior marginal effects using the full retained covariance matrix.
me_bma <- function(base, interaction, values) {
  b <- moments$posterior_mean; V <- moments$posterior_cov
  variance <- V[base, base] + values^2 * V[interaction, interaction] + 2 * values * V[base, interaction]
  data.frame(value = values, effect = b[[base]] + values * b[[interaction]],
             low = b[[base]] + values * b[[interaction]] - z90 * sqrt(variance),
             high = b[[base]] + values * b[[interaction]] + z90 * sqrt(variance))
}
g4a <- me_bma("trust_index", "trust_rule_of_law",
              seq(min(df$rule_of_law), max(df$rule_of_law), length.out = 200))
g4b <- me_bma("rule_of_law", "trust_rule_of_law",
              seq(min(df$trust_index), max(df$trust_index), length.out = 200))
write.csv(g4a, file.path(OUTPUTS, "Figure_4a_data.csv"), row.names = FALSE)
write.csv(g4b, file.path(OUTPUTS, "Figure_4b_data.csv"), row.names = FALSE)
posterior_plot <- function(grid, support, x_label, y_label, fill_color, line_color) ggplot(grid, aes(value, effect)) +
  geom_hline(yintercept = 0, color = "grey35", linewidth = 0.4, linetype = "dashed") +
  geom_ribbon(aes(ymin = low, ymax = high), fill = fill_color, alpha = 0.18) +
  geom_line(color = line_color, linewidth = 0.9) +
  geom_point(data = data.frame(value = support, y = 0), aes(value, y), inherit.aes = FALSE,
             color = "black", alpha = 0.35, size = 1) +
  labs(x = x_label, y = y_label) + theme_minimal(base_size = 12)
p4a <- posterior_plot(g4a, df$rule_of_law, "Rule of law",
                      "Marginal effect of generalized trust on the reported size slope", "#3B7EA1", "#1B5E7A")
p4b <- posterior_plot(g4b, df$trust_index, "Generalized trust",
                      "Marginal effect of rule of law on the reported size slope", "#A65E2E", "#7A4018")

# Figure 5 uses the native BMS image method.
mean_col <- grep("Post.Mean", names(coef_raw), value = TRUE)[1]
sd_col <- grep("Post.SD", names(coef_raw), value = TRUE)[1]
pip_col <- grep("PIP", names(coef_raw), value = TRUE)[1]
coefficient_labels <- c(
  se="Standard error", rel_size="Relative-size coding", avg_year="Sample midyear",
  num_years="Log sample length", any_trim="Any trimming", OLS="OLS estimator",
  FM="Fama-MacBeth estimator", excess="Excess-return specification",
  indstock="Individual-stock sample", jan="January-return specification",
  impact="Journal impact", cits="Log citations", indep_demean="Demeaned size variable",
  value="Value controls", syst_risk="Systematic-risk controls", momentum="Momentum controls",
  leverage_group="Leverage controls", profitability="Profitability controls",
  liquidity="Liquidity controls", volatility="Volatility controls", price_group="Price controls",
  mkt_char="Market-characteristic controls", mkt_state="Market-state controls",
  information="Information controls", ownership="Ownership controls", growth="Growth controls",
  distress_group="Distress controls", size_group="Additional size controls", bond_yield="Bond yield",
  mkt_return_vol="Market-return volatility", credit_gdp="Private credit to GDP",
  gdp_pc="Log GDP per capita", gdp_growth="GDP growth", rule_of_law="Rule of law",
  trust_index="Generalized trust", trust_rule_of_law="Generalized trust x Rule of law"
)
plot_bma_model <- bma_model
plot_bma_model$reg.names <- unname(coefficient_labels[plot_bma_model$reg.names])
stopifnot(!anyNA(plot_bma_model$reg.names))

save_plot(p4a, "Figure_4a", 4.0, 5.2); save_plot(p4b, "Figure_4b", 4.0, 5.2)
png(file.path(FIG_PNG, "Figure_5.png"), width = 2400, height = 1600, res = 220)
image(plot_bma_model, cex.axis = 0.7, order.by.pip = TRUE, yprop2pip = FALSE)
dev.off()
pdf(file.path(FIG_PDF, "Figure_5.pdf"), width = 11, height = 8)
image(plot_bma_model, cex.axis = 0.7, order.by.pip = TRUE, yprop2pip = FALSE)
dev.off()

# Table 11 combines the full OLS results with BMA summaries.
ols_study <- baseline$study[, c("term", "estimate", "std_error", "p_value")]
names(ols_study)[-1] <- paste0(names(ols_study)[-1], "_study")
ols_country <- baseline$country[, c("term", "std_error", "p_value")]
names(ols_country)[-1] <- paste0(names(ols_country)[-1], "_country")
bma_export <- coef_raw %>% transmute(term, posterior_mean = .data[[mean_col]],
                                     posterior_sd = .data[[sd_col]], pip = .data[[pip_col]])
table11 <- Reduce(function(x, y) merge(x, y, by = "term", all = TRUE),
                  list(ols_study, ols_country, bma_export))
write.csv(table11, file.path(TABLES, "Table_11.csv"), row.names = FALSE)
