args <- commandArgs(trailingOnly = FALSE)
file_arg <- grep("^--file=", args, value = TRUE)
analysis_dir <- if (length(file_arg)) dirname(normalizePath(sub("^--file=", "", file_arg[[1]]), mustWork = FALSE)) else file.path(getwd(), "analysis")
source(file.path(analysis_dir, "00_config.R"))
source(file.path(analysis_dir, "lib_timss.R"))

primary_estimator <- function(data, weight, map, threshold) {
  math_t0 <- data[[map$math_t0]]
  math_t1 <- data[[map$math_t1]]
  science_t0 <- data[[map$science_t0]]
  science_t1 <- data[[map$science_t1]]
  state_t0 <- achievement_state(math_t0, science_t0, threshold)
  state_t1 <- achievement_state(math_t1, science_t1, threshold)
  state_levels <- levels(state_t0)

  transitions <- transition_vector(state_t0, state_t1, weight)
  baseline_shares <- setNames(
    vapply(state_levels, function(state) weighted_probability(state_t0 == state, weight), numeric(1)),
    paste0("baseline_", state_levels)
  )
  followup_shares <- setNames(
    vapply(state_levels, function(state) weighted_probability(state_t1 == state, weight), numeric(1)),
    paste0("followup_", state_levels)
  )

  baseline_bb <- state_t0 == "BB"
  baseline_any_deficit <- state_t0 != "AA"
  weakest_t0 <- pmin(math_t0, science_t0) - threshold
  weakest_t1 <- pmin(math_t1, science_t1) - threshold
  extras <- c(
    any_deficit_to_AA = weighted_probability(state_t1[baseline_any_deficit] == "AA", weight[baseline_any_deficit]),
    bb_exit = weighted_probability(state_t1[baseline_bb] != "BB", weight[baseline_bb]),
    weakest_margin_change_BB = weighted_mean((weakest_t1 - weakest_t0)[baseline_bb], weight[baseline_bb]),
    math_change_BB = weighted_mean((math_t1 - math_t0)[baseline_bb], weight[baseline_bb]),
    science_change_BB = weighted_mean((science_t1 - science_t0)[baseline_bb], weight[baseline_bb])
  )
  c(transitions, baseline_shares, followup_shares, extras)
}

analyze_system <- function(data, band, threshold) {
  draw_results <- lapply(seq_len(config$pv_count), function(draw) {
    map <- pv_variable_map(band, draw)
    estimator <- function(dataset, weight) primary_estimator(dataset, weight, map, threshold)
    result <- jrr_estimate(data, estimator)
    result$replicates <- NULL
    result
  })
  combine_pv_results(draw_results)
}

analyze_band <- function(data, band, threshold) {
  systems <- split(data, data$system)
  results <- lapply(systems, analyze_system, band = band, threshold = threshold)
  names(results) <- names(systems)
  results
}

result_to_long <- function(result, band, system, threshold) {
  data.frame(
    band = band,
    system = system,
    threshold = threshold,
    parameter = names(result$estimate),
    estimate = as.numeric(result$estimate),
    standard_error = as.numeric(result$se),
    ci_lower = as.numeric(result$estimate - 1.96 * result$se),
    ci_upper = as.numeric(result$estimate + 1.96 * result$se),
    stringsAsFactors = FALSE
  )
}

message("Estimating Grade 4-5 primary transition models with full PV x JRR inference")
g45 <- readRDS(file.path(paths$derived, "timss_l_g45_analysis.rds"))
g45_results <- analyze_band(g45, "g45", config$primary_benchmark)

message("Estimating Grade 8-9 extension transition models with full PV x JRR inference")
g89 <- readRDS(file.path(paths$derived, "timss_l_g89_analysis.rds"))
g89_results <- analyze_band(g89, "g89", config$primary_benchmark)

primary_results <- list(
  threshold = config$primary_benchmark,
  g45 = g45_results,
  g89 = g89_results
)
saveRDS(primary_results, file.path(paths$derived, "primary_transition_results.rds"), compress = "xz")

primary_long <- do.call(rbind, c(
  Map(function(result, system) result_to_long(result, "g45", system, config$primary_benchmark), g45_results, names(g45_results)),
  Map(function(result, system) result_to_long(result, "g89", system, config$primary_benchmark), g89_results, names(g89_results))
))
rownames(primary_long) <- NULL
write.csv(primary_long, file.path(paths$tables, "primary_transition_estimates.csv"), row.names = FALSE)

headline <- primary_long[primary_long$band == "g45" & primary_long$parameter %in% c("BB_to_AA", "any_deficit_to_AA", "BB_to_BB"), ]
write.csv(headline, file.path(paths$tables, "primary_headline_estimates.csv"), row.names = FALSE)
message("Primary transition analysis complete: ", nrow(primary_long), " estimates")
