experiment.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "numpy>=1.26",
#   "pandas>=2.2",
#   "scikit-learn>=1.4",
#   "scipy>=1.12",
#   "plotly>=5.20",
#   "pyarrow>=15.0",
#   "huggingface_hub>=0.24",
# ]
# ///
"""
Synthetic-proxy reproduction of the offline-experiments claim in
"Large-Scale Notification Dispatch with Bundle Treatments and Multi-Outcome
Uplift Optimization" (ICML 2026, OpenReview 8GH752ZJ5j, BUOPLR / Kuaishou).

No public code, data, or arXiv version exists for this paper. This script
builds a self-contained simulator of the paper's setting (combinatorial
bundle treatments, multiple outcomes, global budget/quota/fatigue
constraints) and implements BUOPLR's two-stage algorithmic idea:

  Stage 1 (uplift estimation): a shared multi-task network predicting all
  outcomes jointly from (features, bundle) so it can capture cross-treatment
  and cross-outcome structure.

  Stage 2 (assignment): restrict each user's candidate set to their top-C
  bundles, then solve the global budget/quota/fatigue constraints via
  Lagrangian relaxation (dual subgradient ascent) instead of an exact LP.

These are compared against baseline uplift estimators (linear-additive,
independent-per-outcome MLP) and baseline assignment algorithms (greedy
myopic, exact LP, no-treatment control) on matched synthetic ground truth.
"""
from __future__ import annotations

import argparse
import itertools
import json
import time
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.neural_network import MLPRegressor

OUTCOMES = ["main", "cost", "fatigue"]
K = 4  # number of binary notification "slots" (timing/style dimensions)
MAX_ACTIVE = 2  # restrict bundles to at most this many simultaneous slots
N_FEAT = 10
N_VENDORS = 3


def enumerate_bundles(k: int = K, max_active: int = MAX_ACTIVE) -> list[tuple[int, ...]]:
    bundles = [tuple([0] * k)]
    for r in range(1, max_active + 1):
        for combo in itertools.combinations(range(k), r):
            b = [0] * k
            for i in combo:
                b[i] = 1
            bundles.append(tuple(b))
    return bundles


BUNDLES = enumerate_bundles()
N_ACTIONS = len(BUNDLES)
NULL_ACTION = 0
BUNDLE_MAT = np.array(BUNDLES, dtype=float)  # (N_ACTIONS, K)


# --------------------------------------------------------------------------
# Simulator: heterogeneous, small-effect, cross-slot-interacting outcomes
# --------------------------------------------------------------------------

def make_effect_params(seed: int) -> dict:
    rng = np.random.default_rng(seed)
    w_main = rng.normal(0, 1, size=(K, N_FEAT)) * 0.9
    w_fatigue = np.abs(rng.normal(0, 1, size=(K, N_FEAT))) * 0.5
    unit_cost = rng.uniform(0.8, 1.5, size=K)
    interactions = {}
    for i in range(K):
        for j in range(i + 1, K):
            sign = rng.choice([-1.0, 1.0])
            mag = rng.uniform(0.3, 0.9)
            interactions[(i, j)] = sign * mag
    base_main_w = rng.normal(0, 1, N_FEAT) * 0.3
    return dict(
        w_main=w_main, w_fatigue=w_fatigue, unit_cost=unit_cost,
        interactions=interactions, base_main_w=base_main_w,
    )


def sample_users(n: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed)
    X = rng.normal(0, 1, size=(n, N_FEAT))
    vendor = rng.integers(0, N_VENDORS, size=n)
    return X, vendor


def true_outcomes_for_bundle(X: np.ndarray, bundle: tuple[int, ...], params: dict) -> np.ndarray:
    """Expected [main, cost, fatigue] for every row of X under a fixed bundle."""
    b = np.array(bundle, dtype=float)
    main_raw = X @ params["base_main_w"]
    fatigue_raw = np.full(X.shape[0], -2.5)  # low base -> small base fatigue after sigmoid
    cost = 0.0
    for k in range(K):
        if b[k]:
            main_raw = main_raw + np.tanh(X @ params["w_main"][k]) * 1.4
            fatigue_raw = fatigue_raw + np.abs(np.tanh(X @ params["w_fatigue"][k])) * 1.1
            cost = cost + params["unit_cost"][k]
    for (i, j), coef in params["interactions"].items():
        if b[i] and b[j]:
            main_raw = main_raw + coef * np.tanh(X @ params["w_main"][i]) * np.tanh(X @ params["w_main"][j])
    main = 1.0 / (1.0 + np.exp(-main_raw)) * 0.02  # small-effect engagement probability, paper-motivated
    fatigue = 1.0 / (1.0 + np.exp(-fatigue_raw))
    cost_arr = np.full(X.shape[0], cost)
    return np.stack([main, cost_arr, fatigue], axis=1)


