# Shared data preparation and validation ------------------------------------
#
# Every analysis script sources this file.  It centralizes transformations so
# that no table or figure can silently use a different sample definition.

suppressPackageStartupMessages({
  library(dplyr)
  library(tidyr)
})

trust_components <- c(
  "z_mean_g007_18_b", "z_mean_g007_33_b", "z_mean_g007_34_b",
  "z_mean_g007_35_b", "z_mean_g007_36_b", "z_mean_g007_01", "z_share_6_a165"
)
placebo_components <- c("z_mean_d001", "z_mean_d001_b")
controls <- c(
  "se", "rel_size", "avg_year", "num_years", "any_trim", "OLS", "FM", "excess",
  "indstock", "jan", "impact", "cits", "indep_demean", "value", "syst_risk",
  "momentum", "leverage_group", "profitability", "liquidity", "volatility",
  "price_group", "mkt_char", "mkt_state", "information", "ownership", "growth",
  "distress_group", "size_group", "bond_yield", "mkt_return_vol", "credit_gdp",
  "gdp_pc", "gdp_growth"
)
institutional_terms <- c("trust_index", "rule_of_law", "trust_rule_of_law")

winsorize <- function(x, probabilities = c(0.01, 0.99)) {
  cutoffs <- quantile(x, probabilities, na.rm = TRUE, names = FALSE)
  pmin(pmax(x, cutoffs[[1]]), cutoffs[[2]])
}

prepare_analysis_data <- function(path = file.path(ROOT, "data", "analysis_ready.rds"),
                                  winsor_probs = c(0.01, 0.99), sample_controls = controls) {
  raw <- readRDS(path)
  required <- unique(c("cn_id", "idstudy", "size", "rule_of_law", sample_controls,
                       trust_components, placebo_components))
  absent <- setdiff(required, names(raw))
  if (length(absent)) stop("Input data omit required variables: ", paste(absent, collapse = ", "))

  prepared <- raw %>%
    filter(!grepl("^GR[0-9]+$", cn_id)) %>%
    mutate(
      avg_year = avg_year - 1983.301,
      num_years = log(num_years), cits = log(cits),
      mkt_return_vol = 100 * mkt_return_vol,
      credit_gdp = credit_gdp / 100, gdp_pc = log(gdp_pc),
      gdp_growth = 100 * gdp_growth,
      trust_index = rowMeans(across(all_of(trust_components)), na.rm = TRUE),
      placebo_index = rowMeans(across(all_of(placebo_components)), na.rm = TRUE)
    )
  prepared$trust_index[rowSums(!is.na(prepared[trust_components])) == 0] <- NA_real_
  prepared$placebo_index[rowSums(!is.na(prepared[placebo_components])) == 0] <- NA_real_

  model_frame <- prepared %>%
    select(cn_id, idstudy, size, rule_of_law, placebo_index,
           control_of_corruption, government_effectiveness, regulatory_quality,
           all_of(controls), trust_index) %>%
    drop_na(cn_id, idstudy, size, rule_of_law, trust_index, all_of(sample_controls)) %>%
    filter(if_all(c(size, rule_of_law, trust_index, all_of(sample_controls)), is.finite))

  # Tail treatment comes after the complete-case frame is fixed.
  model_frame <- model_frame %>%
    mutate(size_unwinsorized = size, se_unwinsorized = se,
           size = winsorize(size, winsor_probs), se = winsorize(se, winsor_probs),
           trust_rule_of_law = trust_index * rule_of_law)
  model_frame
}

baseline_formula <- as.formula(paste(
  "size ~", paste(c(controls, institutional_terms), collapse = " + ")
))

cluster_table <- function(model, cluster) {
  covariance <- sandwich::vcovCL(model, cluster = cluster, type = "HC1")
  test <- lmtest::coeftest(model, vcov. = covariance)
  data.frame(term = rownames(test), estimate = test[, 1], std_error = test[, 2],
             statistic = test[, 3], p_value = test[, 4], row.names = NULL,
             check.names = FALSE)
}

validate_baseline_data <- function(df) {
  checks <- c(
    observations = nrow(df) == 1613L,
    studies = dplyr::n_distinct(df$idstudy) == 105L,
    countries = dplyr::n_distinct(df$cn_id) == 31L,
    us_observations = sum(df$cn_id == "US") == 1115L,
    all_model_values_finite = all(is.finite(as.matrix(df[c("size", controls, institutional_terms)])))
  )
  if (!all(checks)) stop("Baseline data validation failed: ", paste(names(checks)[!checks], collapse = ", "))
  data.frame(check = names(checks), passed = unname(checks))
}
