system_labels <- c(
  `268` = "Georgia",
  `380` = "Italy",
  `400` = "Jordan",
  `410` = "Korea, Rep. of",
  `411` = "Kosovo",
  `499` = "Montenegro",
  `705` = "Slovenia",
  `752` = "Sweden",
  `807` = "North Macedonia"
)

as_plain_numeric <- function(x) {
  suppressWarnings(as.numeric(x))
}

load_single_object <- function(file) {
  environment <- new.env(parent = emptyenv())
  object_names <- load(file, envir = environment)
  if (length(object_names) != 1L) stop("Expected one object in ", file)
  environment[[object_names[[1]]]]
}

pv_variable_map <- function(band, draw) {
  stopifnot(draw %in% 1:5)
  prefix <- if (band == "g45") "A" else if (band == "g89") "B" else stop("Unknown band")
  list(
    math_t0 = sprintf("%sSMMAT1%d", prefix, draw),
    math_t1 = sprintf("%sSMMAT2%d", prefix, draw),
    science_t0 = sprintf("%sSSSCI1%d", prefix, draw),
    science_t1 = sprintf("%sSSSCI2%d", prefix, draw),
    math_change = sprintf("%sSMCHG0%d", prefix, draw),
    science_change = sprintf("%sSSCHG0%d", prefix, draw)
  )
}

all_pv_variables <- function(band) {
  unique(unlist(lapply(1:5, function(draw) unlist(pv_variable_map(band, draw))), use.names = FALSE))
}