def true_outcomes_all_actions(X: np.ndarray, params: dict) -> np.ndarray:
    """(n, N_ACTIONS, 3) true expected outcomes for every candidate bundle."""
    out = np.empty((X.shape[0], N_ACTIONS, 3))
    for a, bundle in enumerate(BUNDLES):
        out[:, a, :] = true_outcomes_for_bundle(X, bundle, params)
    return out


def simulate_logged_data(n: int, seed: int, params: dict) -> pd.DataFrame:
    """Observational-style training log: uniform random exploration bundle per user."""
    rng = np.random.default_rng(seed)
    X, vendor = sample_users(n, seed)
    action_idx = rng.integers(0, N_ACTIONS, size=n)
    B = BUNDLE_MAT[action_idx]
    exp_outcomes = np.stack(
        [true_outcomes_for_bundle(X[i:i + 1], BUNDLES[action_idx[i]], params)[0] for i in range(n)]
    ) if n <= 20000 else _batched_logged_outcomes(X, action_idx, params)
    # Noise stds are scaled to each outcome's own dynamic range: "main" is a
    # deliberately small-effect signal (max ~0.02, paper-motivated), so using
    # a flat noise std here (as in an earlier version of this script) drowned
    # the label in noise and penalized the higher-capacity MLP far more than
    # the low-variance linear baseline purely as a bias-variance artifact,
    # not a genuine test of the modeling claim.
    noise_std = [0.004, 0.08, 0.03]
    noise = np.random.default_rng(seed + 1).normal(0, noise_std, size=exp_outcomes.shape)
    Y = exp_outcomes + noise
    df = {"vendor": vendor, "action_idx": action_idx}
    for f in range(N_FEAT):
        df[f"x{f}"] = X[:, f]
    for k in range(K):
        df[f"b{k}"] = B[:, k]
    for m, name in enumerate(OUTCOMES):
        df[f"y_{name}"] = Y[:, m]
    return pd.DataFrame(df), X, B, Y


def _batched_logged_outcomes(X, action_idx, params, batch=50000):
    n = X.shape[0]
    out = np.empty((n, 3))
    for a in range(N_ACTIONS):
        mask = action_idx == a
        if mask.any():
            out[mask] = true_outcomes_for_bundle(X[mask], BUNDLES[a], params)
    return out


# --------------------------------------------------------------------------
# Stage 1: uplift estimators
# --------------------------------------------------------------------------

def featurize_plain(X: np.ndarray, B: np.ndarray) -> np.ndarray:
    return np.concatenate([X, B], axis=1)


def featurize_additive(X: np.ndarray, B: np.ndarray) -> np.ndarray:
    parts = [X, B]
    for k in range(K):
        parts.append(X * B[:, k:k + 1])
    return np.concatenate(parts, axis=1)


def fit_linear_additive(X, B, Y):
    F = featurize_additive(X, B)
    models = [LinearRegression().fit(F, Y[:, m]) for m in range(3)]

    def predict(Xq, Bq):
        Fq = featurize_additive(Xq, Bq)
        return np.stack([mdl.predict(Fq) for mdl in models], axis=1)

    return predict


def fit_buoplr_net(X, B, Y, seed):
    F = featurize_plain(X, B)
    # Outcomes live on very different scales (main ~1e-2, cost ~1, fatigue
    # ~0-1). A shared-trunk multi-output net trained on raw targets lets the
    # largest-scale outcome dominate the summed loss and starves the others
    # of gradient signal -- standardizing per-outcome is what makes the
    # cross-outcome sharing actually beneficial rather than just noise.
    y_mean, y_std = Y.mean(axis=0), Y.std(axis=0) + 1e-8
    Yz = (Y - y_mean) / y_std
    mlp = MLPRegressor(hidden_layer_sizes=(64, 64), activation="relu", max_iter=600,
                        random_state=seed, early_stopping=True, n_iter_no_change=20)
    mlp.fit(F, Yz)  # single shared-trunk network, multi-output -> cross-outcome sharing

    def predict(Xq, Bq):
        return mlp.predict(featurize_plain(Xq, Bq)) * y_std + y_mean

    return predict


