# ==============================================================================
# Constrained Automated Test Assembly (ATA) via Maximin MILP in R
# Author: Nedim Yel, Ph.D. (nedimyel.com | statisticalsolutions.org)
# ==============================================================================

library(lpSolve)

assemble_irt_short_form <- function(info_matrix,      # Matrix (Q x I): quadrature points x items
                                    item_domains,     # Character vector of length I
                                    domain_min,       # Named numeric vector of min items per domain
                                    reading_grade,    # Numeric vector of length I
                                    ld_enemy_pairs,   # 2-column matrix of locally dependent item indices
                                    k_short = 6,
                                    max_avg_grade = 7.0,
                                    weights = rep(1, nrow(info_matrix)),
                                    domain_max = NULL,
                                    dif_excluded = integer(0),
                                    floor_tolerance = 1e-7) {
  
  Q <- nrow(info_matrix)
  I <- ncol(info_matrix)
  domains <- names(domain_min)
  
  stopifnot(is.matrix(info_matrix), Q > 0, I > 0,
            all(is.finite(info_matrix)), all(info_matrix >= 0),
            length(item_domains) == I, !anyNA(item_domains),
            length(reading_grade) == I, all(is.finite(reading_grade)),
            all(reading_grade >= 0), length(weights) == Q,
            all(is.finite(weights)), all(weights > 0),
            length(k_short) == 1, k_short >= 1, k_short <= I, k_short == floor(k_short),
            is.finite(max_avg_grade), max_avg_grade >= 0,
            is.finite(floor_tolerance), floor_tolerance >= 0)
  for (limits in list(domain_min, domain_max)) {
    if (length(limits)) stopifnot(length(names(limits)) == length(limits),
      all(names(limits) %in% item_domains), all(is.finite(limits)),
      all(limits >= 0), all(limits == floor(limits)))
  }
  stopifnot(all(dif_excluded %in% seq_len(I)))
  if (!is.null(ld_enemy_pairs) && length(ld_enemy_pairs)) {
    stopifnot(is.matrix(ld_enemy_pairs), ncol(ld_enemy_pairs) == 2,
      all(ld_enemy_pairs %in% seq_len(I)),
      all(ld_enemy_pairs[, 1] != ld_enemy_pairs[, 2]))
  }
  # Stage 1 maximizes the floor only.
  f_obj <- c(rep(0, I), 1)

  con_rows <- list()
  con_dirs <- character(0)
  con_rhs  <- numeric(0)
  
  # 1. Maximin Quadrature Floor Constraints: sum_i I_i(theta_q)*x_i - w_q*y >= 0
  for (q in seq_len(Q)) {
    con_rows[[length(con_rows) + 1]] <- c(info_matrix[q, ], -weights[q])
    con_dirs <- c(con_dirs, ">=")
    con_rhs  <- c(con_rhs, 0)
  }
  
  # 2. Exact Short-Form Length: sum(x_i) == k_short
  con_rows[[length(con_rows) + 1]] <- c(rep(1, I), 0)
  con_dirs <- c(con_dirs, "==")
  con_rhs  <- c(con_rhs, k_short)
  
  # 3. Content Blueprint Constraints: sum_{i in c} x_i >= domain_min[c]
  for (dom in domains) {
    dom_indicator <- as.numeric(item_domains == dom)
    con_rows[[length(con_rows) + 1]] <- c(dom_indicator, 0)
    con_dirs <- c(con_dirs, ">=")
    con_rhs  <- c(con_rhs, as.numeric(domain_min[[dom]]))
  }
  
  for (dom in names(domain_max)) {
    con_rows[[length(con_rows) + 1]] <- c(as.numeric(item_domains == dom), 0)
    con_dirs <- c(con_dirs, "<=")
    con_rhs <- c(con_rhs, domain_max[[dom]])
  }
  for (idx in dif_excluded) {
    row <- rep(0, I + 1)
    row[idx] <- 1
    con_rows[[length(con_rows) + 1]] <- row
    con_dirs <- c(con_dirs, "==")
    con_rhs <- c(con_rhs, 0)
  }
  # 4. Reading Grade Ceiling: sum(grade_i * x_i) <= max_avg_grade * k_short
  con_rows[[length(con_rows) + 1]] <- c(reading_grade, 0)
  con_dirs <- c(con_dirs, "<=")
  con_rhs  <- c(con_rhs, max_avg_grade * k_short)
  
  # 5. Yen's Q3 Local Independence Exclusion: x_i + x_j <= 1
  if (!is.null(ld_enemy_pairs) && nrow(ld_enemy_pairs) > 0) {
    for (r in seq_len(nrow(ld_enemy_pairs))) {
      pair_vec <- rep(0, I + 1)
      pair_vec[ld_enemy_pairs[r, 1]] <- 1
      pair_vec[ld_enemy_pairs[r, 2]] <- 1
      con_rows[[length(con_rows) + 1]] <- pair_vec
      con_dirs <- c(con_dirs, "<=")
      con_rhs  <- c(con_rhs, 1)
    }
  }
  
  # 6. Binary Upper Bounds for x_i (x_i <= 1) while y remains continuous
  for (i in seq_len(I)) {
    ub_vec <- rep(0, I + 1)
    ub_vec[i] <- 1
    con_rows[[length(con_rows) + 1]] <- ub_vec
    con_dirs <- c(con_dirs, "<=")
    con_rhs  <- c(con_rhs, 1)
  }
  
  const_mat <- do.call(rbind, con_rows)
  
  sol <- lpSolve::lp(
    direction = "max",
    objective.in = f_obj,
    const.mat = const_mat,
    const.dir = con_dirs,
    const.rhs = con_rhs,
    int.vec = seq_len(I) # First I variables are binary; variable (I+1) is continuous y
  )
  
  if (sol$status != 0) return(list(success = FALSE, stage = 1, status = sol$status))
  optimum_floor <- sol$solution[I + 1]
  sol <- lpSolve::lp("max", c(colSums(info_matrix), 0),
    rbind(const_mat, c(rep(0, I), 1)), c(con_dirs, ">="),
    c(con_rhs, max(0, optimum_floor - floor_tolerance)), int.vec = seq_len(I))
  if (sol$status != 0) return(list(success = FALSE, stage = 2, status = sol$status))
  selected_idx <- which(round(sol$solution[seq_len(I)]) == 1)
  list(
    success = TRUE,
    status = sol$status,
    stage1_floor = optimum_floor,
    floor_tolerance = floor_tolerance,
    selected_items = selected_idx,
    min_target_info_floor = min(rowSums(info_matrix[, selected_idx, drop = FALSE]) / weights),
    target_csem_profile = round(1 / sqrt(rowSums(info_matrix[, selected_idx, drop = FALSE])), 3)
  )
}