load_band_data <- function(raw_directory, band) {
  achievement_prefix <- if (band == "g45") "asa" else if (band == "g89") "bsa" else stop("Unknown band")
  process_prefix <- if (band == "g45") "asp" else "bsp"
  files <- sort(list.files(raw_directory, pattern = paste0("^", achievement_prefix, ".*8l[.]rdata$"), full.names = TRUE))
  if (!length(files)) stop("No achievement files found in ", raw_directory)

  required_achievement <- c(
    "IDCNTRY", "IDSCHOOL", "IDSTUD", "TOTWGT", "JKZONE", "JKREP",
    all_pv_variables(band)
  )
  required_process <- c(
    "IDCNTRY", "IDSTUD", "ITPARTPT1_23", "ITPARTPT2_23",
    "ITPARTPT1_24", "ITPARTPT2_24"
  )

  pieces <- lapply(files, function(file) {
    achievement <- load_single_object(file)
    absent <- setdiff(required_achievement, names(achievement))
    if (length(absent)) stop("Missing achievement fields in ", basename(file), ": ", paste(absent, collapse = ", "))

    process_file <- file.path(raw_directory, sub(paste0("^", achievement_prefix), process_prefix, basename(file)))
    if (!file.exists(process_file)) stop("Missing process file: ", process_file)
    process <- load_single_object(process_file)
    absent_process <- setdiff(required_process, names(process))
    if (length(absent_process)) stop("Missing process fields in ", basename(process_file), ": ", paste(absent_process, collapse = ", "))
    participation_fields <- required_process[grepl("^ITPARTPT", required_process)]
    participation_codes <- do.call(cbind, lapply(participation_fields, function(variable) {
      as_plain_numeric(process[[variable]])
    }))
    nonmissing_participation_codes <- participation_codes[!is.na(participation_codes)]
    if (
      any(!is.finite(nonmissing_participation_codes)) ||
      any(!(nonmissing_participation_codes %in% c(1, 2)))
    ) {
      stop("Nonmissing participation codes must be finite and exactly 1 or 2 in ", basename(process_file))
    }

    key_achievement <- paste(as_plain_numeric(achievement$IDCNTRY), as_plain_numeric(achievement$IDSTUD), sep = ":")
    key_process <- paste(as_plain_numeric(process$IDCNTRY), as_plain_numeric(process$IDSTUD), sep = ":")
    if (anyDuplicated(key_achievement) || anyDuplicated(key_process)) stop("Duplicate student key")
    process_index <- match(key_achievement, key_process)
    if (anyNA(process_index)) stop("Process linkage failed in ", basename(file))

    output <- data.frame(
      band = band,
      system_id = as_plain_numeric(achievement$IDCNTRY),
      school_id = as_plain_numeric(achievement$IDSCHOOL),
      student_id = as_plain_numeric(achievement$IDSTUD),
      weight = as_plain_numeric(achievement$TOTWGT),
      jkzone = as_plain_numeric(achievement$JKZONE),
      jkrep = as_plain_numeric(achievement$JKREP),
      stringsAsFactors = FALSE
    )
    output$system <- unname(system_labels[as.character(output$system_id)])
    design_values <- as.matrix(output[c(
      "system_id", "school_id", "student_id", "weight", "jkzone", "jkrep"
    )])
    if (any(!is.finite(design_values)) || anyNA(output$system)) {
      stop("Identifiers and design variables must be finite and map to configured systems in ", basename(file))
    }
    if (any(output$weight <= 0)) stop("Final student weights must be strictly positive in ", basename(file))
    if (any(!(output$jkrep %in% c(0, 1)))) stop("JKREP must be exactly 0 or 1 in ", basename(file))

    for (variable in all_pv_variables(band)) {
      output[[variable]] <- as_plain_numeric(achievement[[variable]])
    }

    p1_23 <- as_plain_numeric(process$ITPARTPT1_23[process_index])
    p2_23 <- as_plain_numeric(process$ITPARTPT2_23[process_index])
    p1_24 <- as_plain_numeric(process$ITPARTPT1_24[process_index])
    p2_24 <- as_plain_numeric(process$ITPARTPT2_24[process_index])
    output$participated_2023 <- as.integer(rowSums(cbind(p1_23 == 1, p2_23 == 1), na.rm = TRUE) > 0)
    output$participated_2024 <- as.integer(rowSums(cbind(p1_24 == 1, p2_24 == 1), na.rm = TRUE) > 0)
    output$participation_pattern <- ifelse(
      output$participated_2023 == 1 & output$participated_2024 == 1, "both",
      ifelse(output$participated_2023 == 1, "2023_only", ifelse(output$participated_2024 == 1, "2024_only", "neither"))
    )
    output
  })

  combined <- do.call(rbind, pieces)
  rownames(combined) <- NULL
  zone_key <- paste(combined$system_id, combined$jkzone, sep = ":")
  members_by_zone <- split(combined$jkrep, zone_key)
  complete_member_pair <- vapply(
    members_by_zone,
    function(member) setequal(unique(member), c(0, 1)),
    logical(1)
  )
  if (!all(complete_member_pair)) {
    stop(
      "Every system-specific JKZONE must contain both JKREP members; invalid zone(s): ",
      paste(names(complete_member_pair)[!complete_member_pair], collapse = ", ")
    )
  }
  combined
}

weighted_mean <- function(x, weight) {
  valid <- is.finite(x) & is.finite(weight) & weight > 0
  if (!any(valid)) return(NA_real_)
  sum(x[valid] * weight[valid]) / sum(weight[valid])
}

weighted_probability <- function(condition, weight) {
  weighted_mean(as.numeric(condition), weight)
}

achievement_state <- function(math, science, threshold = 400) {
  state <- rep(NA_character_, length(math))
  valid <- is.finite(math) & is.finite(science)
  state[valid & math < threshold & science < threshold] <- "BB"
  state[valid & math < threshold & science >= threshold] <- "MB"
  state[valid & math >= threshold & science < threshold] <- "SB"
  state[valid & math >= threshold & science >= threshold] <- "AA"
  factor(state, levels = c("BB", "MB", "SB", "AA"))
}

replicate_weight <- function(weight, jkzone, jkrep, zone, included_member) {
  factor <- rep(1, length(weight))
  in_zone <- jkzone == zone
  factor[in_zone & jkrep == included_member] <- 2
  factor[in_zone & jkrep != included_member] <- 0
  weight * factor
}