def fit_independent_mlp(X, B, Y, seed):
    F = featurize_plain(X, B)
    models = []
    for m in range(3):
        mlp = MLPRegressor(hidden_layer_sizes=(64, 64), activation="relu", max_iter=400,
                            random_state=seed + m, early_stopping=True, n_iter_no_change=15)
        mlp.fit(F, Y[:, m])
        models.append(mlp)

    def predict(Xq, Bq):
        Fq = featurize_plain(Xq, Bq)
        return np.stack([mdl.predict(Fq) for mdl in models], axis=1)

    return predict


def predict_all_actions(predict_fn, X: np.ndarray) -> np.ndarray:
    """(n, N_ACTIONS, 3) predicted outcomes for every candidate bundle."""
    n = X.shape[0]
    out = np.empty((n, N_ACTIONS, 3))
    for a, bundle in enumerate(BUNDLES):
        Ba = np.tile(np.array(bundle, dtype=float), (n, 1))
        out[:, a, :] = predict_fn(X, Ba)
    return out


def to_uplift(pred_all: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    U = pred_all[:, :, 0] - pred_all[:, 0, 0:1]  # main uplift vs null
    C = pred_all[:, :, 1]
    Fat = pred_all[:, :, 2]
    return U, C, Fat


# --------------------------------------------------------------------------
# Stage 2: assignment algorithms
# --------------------------------------------------------------------------

def restrict_candidates(U: np.ndarray, n_candidates: int) -> np.ndarray:
    """Top-C actions per user by predicted main uplift, always including null."""
    order = np.argsort(-U, axis=1)[:, :n_candidates]
    has_null = (order == NULL_ACTION).any(axis=1)
    order[~has_null, -1] = NULL_ACTION
    return order  # (n, n_candidates) action indices


def assign_lagrangian(U, C, Fat, vendor, budget, quotas, fatigue_cap, n_candidates=3,
                       iters=60):
    """Restricted-decision-space Lagrangian relaxation (BUOPLR stage 2).

    Dual step sizes are self-calibrated from the uplift/cost/fatigue magnitudes
    so the method works regardless of the absolute scale of U (which is
    deliberately tiny here, mirroring the paper's "small-effect uplift"
    framing) — a fixed step size tuned for one scale silently collapses every
    user to the null action at another. We also track the best *feasible*
    primal iterate seen across the dual trajectory, since subgradient duals
    oscillate around the constraint boundary rather than converging exactly.
    """
    n = U.shape[0]
    cand = restrict_candidates(U, n_candidates)
    Uc = np.take_along_axis(U, cand, axis=1)
    Cc = np.take_along_axis(C, cand, axis=1)
    Fc = np.take_along_axis(Fat, cand, axis=1)
    is_active = (cand != NULL_ACTION).astype(float)

    u_scale = np.median(np.abs(Uc[Uc > 0])) if (Uc > 0).any() else 1e-3
    c_scale = max(np.median(Cc[Cc > 0]) if (Cc > 0).any() else 1.0, 1e-6)
    f_scale = max(np.median(Fc[Fc > 0]) if (Fc > 0).any() else 1.0, 1e-6)
    step_b0 = u_scale / c_scale
    step_q0 = u_scale
    step_f0 = u_scale / f_scale

    lam_b, lam_f = 0.0, 0.0
    lam_q = np.zeros(N_VENDORS)
    idx = np.arange(n)

    best_chosen, best_uplift = None, -np.inf
    chosen = np.zeros(n, dtype=int)
    for it in range(iters):
        decay = 1.0 / (1.0 + 0.15 * it)
        quota_pen = lam_q[vendor][:, None] * is_active
        score = Uc - lam_b * Cc - quota_pen - lam_f * Fc
        best_j = np.argmax(score, axis=1)
        chosen = np.take_along_axis(cand, best_j[:, None], axis=1)[:, 0]

        total_cost = C[idx, chosen].sum()
        active_mask = chosen != NULL_ACTION
        quota_usage = np.array([(active_mask & (vendor == v)).sum() for v in range(N_VENDORS)])
        avg_fatigue = Fat[idx, chosen].mean()

        feasible = (total_cost <= budget * 1.005 and (quota_usage <= quotas * 1.005).all()
                    and avg_fatigue <= fatigue_cap * 1.005)
        total_uplift = U[idx, chosen].sum()
        if feasible and total_uplift > best_uplift:
            best_uplift, best_chosen = total_uplift, chosen.copy()

        rel_budget = (total_cost - budget) / max(budget, 1e-6)
        rel_quota = (quota_usage - quotas) / np.maximum(quotas, 1e-6)
        rel_fatigue = (avg_fatigue - fatigue_cap) / max(fatigue_cap, 1e-6)
        lam_b = max(0.0, lam_b + decay * step_b0 * rel_budget)
        lam_q = np.maximum(0.0, lam_q + decay * step_q0 * rel_quota)
        lam_f = max(0.0, lam_f + decay * step_f0 * rel_fatigue)

    if best_chosen is None:
        best_chosen = chosen  # dual trajectory never found a feasible iterate; repair below
    chosen = _repair_feasibility(best_chosen, U, C, Fat, vendor, budget, quotas, fatigue_cap)
    return chosen


def _repair_feasibility(chosen, U, C, Fat, vendor, budget, quotas, fatigue_cap, max_steps=None):
    n = len(chosen)
    chosen = chosen.copy()
    max_steps = max_steps or n

    def violated():
        total_cost = C[np.arange(n), chosen].sum()
        active = chosen != NULL_ACTION
        quota_usage = np.array([(active & (vendor == v)).sum() for v in range(N_VENDORS)])
        avg_fatigue = Fat[np.arange(n), chosen].mean() if active.any() or True else 0.0
        return total_cost > budget * 1.01 or (quota_usage > quotas).any() or avg_fatigue > fatigue_cap * 1.01

    steps = 0
    active_idx = np.where(chosen != NULL_ACTION)[0]
    uplift_of_choice = U[np.arange(n), chosen]
    order = active_idx[np.argsort(uplift_of_choice[active_idx])]  # lowest uplift first
    ptr = 0
    while violated() and steps < max_steps and ptr < len(order):
        chosen[order[ptr]] = NULL_ACTION
        ptr += 1
        steps += 1
    return chosen


def assign_greedy(U, C, Fat, vendor, budget, quotas, fatigue_cap):
    n = U.shape[0]
    best_action = np.argmax(U, axis=1)
    best_uplift = U[np.arange(n), best_action]
    order = np.argsort(-best_uplift)
    chosen = np.zeros(n, dtype=int)
    remaining_budget = budget
    remaining_quota = quotas.copy().astype(float)
    fatigue_sum = 0.0
    for i in order:
        for cand_rank in range(N_ACTIONS):
            a = np.argsort(-U[i])[cand_rank]
            if a == NULL_ACTION:
                chosen[i] = NULL_ACTION
                break
            c = C[i, a]
            v = vendor[i]
            new_fatigue_avg = (fatigue_sum + Fat[i, a]) / (i + 1) if i > 0 else Fat[i, a]
            if c <= remaining_budget and remaining_quota[v] >= 1 and new_fatigue_avg <= fatigue_cap * 1.05:
                chosen[i] = a
                remaining_budget -= c
                remaining_quota[v] -= 1
                fatigue_sum += Fat[i, a]
                break
        else:
            chosen[i] = NULL_ACTION
        if chosen[i] == NULL_ACTION:
            fatigue_sum += Fat[i, NULL_ACTION]
    return chosen


def assign_lp(U, C, Fat, vendor, budget, quotas, fatigue_cap, n_candidates=4):
    from scipy.optimize import linprog
    from scipy.sparse import lil_matrix

    n = U.shape[0]
    cand = restrict_candidates(U, n_candidates)
    n_vars = n * n_candidates
    c_obj = -np.take_along_axis(U, cand, axis=1).reshape(-1)  # minimize -uplift

    rows_eq, cols_eq, data_eq, rhs_eq = [], [], [], []
    for i in range(n):
        for j in range(n_candidates):
            rows_eq.append(i)
            cols_eq.append(i * n_candidates + j)
            data_eq.append(1.0)
        rhs_eq.append(1.0)
    A_eq = lil_matrix((n, n_vars))
    for r, c_, v in zip(rows_eq, cols_eq, data_eq):
        A_eq[r, c_] = v

    Cc = np.take_along_axis(C, cand, axis=1).reshape(-1)
    Fc = np.take_along_axis(Fat, cand, axis=1).reshape(-1)
    active = (cand != NULL_ACTION).astype(float).reshape(-1)
    vendor_rep = np.repeat(vendor, n_candidates)

    A_ub_rows = [Cc]
    b_ub = [budget]
    for v in range(N_VENDORS):
        A_ub_rows.append(active * (vendor_rep == v).astype(float))
        b_ub.append(quotas[v])
    A_ub_rows.append(Fc)
    b_ub.append(fatigue_cap * n)
    A_ub = np.stack(A_ub_rows, axis=0)

    res = linprog(c_obj, A_ub=A_ub, b_ub=np.array(b_ub), A_eq=A_eq.toarray(), b_eq=np.array(rhs_eq),
                   bounds=(0, 1), method="highs")
    x = res.x.reshape(n, n_candidates)
    best_j = np.argmax(x, axis=1)
    chosen = np.take_along_axis(cand, best_j[:, None], axis=1)[:, 0]
    chosen = _repair_feasibility(chosen, U, C, Fat, vendor, budget, quotas, fatigue_cap)
    return chosen


# --------------------------------------------------------------------------
# Evaluation
# --------------------------------------------------------------------------

def evaluate_assignment(chosen, X, vendor, params, budget, quotas, fatigue_cap):
    true_all = true_outcomes_all_actions(X, params)  # (n, N_ACTIONS, 3)
    n = X.shape[0]
    idx = np.arange(n)
    main = true_all[idx, chosen, 0]
    cost = true_all[idx, chosen, 1]
    fatigue = true_all[idx, chosen, 2]
    null_main = true_all[idx, NULL_ACTION, 0]
    total_uplift = float((main - null_main).sum())
    total_cost = float(cost.sum())
    active = chosen != NULL_ACTION
    quota_usage = np.array([(active & (vendor == v)).sum() for v in range(N_VENDORS)])
    avg_fatigue = float(fatigue.mean())
    return dict(
        total_main_uplift=total_uplift,
        total_cost=total_cost,
        budget=budget,
        budget_ok=bool(total_cost <= budget * 1.02),
        quota_usage=quota_usage.tolist(),
        quotas=quotas.tolist(),
        quota_ok=bool((quota_usage <= quotas * 1.02).all()),
        avg_fatigue=avg_fatigue,
        fatigue_cap=fatigue_cap,
        fatigue_ok=bool(avg_fatigue <= fatigue_cap * 1.02),
        n_treated=int(active.sum()),
        n_users=n,
    )


def oracle_assignment(X, vendor, params, budget, quotas, fatigue_cap, n_candidates=4):
    true_all = true_outcomes_all_actions(X, params)
    U = true_all[:, :, 0] - true_all[:, 0, 0:1]
    C = true_all[:, :, 1]
    Fat = true_all[:, :, 2]
    return assign_lagrangian(U, C, Fat, vendor, budget, quotas, fatigue_cap, n_candidates=n_candidates, iters=60)


# --------------------------------------------------------------------------
# Main experiment
# --------------------------------------------------------------------------

def run(n_train: int, n_eval: int, seed: int, out_dir: Path, run_lp: bool, lp_max_n: int = 5000,
        scaling_ns: list[int] | None = None):
    out_dir.mkdir(parents=True, exist_ok=True)
    params = make_effect_params(seed)

    t0 = time.time()
    df_train, Xtr, Btr, Ytr = simulate_logged_data(n_train, seed, params)
    Xev, vendor_ev = sample_users(n_eval, seed + 999)
    print(f"[data] train={n_train} eval={n_eval} build_time={time.time() - t0:.1f}s")

    budget = 0.55 * n_eval * float(np.mean(params["unit_cost"]))
    quotas = np.array([0.55 * (vendor_ev == v).sum() for v in range(N_VENDORS)])
    fatigue_cap = 0.30

    uplift_models = {}
    fit_times = {}
    for name, fitter in [
        ("linear_additive", lambda: fit_linear_additive(Xtr, Btr, Ytr)),
        ("independent_mlp", lambda: fit_independent_mlp(Xtr, Btr, Ytr, seed)),
        ("buoplr_net", lambda: fit_buoplr_net(Xtr, Btr, Ytr, seed)),
    ]:
        t0 = time.time()
        predict_fn = fitter()
        fit_times[name] = time.time() - t0
        pred_all = predict_all_actions(predict_fn, Xev)
        U, C, Fat = to_uplift(pred_all)
        uplift_models[name] = dict(U=U, C=C, Fat=Fat)
        print(f"[fit] {name} fit_time={fit_times[name]:.1f}s")

    pipelines = {}

    # Control: no treatment
    chosen_ctrl = np.zeros(n_eval, dtype=int)
    pipelines["control_no_treatment"] = (chosen_ctrl, 0.0)

    # BUOPLR proposed pipeline: buoplr_net + lagrangian restricted assignment
    for uplift_name in uplift_models:
        m = uplift_models[uplift_name]
        t0 = time.time()
        chosen = assign_lagrangian(m["U"], m["C"], m["Fat"], vendor_ev, budget, quotas, fatigue_cap)
        rt = time.time() - t0
        pipelines[f"{uplift_name}__lagrangian"] = (chosen, rt)

        t0 = time.time()
        chosen_g = assign_greedy(m["U"], m["C"], m["Fat"], vendor_ev, budget, quotas, fatigue_cap)
        rt_g = time.time() - t0
        pipelines[f"{uplift_name}__greedy"] = (chosen_g, rt_g)

    # Oracle upper bound (true effects, Lagrangian assignment) on the full eval set
    t0 = time.time()
    chosen_oracle = oracle_assignment(Xev, vendor_ev, params, budget, quotas, fatigue_cap)
    rt_oracle = time.time() - t0
    pipelines["oracle__lagrangian"] = (chosen_oracle, rt_oracle)

    results = []
    for name, (chosen, rt) in pipelines.items():
        ev = evaluate_assignment(chosen, Xev, vendor_ev, params, budget, quotas, fatigue_cap)
        ev["pipeline"] = name
        ev["assignment_runtime_s"] = rt
        uplift_name = name.split("__")[0]
        ev["fit_time_s"] = fit_times.get(uplift_name, None)
        results.append(ev)

    res_df = pd.DataFrame(results).sort_values("total_main_uplift", ascending=False)
    res_df.to_csv(out_dir / "results.csv", index=False)
    print(res_df[["pipeline", "total_main_uplift", "budget_ok", "quota_ok", "fatigue_ok",
                   "assignment_runtime_s"]].to_string(index=False))

    # LP quality-gap comparison on a bounded subset (LP does not scale to the
    # full eval set -- that intractability is itself part of what we're
    # testing: does BUOPLR's restricted Lagrangian relaxation match near-exact
    # LP quality while staying orders of magnitude faster?).
    lp_rows = []
    if run_lp:
        n_lp = min(lp_max_n, n_eval)
        Xlp, vendor_lp = Xev[:n_lp], vendor_ev[:n_lp]
        budget_lp = 0.55 * n_lp * float(np.mean(params["unit_cost"]))
        quotas_lp = np.array([0.55 * (vendor_lp == v).sum() for v in range(N_VENDORS)])
        for uplift_name, m in uplift_models.items():
            Ulp, Clp, Flp = m["U"][:n_lp], m["C"][:n_lp], m["Fat"][:n_lp]
            t0 = time.time()
            chosen_lag = assign_lagrangian(Ulp, Clp, Flp, vendor_lp, budget_lp, quotas_lp, fatigue_cap)
            rt_lag = time.time() - t0
            t0 = time.time()
            chosen_lp = assign_lp(Ulp, Clp, Flp, vendor_lp, budget_lp, quotas_lp, fatigue_cap)
            rt_lp = time.time() - t0
            ev_lag = evaluate_assignment(chosen_lag, Xlp, vendor_lp, params, budget_lp, quotas_lp, fatigue_cap)
            ev_lp = evaluate_assignment(chosen_lp, Xlp, vendor_lp, params, budget_lp, quotas_lp, fatigue_cap)
            lp_rows.append(dict(uplift_model=uplift_name, n_users=n_lp,
                                 lagrangian_uplift=ev_lag["total_main_uplift"], lagrangian_runtime_s=rt_lag,
                                 lp_uplift=ev_lp["total_main_uplift"], lp_runtime_s=rt_lp,
                                 quality_gap_pct=100.0 * (ev_lp["total_main_uplift"] - ev_lag["total_main_uplift"])
                                 / max(abs(ev_lp["total_main_uplift"]), 1e-9),
                                 speedup_x=rt_lp / max(rt_lag, 1e-9)))
        lp_df = pd.DataFrame(lp_rows)
        lp_df.to_csv(out_dir / "lp_quality_gap.csv", index=False)
        print("\n[LP quality-gap @ n=%d]" % n_lp)
        print(lp_df.to_string(index=False))

    scaling_rows = []
    if scaling_ns:
        for n_s in scaling_ns:
            Xs, vendor_s = sample_users(n_s, seed + 12345 + n_s)
            budget_s = 0.55 * n_s * float(np.mean(params["unit_cost"]))
            quotas_s = np.array([0.55 * (vendor_s == v).sum() for v in range(N_VENDORS)])
            # reuse the already-fit buoplr_net predictor for inference + assignment timing
            t0 = time.time()
            pred_all_s = predict_all_actions(_LAST_BUOPLR_PREDICT[0], Xs)
            infer_time = time.time() - t0
            U_s, C_s, Fat_s = to_uplift(pred_all_s)
            t0 = time.time()
            assign_lagrangian(U_s, C_s, Fat_s, vendor_s, budget_s, quotas_s, fatigue_cap)
            lag_time = time.time() - t0
            t0 = time.time()
            assign_greedy(U_s, C_s, Fat_s, vendor_s, budget_s, quotas_s, fatigue_cap)
            greedy_time = time.time() - t0
            scaling_rows.append(dict(n_users=n_s, infer_time_s=infer_time,
                                      lagrangian_assign_s=lag_time, greedy_assign_s=greedy_time))
            print(f"[scaling] n={n_s} infer={infer_time:.2f}s lagrangian={lag_time:.2f}s greedy={greedy_time:.2f}s")
    scaling_df = pd.DataFrame(scaling_rows)
    if len(scaling_df):
        scaling_df.to_csv(out_dir / "scaling.csv", index=False)

    with open(out_dir / "config.json", "w") as f:
        json.dump(dict(n_train=n_train, n_eval=n_eval, seed=seed, budget=budget,
                        quotas=quotas.tolist(), fatigue_cap=fatigue_cap,
                        n_actions=N_ACTIONS, k_slots=K, max_active=MAX_ACTIVE,
                        fit_times=fit_times, scaling_ns=scaling_ns or []), f, indent=2)

    return res_df, scaling_df


def push_outputs_to_hub(out_dir: Path, repo_id: str):
    import os
    from huggingface_hub import HfApi

    token = os.environ.get("HF_TOKEN")
    api = HfApi(token=token)
    api.create_repo(repo_id, repo_type="dataset", exist_ok=True, private=False)
    for f in out_dir.iterdir():
        if f.is_file():
            api.upload_file(path_or_fileobj=str(f), path_in_repo=f.name,
                             repo_id=repo_id, repo_type="dataset")
    print(f"Pushed outputs to https://huggingface.co/datasets/{repo_id}")


_LAST_BUOPLR_PREDICT = [None]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-train", type=int, default=5000)
    ap.add_argument("--n-eval", type=int, default=2000)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out-dir", type=str, default="outputs/local")
    ap.add_argument("--run-lp", action="store_true")
    ap.add_argument("--lp-max-n", type=int, default=5000)
    ap.add_argument("--scaling-ns", type=str, default="")
    ap.add_argument("--push-repo", type=str, default="")
    args = ap.parse_args()

    scaling_ns = [int(x) for x in args.scaling_ns.split(",") if x.strip()] if args.scaling_ns else None

    # monkey-patch to capture the buoplr predictor for the scaling sweep
    orig_fit = fit_buoplr_net

    def fit_buoplr_net_capture(*a, **kw):
        fn = orig_fit(*a, **kw)
        _LAST_BUOPLR_PREDICT[0] = fn
        return fn

    globals()["fit_buoplr_net"] = fit_buoplr_net_capture

    out_dir = Path(args.out_dir)
    run(args.n_train, args.n_eval, args.seed, out_dir, args.run_lp, args.lp_max_n, scaling_ns)

    if args.push_repo:
        push_outputs_to_hub(out_dir, args.push_repo)


if __name__ == "__main__":
    main()