run_synthetic_example <- function() {
  information <- matrix(c(4,2,1,3,1,2, 2,3,4,2,1,2, 1,2,3,1,4,2),
                        nrow = 3, byrow = TRUE)
  domains <- c("A", "A", "B", "B", "C", "C")
  arguments <- list(info_matrix = information, item_domains = domains,
                    domain_min = c(A=1, B=1, C=1), domain_max = c(A=1, B=1, C=1),
                    reading_grade = rep(6, 6), ld_enemy_pairs = matrix(c(1,3), ncol=2),
                    dif_excluded = 6, k_short = 3)
  result <- do.call(assemble_irt_short_form, arguments)
  stopifnot(result$success)
  subsets <- combn(seq_len(6), 3, simplify = FALSE)
  feasible <- Filter(function(x) !6 %in% x && !all(c(1,3) %in% x) &&
                       all(vapply(c("A","B","C"), function(d) sum(domains[x] == d) == 1,
                                  logical(1))), subsets)
  floors <- vapply(feasible, function(x) min(rowSums(information[,x])), numeric(1))
  totals <- vapply(feasible, function(x) sum(information[,x]), numeric(1))
  best_floor <- max(floors)
  best_total <- max(totals[floors >= best_floor - 1e-7])
  stopifnot(identical(result$selected_items, c(2L,4L,5L)),
            abs(result$min_target_info_floor - best_floor) < 1e-6,
            abs(sum(information[,result$selected_items]) - best_total) < 1e-6)
  arguments$dif_excluded <- c(1,2)
  stopifnot(!do.call(assemble_irt_short_form, arguments)$success)
  cat("Synthetic example only: no fitted responses or clinical validation.\n")
  cat("Selected item numbers (1-based):", result$selected_items, "\n")
  cat("Information at the three target points:", rowSums(information[,result$selected_items]), "\n")
  cat("Minimum information:", result$min_target_info_floor, "\n")
  cat("PASS: exhaustive objective comparison and infeasible-case handling.\n")
  invisible(result)
}

if (sys.nframe() == 0L) run_synthetic_example()