import itertools
import json
import numpy as np


def solve_allocation(R, cost, weight, eta, budget):
    """Exact multiple-choice knapsack; one expert per workload."""
    R = np.asarray(R, float)
    cost = np.asarray(cost, float)
    weight = np.asarray(weight, float)
    eta = np.asarray(eta, float)
    K, J = R.shape
    best = None
    best_obj = -np.inf
    for choice in itertools.product(range(K), repeat=J):
        choice = np.asarray(choice, int)
        total_cost = np.sum(weight * cost[choice, np.arange(J)])
        if total_cost <= budget + 1e-12:
            objective = np.sum(weight * (R[choice, np.arange(J)] - eta))
            # deterministic tie break by lexicographic choice
            key = tuple(choice.tolist())
            if objective > best_obj + 1e-12 or (abs(objective-best_obj) <= 1e-12 and (best is None or key < tuple(best.tolist()))):
                best_obj = objective
                best = choice.copy()
    if best is None:
        raise ValueError("No feasible allocation")
    return best, best_obj, np.sum(weight * cost[best, np.arange(J)])


def allocation_score(q, choice, weight):
    q = np.asarray(q)
    return float(np.sum(np.asarray(weight) * q[choice, np.arange(q.shape[1])]))


def run(seed=7):
    rng = np.random.default_rng(seed)
    # Each row is an expert and each column a workload cluster.
    J, K = 6, 4
    weight = rng.uniform(.5, 1.5, J)
    cost = rng.uniform(.7, 2.0, (K, J))
    R = rng.uniform(.45, .95, (K, J))
    true = np.clip(R + rng.normal(0, .08, (K, J)), 0, 1)
    budget = .60 * np.sum(weight * np.min(cost, axis=0)) + .40 * np.sum(weight * np.max(cost, axis=0))

    nominal, nom_obj, nom_cost = solve_allocation(R, cost, weight, np.zeros(J), budget)
    rows = []
    for eta_level in np.linspace(0, .8, 9):
        eta = np.full(J, eta_level)
        pess, pess_obj, pess_cost = solve_allocation(R, cost, weight, eta, budget)
        rows.append({
            "eta": float(eta_level),
            "disputed_clusters": int(np.sum(nominal != pess)),
            "same_route": bool(np.array_equal(nominal, pess)),
            "nominal_true_score": allocation_score(true, nominal, weight),
            "pessimistic_true_score": allocation_score(true, pess, weight),
            "objective_shift": float(nom_obj - pess_obj),
        })

    # Broad exact sweep: prediction is zero disagreements for every common eta.
    total = 0
    violations = 0
    for _ in range(40):
        jj = int(rng.integers(2, 5)); kk = int(rng.integers(2, 5))
        w = rng.uniform(.1, 2, jj)
        c = rng.uniform(.1, 3, (kk, jj))
        r = rng.uniform(-2, 2, (kk, jj))
        min_cost = np.sum(w * np.min(c, axis=0))
        b = min_cost + rng.uniform(0, 1) * np.sum(w * np.ptp(c, axis=0))
        a, _, _ = solve_allocation(r, c, w, np.zeros(jj), b)
        for e in [0, .1, 1, 100]:
            z, _, _ = solve_allocation(r, c, w, np.full(jj, e), b)
            total += 1
            violations += int(not np.array_equal(a, z))

    # A small comparison with a noisy nominal router and an always-cheapest baseline.
    # Robust policy is mathematically identical to nominal here.
    true_nom = allocation_score(true, nominal, weight)
    cheapest = np.argmin(cost, axis=0)
    cheap_score = allocation_score(true, cheapest, weight)
    predicted_shift = float(np.sum(weight * np.full(J, .8)))
    shift_error = max(abs(x["objective_shift"] - float(np.sum(weight * np.full(J, x["eta"])))) for x in rows)
    checks = {
        "prediction_1_zero_disagreements": violations == 0,
        "prediction_2_shift_error": float(shift_error),
        "prediction_2_pass": shift_error < 1e-9,
        "prediction_3_identical_quality": abs(true_nom - allocation_score(true, pess, weight)) < 1e-12,
    }
    return {
        "seed": seed,
 "shape": [J, K], "budget": float(budget),
        "nominal_choice": nominal.tolist(), "nominal_cost": float(nom_cost),
        "cheapest_choice": cheapest.tolist(),
        "cheapest_cost": float(np.sum(weight * cost[cheapest, np.arange(J)])),
        "sweep": rows,
        "random_sweep_instances": total, "random_sweep_disagreements": violations,
        "nominal_true_weighted_quality": true_nom,
        "robust_true_weighted_quality": true_nom,
        "cheapest_true_weighted_quality": cheap_score,
        "checks": checks,
    }

if __name__ == "__main__":
    print(json.dumps(run(), indent=2))