# ==============================================================================
# ata_maximin_solver.py
# Annotated Python Implementation of Maximin IRT Automated Test Assembly (ATA)
# Author: Nedim Yel, Ph.D. (nedimyel.com | statisticalsolutions.org)
# ==============================================================================

from typing import Dict, List, Tuple
import numpy as np
from scipy.optimize import Bounds, LinearConstraint, milp


def assemble_irt_short_form_py(
    info_matrix: np.ndarray,                 # 2D array of shape (Q, I): quadrature points x items
    item_domains: np.ndarray,                # 1D array of strings of length I
    domain_min: Dict[str, int],              # Dict mapping each domain name to min required items
    reading_grade: np.ndarray,               # 1D float array of length I
    ld_enemy_pairs: List[Tuple[int, int]],   # List of 0-indexed tuples (i, j) reviewed for local dependence
    k_short: int = 6,
    max_avg_grade: float = 7.0,
    weights: np.ndarray = None,
    domain_max: Dict[str, int] = None,
    dif_excluded: List[int] = None,
    floor_tolerance: float = 1e-7,
) -> Dict:
    """
    Solves the Maximin IRT Short-Form Assembly problem using scipy.optimize.milp.
    
    Decision vector z = [x_0, x_1, ..., x_{I-1}, y] of length (I + 1):
      - x_i in {0, 1} indicates whether item i is selected.
      - y >= 0 is the continuous auxiliary variable representing the minimum
        ratio of information to requirement across the target quadrature grid.
    """
    info_matrix = np.asarray(info_matrix, dtype=float)
    item_domains = np.asarray(item_domains)
    reading_grade = np.asarray(reading_grade, dtype=float)
    if info_matrix.ndim != 2 or min(info_matrix.shape) < 1:
        raise ValueError("Information must be a nonempty Q by I matrix")
    Q, I = info_matrix.shape
    if (not np.isfinite(info_matrix).all() or (info_matrix < 0).any()
        or item_domains.shape != (I,) or reading_grade.shape != (I,)
        or not np.isfinite(reading_grade).all() or (reading_grade < 0).any()
        or not 1 <= k_short <= I or int(k_short) != k_short
        or not np.isfinite(max_avg_grade) or max_avg_grade < 0
        or not np.isfinite(floor_tolerance) or floor_tolerance < 0):
        raise ValueError("Invalid information, metadata, length or tolerance")
    if weights is None:
        weights = np.ones(Q, dtype=float)

    weights = np.asarray(weights, dtype=float)
    if weights.shape != (Q,) or not np.isfinite(weights).all() or (weights <= 0).any():
        raise ValueError("Requirements must be finite and positive")
    for limits in (domain_min, domain_max or {}):
        for domain, count in limits.items():
            if domain not in item_domains or not np.isfinite(count) or count < 0 or int(count) != count:
                raise ValueError("Invalid domain bound")
    c_obj = np.concatenate([np.zeros(I), [-1.0]])

    # Lists to accumulate linear constraint rows, lower bounds (lb), and upper bounds (ub)
    # SciPy's LinearConstraint takes the unified form: lb <= A @ z <= ub
    rows = []
    lb_list = []
    ub_list = []

    # 2. Maximin Quadrature Floor Constraints:
    # sum_i I_i(theta_q) * x_i - w_q * y >= 0  --->  0 <= [I(theta_q), -w_q] @ z <= +inf
    for q in range(Q):
        row = np.concatenate([info_matrix[q, :], [-weights[q]]])
        rows.append(row)
        lb_list.append(0.0)
        ub_list.append(np.inf)

    # 3. Exact Short-Form Length Constraint:
    # sum_i x_i == k_short  --->  k_short <= [1, 1, ..., 1, 0] @ z <= k_short
    rows.append(np.concatenate([np.ones(I), [0.0]]))
    lb_list.append(float(k_short))
    ub_list.append(float(k_short))

    # 4. Content Blueprint Constraints:
    # For each subdomain, sum_{i in dom} x_i >= min_count
    for dom_name, min_count in domain_min.items():
        dom_mask = (item_domains == dom_name).astype(float)
        rows.append(np.concatenate([dom_mask, [0.0]]))
        lb_list.append(float(min_count))
        ub_list.append(np.inf)

    for domain, count in (domain_max or {}).items():
        rows.append(np.concatenate([(item_domains == domain).astype(float), [0.0]]))
        lb_list.append(-np.inf)
        ub_list.append(float(count))
    for idx in dif_excluded or []:
        if not isinstance(idx, (int, np.integer)) or not 0 <= idx < I:
            raise ValueError("Invalid DIF item index")
        row = np.zeros(I + 1)
        row[idx] = 1
        rows.append(row)
        lb_list.append(0.0)
        ub_list.append(0.0)

    # 5. Reading Grade Ceiling Constraint:
    # sum_i grade_i * x_i <= max_avg_grade * k_short
    rows.append(np.concatenate([reading_grade.astype(float), [0.0]]))
    lb_list.append(-np.inf)
    ub_list.append(float(max_avg_grade * k_short))

    # 6. Yen's Q3 Local Independence Exclusion (Enemy Pairs):
    # x_i + x_j <= 1.0
    for i_idx, j_idx in ld_enemy_pairs:
        if (not all(isinstance(j, (int, np.integer)) and 0 <= j < I for j in (i_idx, j_idx))
            or i_idx == j_idx):
            raise ValueError("Invalid enemy pair")
        pair_row = np.zeros(I + 1, dtype=float)
        pair_row[i_idx] = 1.0
        pair_row[j_idx] = 1.0
        rows.append(pair_row)
        lb_list.append(-np.inf)
        ub_list.append(1.0)

    # Stack all constraint rows into a 2D matrix A (equivalent to do.call(rbind, con_rows) in R)
    A_mat = np.vstack(rows)
    constraints = LinearConstraint(A_mat, lb_list, ub_list)

    # 7. Variable Bounds & Integrality:
    # First I variables (x_i) are binary in [0, 1]; the last variable (y) is continuous in [0, +inf)
    var_lb = np.zeros(I + 1)
    var_ub = np.concatenate([np.ones(I), [np.inf]])
    bounds = Bounds(var_lb, var_ub)

    # integrality = 1 enforces integer values; integrality = 0 allows continuous real values
    integrality = np.concatenate([np.ones(I), [0.0]])

    # 8. Execute the HiGHS MILP Solver
    res = milp(c=c_obj, integrality=integrality, bounds=bounds, constraints=constraints, options={"mip_rel_gap": 0.0})

    if not res.success or res.x is None:
        return {"success": False, "stage": 1, "status": int(res.status), "message": res.message}
    optimum_floor = float(res.x[I])
    floor_row = np.zeros(I + 1)
    floor_row[-1] = 1
    constraints2 = LinearConstraint(
        np.vstack([A_mat, floor_row]),
        lb_list + [max(0.0, optimum_floor - floor_tolerance)], ub_list + [np.inf])
    res = milp(c=np.concatenate([-info_matrix.sum(axis=0), [0.0]]),
               integrality=integrality, bounds=bounds, constraints=constraints2,
               options={"mip_rel_gap": 0.0})
    if not res.success or res.x is None:
        return {"success": False, "stage": 2, "status": int(res.status), "message": res.message}

    # Extract binary item selections and compute resulting conditional standard errors
    x_sol = np.round(res.x[:I]).astype(bool)
    selected_indices = np.where(x_sol)[0].tolist()
    tif_target = info_matrix[:, x_sol].sum(axis=1)
    csem_target = np.full(Q, np.inf)
    np.divide(1.0, np.sqrt(tif_target), out=csem_target, where=tif_target > 0)

    return {
        "success": bool(res.success),
        "stage1_floor": optimum_floor,
        "floor_tolerance": floor_tolerance,
        "selected_item_indices_0based": selected_indices,
        "min_target_info_floor": float(np.min(tif_target / weights)),
        "target_quadrature_csem": csem_target.tolist(),
    }