weighted_quantile <- function(x, weight, probabilities) {
  valid <- is.finite(x) & is.finite(weight) & weight > 0
  x <- x[valid]
  weight <- weight[valid]
  if (!length(x)) return(rep(NA_real_, length(probabilities)))
  order_index <- order(x)
  x <- x[order_index]
  weight <- weight[order_index]
  cumulative <- cumsum(weight) / sum(weight)
  vapply(probabilities, function(probability) x[which(cumulative >= probability)[1]], numeric(1))
}

weighted_effective_n <- function(weight) {
  weight <- weight[is.finite(weight) & weight > 0]
  if (!length(weight)) return(0)
  sum(weight)^2 / sum(weight^2)
}

jrr_estimate <- function(data, estimator, weight = data$weight) {
  point_named <- estimator(data, weight)
  point <- as.numeric(point_named)
  parameter_names <- names(point_named)
  if (is.null(parameter_names)) parameter_names <- paste0("parameter_", seq_along(point))
  names(point) <- parameter_names
  zones <- sort(unique(data$jkzone[is.finite(data$jkzone)]))
  replicates <- matrix(NA_real_, nrow = 2L * length(zones), ncol = length(point), dimnames = list(NULL, parameter_names))
  row <- 1L
  for (zone in zones) {
    for (included in c(0, 1)) {
      replicate <- replicate_weight(weight, data$jkzone, data$jkrep, zone, included)
      replicates[row, ] <- as.numeric(estimator(data, replicate))
      row <- row + 1L
    }
  }
  deviations <- sweep(replicates, 2L, point, "-")
  variance <- crossprod(deviations) / 2
  dimnames(variance) <- list(parameter_names, parameter_names)
  list(estimate = point, variance = variance, replicates = replicates, zones = zones)
}

combine_pv_results <- function(results) {
  estimates <- do.call(rbind, lapply(results, `[[`, "estimate"))
  within <- Reduce(`+`, lapply(results, `[[`, "variance")) / length(results)
  between <- if (nrow(estimates) > 1L) stats::cov(estimates) else matrix(0, ncol(estimates), ncol(estimates))
  total <- within + (1 + 1 / nrow(estimates)) * between
  point <- colMeans(estimates)
  standard_error <- sqrt(diag(total))
  list(
    estimate = point,
    within = within,
    between = between,
    variance = total,
    se = standard_error,
    draw_estimates = estimates
  )
}

transition_vector <- function(state_t0, state_t1, weight) {
  state_levels <- c("BB", "MB", "SB", "AA")
  parameter_names <- unlist(lapply(state_levels, function(origin) paste(origin, "to", state_levels, sep = "_")), use.names = FALSE)
  values <- setNames(numeric(length(parameter_names)), parameter_names)
  valid <- !is.na(state_t0) & !is.na(state_t1) &
    as.character(state_t0) %in% state_levels & as.character(state_t1) %in% state_levels &
    is.finite(weight) & weight > 0
  index <- 1L
  for (origin in state_levels) {
    origin_valid <- valid & state_t0 == origin
    denominator <- sum(weight[origin_valid])
    if (!is.finite(denominator) || denominator <= 0) {
      stop("No finite positive-weight jointly valid records for transition origin ", origin)
    }
    row_values <- numeric(length(state_levels))
    for (destination_index in seq_along(state_levels)) {
      destination <- state_levels[destination_index]
      numerator <- sum(weight[origin_valid & state_t1 == destination])
      row_values[destination_index] <- numerator / denominator
      values[index] <- row_values[destination_index]
      index <- index + 1L
    }
    if (any(!is.finite(row_values)) || any(row_values < 0 | row_values > 1) ||
        abs(sum(row_values) - 1) > 1e-10) {
      stop("Transition row failed the jointly valid row-sum gate for origin ", origin)
    }
  }
  values
}