def run_synthetic_example():
    """Check an invented six-item bank against exhaustive enumeration."""
    from itertools import combinations
    information = np.array([[4, 2, 1, 3, 1, 2],
                            [2, 3, 4, 2, 1, 2],
                            [1, 2, 3, 1, 4, 2]], dtype=float)
    domains = np.array(["A", "A", "B", "B", "C", "C"])
    arguments = dict(info_matrix=information, item_domains=domains,
                     domain_min={"A": 1, "B": 1, "C": 1},
                     domain_max={"A": 1, "B": 1, "C": 1},
                     reading_grade=np.full(6, 6.0), ld_enemy_pairs=[(0, 2)],
                     dif_excluded=[5], k_short=3)
    result = assemble_irt_short_form_py(**arguments)
    assert result["success"], result
    feasible = [subset for subset in combinations(range(6), 3)
                if 5 not in subset and not {0, 2}.issubset(subset)
                and all(sum(domains[i] == d for i in subset) == 1 for d in "ABC")]
    # Compare the primary and secondary objectives with all feasible subsets.
    scores = [(information[:, subset].sum(axis=1).min(),
               information[:, subset].sum()) for subset in feasible]
    best_floor = max(floor for floor, total in scores)
    best_total = max(total for floor, total in scores if floor >= best_floor - 1e-7)
    selected = result["selected_item_indices_0based"]
    assert selected == [1, 3, 4], selected
    assert np.isclose(result["min_target_info_floor"], best_floor)
    assert np.isclose(information[:, selected].sum(), best_total)
    infeasible = assemble_irt_short_form_py(**(arguments | {"dif_excluded": [0, 1]}))
    assert not infeasible["success"], infeasible
    print("Synthetic example only: no fitted responses or clinical validation.")
    print("Selected item numbers (1-based):", [i + 1 for i in selected])
    print("Information at the three target points:", information[:, selected].sum(axis=1).tolist())
    print("Minimum information:", result["min_target_info_floor"])
    print("PASS: exhaustive objective comparison and infeasible-case handling.")
    return result


if __name__ == "__main__":
    run_synthetic_example()