#!/usr/bin/env python3
"""
apub_paper_models.py -- the paper's OWN Section 5 / Appendix C application
models for "Minimizing Upper Confidence Bounds: A Data-Driven Framework for
Stochastic Programming" (arXiv 2403.08966), orid eXLcL70GXO.

Everything in this module is transcribed from the arXiv e-print LaTeX source
(Section5-NumericalAnalysis-ICML.tex, SectionC-Experiment_Parameters.tex,
SectionD-Sensitivity.tex, Section4-SolutionMethod-ICML.tex). No parameter is
invented.

Three application arms:

  A. Two-stage product mix with RANDOM recourse (Section 5, baseline instance
     |I| = 20 products, |J| = 8 departments), two-regime uniform marginals
     coupled by a Gumbel copula (lambda^r = 2.0 regular, lambda^w = 5.0
     worst-case, p = 0.9).  Appendix C gives c (20-vector) and T (8x20).

  B. Two-stage product mix with FIXED recourse (Section 5.3 / Appendix C.3),
     |I| = 4, |J| = 2, gamma ~ 0.7 N(mu_a, S_a) + 0.3 N(mu_b, S_b).  This is
     the instance the paper uses for its Wasserstein-DRO comparison.

  C. Multi-product newsvendor, 10 products (Appendix D + Appendix C.4),
     F(x, xi) = p'x + h'(x - xi)_+ + b'(xi - x)_+, p = -2, h = 9, b = 5,
     Case I  = Gaussian mixture N(mu1, Sigma1) / N(mu2, Sigma2),
     Case II = Case I + independent biased uniform noise.

Solution method: the deterministic equivalent of the paper's own bootstrap
approximation (Section 4, eq. mod:first-bs)

    min_{x in X, t}  c'x + t + 1/(alpha M) sum_m [ (1/N) sum_n V_mn Q(x, xi_n) - t ]_+

is written as ONE linear program and solved with HiGHS.  The paper solves the
same program with an adapted L-shaped (Benders) decomposition; L-shaped is an
algorithm for this LP, so solving the LP directly gives the same optimum.  We
solve the monolithic form because HiGHS handles these sizes directly.
"""
import numpy as np
from scipy import sparse
from scipy.optimize import linprog

# =============================================================================
# Appendix C -- deterministic parameters, random-recourse product mix
# =============================================================================

# C.1 item 1: unit cost (negative profit) for all 20 products
PM_C = np.array([-14, -9, -20, -15, -4, -40, -18, -11, -13, -16,
                 -17, -8, -9, -24, -10, -7, -12, -3, -4, -5], dtype=float)

# C.1 item 2: labor of department j required per unit of product i  (8 x 20)
PM_T = np.array([
    [10, 6, 8, 4, 10, 6, 8, 4, 6, 8, 4, 10, 7, 9, 12, 8, 11, 13, 16, 17],
    [6, 2, 3, 2, 6, 2, 3, 2, 3, 2, 6, 2, 5, 3, 7, 4, 6, 5, 8, 9],
    [10, 6, 8, 4, 10, 6, 8, 4, 8, 4, 10, 6, 7, 9, 12, 8, 11, 13, 16, 17],
    [6, 2, 3, 2, 6, 2, 3, 2, 2, 6, 2, 3, 5, 3, 7, 4, 6, 5, 8, 9],
    [0, 2, 3, 2, 2, 6, 2, 3, 0, 0, 0, 0, 1, 4, 0, 2, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 10, 6, 8, 4, 0, 0, 0, 0, 0, 0, 9, 0, 0],
    [0, 0, 0, 1, 4, 0, 2, 0, 0, 0, 0, 0, 3, 0, 0, 0, 5, 0, 0, 0],
    [6, 8, 4, 6, 0, 0, 0, 0, 0, 4, 6, 8, 4, 10, 7, 0, 0, 0, 0, 0],
], dtype=float)

# C.2 regular period (probability p = 0.9, Gumbel lambda^r = 2.0)
PM_P_REGULAR = 0.9
PM_LAMBDA_R = 2.0
PM_H_R = np.array([[8000.0, 8500.0], [10000.0, 120000.0]])          # h1, h2
PM_Q_R = np.array([[3, 5], [13, 16], [4, 7], [14, 17],
                   [15, 17], [4, 8], [15, 19], [18, 20]], dtype=float)
PM_W_R = np.array([[0.8, 1.0], [0.8, 1.0], [0.9, 1.0], [0.8, 1.0],
                   [0.85, 1.0], [0.85, 1.0], [0.9, 1.0], [0.9, 1.0]])

# C.2 worst-case period (Gumbel lambda^w = 5.0)
PM_LAMBDA_W = 5.0
PM_H_W = np.array([[2000.0, 3000.0], [10000.0, 120000.0]])
PM_Q_W = np.array([[9, 12], [21, 25], [10, 12], [22, 24],
                   [18, 20], [18, 21], [18, 20], [22, 25]], dtype=float)
PM_W_W = np.array([[0.5, 0.6], [0.5, 0.6], [0.6, 0.7], [0.4, 0.6],
                   [0.55, 0.65], [0.55, 0.65], [0.6, 0.7], [0.6, 0.7]])

# =============================================================================
# Appendix C.3 -- fixed-recourse product mix (the WassDRO comparison instance)
# =============================================================================
FR_C = np.array([-12.0, -20.0, -18.0, -40.0])
FR_QCOST = np.array([6.0, 12.0, 0.0, 0.0])
FR_W = np.array([[-0.9, 0.0, 1.0, 0.0],
                 [0.0, -0.9, 0.0, 1.0]])
FR_T_BASE = np.array([[4.0, 9.0, 7.0, 10.0],
                      [3.0, 1.0, 3.0, 6.0]])
FR_MIX_W = 0.7
FR_MU_A = np.array([12.0, 8.0])
FR_SIG_A = np.array([[5.76, 1.92], [1.92, 2.56]])
FR_MU_B = np.array([2.0, 1.0])
FR_SIG_B = np.array([[0.16, 0.04], [0.04, 0.04]])

# =============================================================================
# Appendix C.4 -- 10-product newsvendor
# =============================================================================
NV_P, NV_H, NV_B = -2.0, 9.0, 5.0
NV_MU1 = np.array([60.89, 48.58, 46.81, 56.54, 61.58,
                   52.69, 69.42, 60.54, 54.43, 51.76])
NV_MU2 = np.array([50.30, 61.87, 53.16, 41.79, 51.94,
                   62.14, 45.47, 45.26, 55.95, 55.95])
NV_SIGMA1 = np.array([
    [9.27, 2.84, -0.07, 1.19, -0.48, 1.40, 2.87, 4.06, -1.40, -1.96],
    [2.84, 5.90, -2.83, 0.21, 2.27, -2.40, -0.89, 4.22, 3.43, 2.78],
    [-0.07, -2.83, 5.48, -0.30, 0.90, 3.54, -4.51, -2.45, -2.91, -4.95],
    [1.19, 0.21, -0.30, 7.99, -1.02, -1.27, -0.15, -1.55, -1.69, -0.36],
    [-0.48, 2.27, 0.90, -1.02, 9.48, -0.08, -3.69, 2.71, -0.69, -0.34],
    [1.40, -2.40, 3.54, -1.27, -0.08, 6.94, -1.26, -2.73, 0.01, -5.19],
    [2.87, -0.89, -4.51, -0.15, -3.69, -1.26, 12.05, -0.16, -0.16, 2.44],
    [4.06, 4.22, -2.45, -1.55, 2.71, -2.73, -0.16, 9.16, -0.77, 1.94],
    [-1.40, 3.43, -2.91, -1.69, -0.69, 0.01, -0.16, -0.77, 7.41, 2.24],
    [-1.96, 2.78, -4.95, -0.36, -0.34, -5.19, 2.44, 1.94, 2.24, 6.70],
])
NV_SIGMA2 = np.array([
    [6.32, 2.99, -0.06, 0.73, -0.33, 1.36, 1.55, 2.51, -1.19, -1.75],
    [2.99, 9.57, -4.09, 0.19, 2.44, -3.60, -0.74, 4.02, 4.49, 3.83],
    [-0.06, -4.09, 7.06, -0.25, 0.86, 4.74, -3.35, -2.08, -3.40, -6.08],
    [0.73, 0.19, -0.25, 4.37, -0.64, -1.11, -0.07, -0.86, -1.29, -0.29],
    [-0.33, 2.44, 0.86, -0.64, 6.74, -0.08, -2.04, 1.71, -0.60, -0.31],
    [1.36, -3.60, 4.74, -1.11, -0.08, 9.65, -0.98, -2.41, 0.01, -6.62],
    [1.55, -0.74, -3.35, -0.07, -2.04, -0.98, 5.17, -0.08, -0.10, 1.72],
    [2.51, 4.02, -2.08, -0.86, 1.71, -2.41, -0.08, 5.12, -0.59, 1.57],
    [-1.19, 4.49, -3.40, -1.29, -0.60, 0.01, -0.10, -0.59, 7.83, 2.49],
    [-1.75, 3.83, -6.08, -0.29, -0.31, -6.62, 1.72, 1.57, 2.49, 7.83],
])
NV_EPS_LO = np.array([-5.37, 6.74, 3.22, -7.48, -4.89,
                      -0.21, -12.14, -7.74, 0.77, 2.13])
NV_EPS_HI = np.array([26.27, 14.16, 17.68, 28.38, 25.79,
                      16.11, 32.99, 28.64, 20.13, 18.77])


# =============================================================================
# Gumbel copula sampling (Marshall-Olkin with Kanter's positive-stable draw)
# =============================================================================
def positive_stable(a, size, rng):
    """Positive stable S with Laplace transform E[exp(-tS)] = exp(-t^a),
    a in (0,1).  Kanter (1975) / Chambers-Mallows-Stuck."""
    u = rng.uniform(0.0, np.pi, size=size)
    w = rng.exponential(1.0, size=size)
    return (np.sin(a * u) / np.sin(u) ** (1.0 / a)) * \
           (np.sin((1.0 - a) * u) / w) ** ((1.0 - a) / a)


def gumbel_copula_uniforms(n, d, lam, rng):
    """n x d matrix of U(0,1) marginals with Gumbel copula dependence,
    C(u; lam) = exp(-(sum_k (-log u_k)^lam)^(1/lam))."""
    a = 1.0 / lam
    s = positive_stable(a, size=(n, 1), rng=rng)
    e = rng.exponential(1.0, size=(n, d))
    return np.exp(-((e / s) ** (1.0 / lam)))


def sample_product_mix_xi(n, rng):
    """Section 5 DGP: two regimes (regular w.p. p = 0.9, worst-case otherwise);
    within a regime the 18 marginals (h1, h2, q1..q8, w1..w8) are uniform on
    the Appendix-C intervals and coupled by a Gumbel copula.
    Returns dict with h1 (n,), h2 (n,), q (n,8), w (n,8)."""
    regular = rng.random(n) < PM_P_REGULAR
    n_r, n_w = int(regular.sum()), int((~regular).sum())
    out_h = np.empty((n, 2))
    out_q = np.empty((n, 8))
    out_w = np.empty((n, 8))
    for mask, cnt, lam, hb, qb, wb in (
            (regular, n_r, PM_LAMBDA_R, PM_H_R, PM_Q_R, PM_W_R),
            (~regular, n_w, PM_LAMBDA_W, PM_H_W, PM_Q_W, PM_W_W)):
        if cnt == 0:
            continue
        u = gumbel_copula_uniforms(cnt, 18, lam, rng)
        out_h[mask] = hb[:, 0] + (hb[:, 1] - hb[:, 0]) * u[:, 0:2]
        out_q[mask] = qb[:, 0] + (qb[:, 1] - qb[:, 0]) * u[:, 2:10]
        out_w[mask] = wb[:, 0] + (wb[:, 1] - wb[:, 0]) * u[:, 10:18]
    return {"h1": out_h[:, 0], "h2": out_h[:, 1], "q": out_q, "w": out_w,
            "regular": regular}


# =============================================================================
# Arm A recourse:  Q(x, xi) = min_{y,z >= 0} q'y
#                  s.t. w_j y_j + z_j >= (Tx)_j,  sum_j z_j = h1,  sum_j y_j <= h2
# =============================================================================
def recourse_pm_closed_form(x, xi, h2_cap=True):
    """Vectorised exact Q(x, xi_n) for every scenario n.

    Free permanent labor (total h1) should be spent on the departments with the
    most expensive temp labor per unit of work, q_j / w_j; whatever demand is
    left is covered by temps.  This greedy allocation is exactly optimal (it is
    a continuous knapsack), and is unit-tested against linprog.

    The greedy ignores the sum_j y_j <= h2 cap; it is returned together with a
    `feasible` flag.  Where the flag is True the greedy solution is optimal for
    the capped problem as well (it is optimal for a relaxation and feasible for
    the original).  Appendix C sets h2 ~ U[10000, 120000] precisely to suppress
    that cap, and the flag is asserted in every reported evaluation."""
    d = PM_T @ np.asarray(x, dtype=float)            # (8,) department demand
    q, w, h1, h2 = xi["q"], xi["w"], xi["h1"], xi["h2"]
    rate = q / w                                     # (n,8) cost per unit work
    order = np.argsort(-rate, axis=1)                # expensive first
    d_sorted = np.take_along_axis(np.broadcast_to(d, rate.shape), order, axis=1)
    csum = np.cumsum(d_sorted, axis=1)
    remaining = np.maximum(h1[:, None] - (csum - d_sorted), 0.0)
    z_sorted = np.minimum(d_sorted, remaining)
    uncovered_sorted = d_sorted - z_sorted
    rate_sorted = np.take_along_axis(rate, order, axis=1)
    q_sorted = np.take_along_axis(q, order, axis=1)
    w_sorted = np.take_along_axis(w, order, axis=1)
    y_sorted = uncovered_sorted / w_sorted
    cost = (q_sorted * y_sorted).sum(axis=1)
    feasible = np.ones_like(cost, dtype=bool) if not h2_cap \
        else (y_sorted.sum(axis=1) <= h2 + 1e-9)
    return cost, feasible


def recourse_pm_lp(x, q, w, h1, h2):
    """Reference single-scenario recourse LP (used only in the unit tests)."""
    d = PM_T @ np.asarray(x, dtype=float)
    nJ = 8
    # variables [y (8), z (8)]
    c = np.concatenate([q, np.zeros(nJ)])
    rows, cols, vals = [], [], []
    for j in range(nJ):                                # -w_j y_j - z_j <= -d_j
        rows += [j, j]
        cols += [j, nJ + j]
        vals += [-w[j], -1.0]
    rows += [nJ] * nJ                                  # sum_j y_j <= h2
    cols += list(range(nJ))
    vals += [1.0] * nJ
    A_ub = sparse.coo_matrix((vals, (rows, cols)), shape=(nJ + 1, 2 * nJ))
    b_ub = np.concatenate([-d, [h2]])
    A_eq = sparse.coo_matrix(([1.0] * nJ, ([0] * nJ, list(range(nJ, 2 * nJ)))),
                             shape=(1, 2 * nJ))
    res = linprog(c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=np.array([h1]),
                  bounds=[(0, None)] * (2 * nJ), method="highs")
    return res.fun if res.status == 0 else np.inf


# =============================================================================
# Bootstrap multiplicities V (Section 4: V ~ Multinomial(N, 1/N))
# =============================================================================
def bootstrap_multiplicities(N, M, rng):
    return rng.multinomial(N, np.full(N, 1.0 / N), size=M).astype(float)


def apub_from_costs(theta, V, alpha):
    """Direct evaluation of the bootstrap APUB
       min_t t + 1/(alpha M) sum_m [ (1/N) sum_n V_mn theta_n - t ]_+
    which is the empirical CVaR at level alpha of the bootstrap means."""
    N = theta.shape[0]
    zm = (V @ theta) / N
    if alpha >= 1.0:
        return float(zm.mean())
    zs = np.sort(zm)
    M = zs.shape[0]
    k = alpha * M
    kf = int(np.floor(k))
    tail = zs[M - kf:].sum() if kf > 0 else 0.0
    frac = k - kf
    if frac > 1e-12 and M - kf - 1 >= 0:
        tail += frac * zs[M - kf - 1]
    return float(tail / k)


# =============================================================================
# Arm A -- APUB-SP / SAA deterministic-equivalent LP (20 products, 8 depts)
# =============================================================================
def solve_pm_random_recourse(xi, alpha, V=None, x_ub=None, h2_cap=True):
    """alpha = 1.0 (and V = None) gives the classical SAA model.

    h2_cap=False drops the sum_j y_j <= h2 row.  Appendix C says verbatim:
    "Without loss, we suppress the constraint on the capacity of total
    outsourced labor by setting a large value for h2", and with the cap active
    the second stage is genuinely infeasible for the low tail of
    h2 ~ U[10000, 120000] at the models' own optimal x, i.e. the program has no
    relatively complete recourse.  Reproducing the paper therefore means
    honouring that sentence; the rate at which the cap would have bound is
    reported alongside every result."""
    N = xi["h1"].shape[0]
    nI, nJ = PM_T.shape[1], PM_T.shape[0]
    M = 0 if V is None else V.shape[0]
    q, w, h1, h2 = xi["q"], xi["w"], xi["h1"], xi["h2"]

    off_x, off_y = 0, nI
    off_z = off_y + N * nJ
    off_th = off_z + N * nJ
    off_t = off_th + N
    off_u = off_t + 1
    nvar = off_u + M

    obj = np.zeros(nvar)
    obj[off_x:off_x + nI] = PM_C
    if M == 0:
        obj[off_th:off_th + N] = 1.0 / N
    else:
        obj[off_t] = 1.0
        obj[off_u:off_u + M] = 1.0 / (alpha * M)

    rows, cols, vals, b_ub = [], [], [], []
    r = 0
    # (1) T x - w_nj y_nj - z_nj <= 0
    Tc = sparse.coo_matrix(PM_T)
    for n in range(N):
        for j in range(nJ):
            for i in range(nI):
                if PM_T[j, i] != 0.0:
                    rows.append(r); cols.append(off_x + i); vals.append(PM_T[j, i])
            rows.append(r); cols.append(off_y + n * nJ + j); vals.append(-w[n, j])
            rows.append(r); cols.append(off_z + n * nJ + j); vals.append(-1.0)
            b_ub.append(0.0)
            r += 1
    # (2) sum_j y_nj <= h2_n   (suppressed per Appendix C when h2_cap=False)
    if h2_cap:
        for n in range(N):
            for j in range(nJ):
                rows.append(r); cols.append(off_y + n * nJ + j); vals.append(1.0)
            b_ub.append(h2[n]); r += 1
    # (3) (1/N) sum_n V_mn theta_n - u_m - t <= 0
    if M > 0:
        for m in range(M):
            nz = np.nonzero(V[m])[0]
            for n in nz:
                rows.append(r); cols.append(off_th + n); vals.append(V[m, n] / N)
            rows.append(r); cols.append(off_u + m); vals.append(-1.0)
            rows.append(r); cols.append(off_t); vals.append(-1.0)
            b_ub.append(0.0); r += 1
    A_ub = sparse.coo_matrix((vals, (rows, cols)), shape=(r, nvar)).tocsr()

    rows, cols, vals, b_eq = [], [], [], []
    re = 0
    for n in range(N):                                  # sum_j z_nj = h1_n
        for j in range(nJ):
            rows.append(re); cols.append(off_z + n * nJ + j); vals.append(1.0)
        b_eq.append(h1[n]); re += 1
    for n in range(N):                                  # theta_n - q_n' y_n = 0
        rows.append(re); cols.append(off_th + n); vals.append(1.0)
        for j in range(nJ):
            rows.append(re); cols.append(off_y + n * nJ + j); vals.append(-q[n, j])
        b_eq.append(0.0); re += 1
    A_eq = sparse.coo_matrix((vals, (rows, cols)), shape=(re, nvar)).tocsr()

    bounds = [(0.0, x_ub)] * nI + [(0.0, None)] * (2 * N * nJ) + \
             [(None, None)] * N + [(None, None)] + [(0.0, None)] * M
    res = linprog(obj, A_ub=A_ub, b_ub=np.array(b_ub), A_eq=A_eq,
                  b_eq=np.array(b_eq), bounds=bounds, method="highs")
    if res.status != 0:
        raise RuntimeError(f"pm LP failed: {res.message}")
    x = res.x[off_x:off_x + nI]
    return {"x": x, "obj": float(res.fun),
            "theta": res.x[off_th:off_th + N]}


# =============================================================================
# Arm B -- fixed-recourse product mix (4 products, 2 departments)
# =============================================================================
def sample_fixed_recourse_gamma(n, rng):
    comp = rng.random(n) < FR_MIX_W
    a = rng.multivariate_normal(FR_MU_A, FR_SIG_A, size=n)
    b = rng.multivariate_normal(FR_MU_B, FR_SIG_B, size=n)
    return np.where(comp[:, None], a, b)


def fr_recourse(x, gam):
    """Q(x, gamma) = min_y q'y s.t. W y = h - T(gamma) x, y >= 0, closed form.
    W = [[-0.9,0,1,0],[0,-0.9,0,1]], q = (6,12,0,0), h = 500*gamma,
    T(gamma)_ji = T_base_ji - gamma_j/4  =>  (Tx)_j - h_j = a_j(x) - gamma_j k(x)
    with a_j(x) = T_base[j] . x and k(x) = sum(x)/4 + 500."""
    x = np.asarray(x, dtype=float)
    a = FR_T_BASE @ x                                  # (2,)
    k = x.sum() / 4.0 + 500.0
    short = np.maximum(a[None, :] - gam * k, 0.0)      # (n,2)
    return (FR_QCOST[0] / 0.9) * short[:, 0] + (FR_QCOST[1] / 0.9) * short[:, 1]


def fr_lipschitz(x):
    """Lipschitz modulus of F(x, .) w.r.t. gamma in the dual (l-inf) norm of the
    l1 ground metric on the Wasserstein ball.  The subgradient wrt gamma of an
    active piece is (-(6/0.9)k, -(12/0.9)k), so Lip = (12/0.9) k(x) -- genuinely
    DECISION DEPENDENT, unlike the newsvendor case."""
    k = np.asarray(x, dtype=float).sum() / 4.0 + 500.0
    return (FR_QCOST[1] / 0.9) * k


def solve_fr(gam, alpha=1.0, V=None, dro_eps=None, x_ub=1e5):
    """SAA (alpha=1, V=None, dro_eps=None), APUB-M (V given), or exact
    Wasserstein-1 DRO (dro_eps given): worst-case expectation over a
    Wasserstein-1 ball equals SAA(x) + eps * Lip(x) for a cost that is a max of
    finitely many affine functions of xi on an unbounded support
    (Mohajerin Esfahani & Kuhn 2018, Cor. 5.1 / Remark 6.6)."""
    N = gam.shape[0]
    nI = 4
    M = 0 if V is None else V.shape[0]
    off_x = 0
    off_v = nI                       # v1_n, v2_n interleaved: off_v + 2n + j
    off_th = off_v + 2 * N
    off_t = off_th + N
    off_u = off_t + 1
    nvar = off_u + M

    obj = np.zeros(nvar)
    obj[off_x:off_x + nI] = FR_C
    if M == 0:
        obj[off_th:off_th + N] = 1.0 / N
        if dro_eps:
            # + eps * (12/0.9) * (sum(x)/4 + 500)   [affine in x]
            obj[off_x:off_x + nI] += dro_eps * (FR_QCOST[1] / 0.9) * 0.25
    else:
        obj[off_t] = 1.0
        obj[off_u:off_u + M] = 1.0 / (alpha * M)

    rows, cols, vals, b_ub = [], [], [], []
    r = 0
    for n in range(N):
        for j in range(2):
            # T_base[j] . x - gamma_nj * (sum(x)/4 + 500) - v_nj <= 0
            for i in range(nI):
                rows.append(r); cols.append(off_x + i)
                vals.append(FR_T_BASE[j, i] - gam[n, j] * 0.25)
            rows.append(r); cols.append(off_v + 2 * n + j); vals.append(-1.0)
            b_ub.append(500.0 * gam[n, j]); r += 1
    if M > 0:
        for m in range(M):
            nz = np.nonzero(V[m])[0]
            for n in nz:
                rows.append(r); cols.append(off_th + n); vals.append(V[m, n] / N)
            rows.append(r); cols.append(off_u + m); vals.append(-1.0)
            rows.append(r); cols.append(off_t); vals.append(-1.0)
            b_ub.append(0.0); r += 1
    A_ub = sparse.coo_matrix((vals, (rows, cols)), shape=(r, nvar)).tocsr()

    rows, cols, vals, b_eq = [], [], [], []
    for n in range(N):
        rows.append(n); cols.append(off_th + n); vals.append(1.0)
        rows.append(n); cols.append(off_v + 2 * n); vals.append(-FR_QCOST[0] / 0.9)
        rows.append(n); cols.append(off_v + 2 * n + 1); vals.append(-FR_QCOST[1] / 0.9)
        b_eq.append(0.0)
    A_eq = sparse.coo_matrix((vals, (rows, cols)), shape=(N, nvar)).tocsr()

    bounds = [(0.0, x_ub)] * nI + [(0.0, None)] * (2 * N) + \
             [(None, None)] * N + [(None, None)] + [(0.0, None)] * M
    res = linprog(obj, A_ub=A_ub, b_ub=np.array(b_ub), A_eq=A_eq,
                  b_eq=np.array(b_eq), bounds=bounds, method="highs")
    if res.status != 0:
        raise RuntimeError(f"fr LP failed: {res.message}")
    return {"x": res.x[off_x:off_x + nI], "obj": float(res.fun),
            "theta": res.x[off_th:off_th + N]}


# =============================================================================
# Arm C -- 10-product newsvendor
# =============================================================================
def sample_newsvendor(n, rng, case):
    comp = rng.random(n) < 0.5
    a = rng.multivariate_normal(NV_MU1, NV_SIGMA1, size=n)
    b = rng.multivariate_normal(NV_MU2, NV_SIGMA2, size=n)
    d = np.where(comp[:, None], a, b)
    if case == 2:
        d = d + rng.uniform(NV_EPS_LO, NV_EPS_HI, size=(n, 10))
    return d


def nv_cost(x, xi):
    x = np.asarray(x, dtype=float)
    return (NV_P * x).sum() + (NV_H * np.maximum(x - xi, 0.0)
                               + NV_B * np.maximum(xi - x, 0.0)).sum(axis=1)


def solve_nv(xi, alpha=1.0, V=None, dro_eps=None, x_ub=500.0):
    N, nI = xi.shape
    M = 0 if V is None else V.shape[0]
    off_x = 0
    off_s = nI                                  # s_ni at off_s + n*nI + i
    off_th = off_s + N * nI
    off_t = off_th + N
    off_u = off_t + 1
    nvar = off_u + M

    obj = np.zeros(nvar)
    const = 0.0
    if M == 0:
        obj[off_th:off_th + N] = 1.0 / N
        if dro_eps:
            # Lipschitz modulus of F(x, .) wrt xi in the dual l-inf norm is
            # max(h, b) = 9 for EVERY x -> constant shift, x_DRO = x_SAA.
            const = dro_eps * max(NV_H, NV_B)
    else:
        obj[off_t] = 1.0
        obj[off_u:off_u + M] = 1.0 / (alpha * M)

    rows, cols, vals, b_ub = [], [], [], []
    r = 0
    for n in range(N):
        for i in range(nI):
            rows += [r, r]; cols += [off_x + i, off_s + n * nI + i]
            vals += [NV_H, -1.0]; b_ub.append(NV_H * xi[n, i]); r += 1
            rows += [r, r]; cols += [off_x + i, off_s + n * nI + i]
            vals += [-NV_B, -1.0]; b_ub.append(-NV_B * xi[n, i]); r += 1
    if M > 0:
        for m in range(M):
            nz = np.nonzero(V[m])[0]
            for n in nz:
                rows.append(r); cols.append(off_th + n); vals.append(V[m, n] / N)
            rows.append(r); cols.append(off_u + m); vals.append(-1.0)
            rows.append(r); cols.append(off_t); vals.append(-1.0)
            b_ub.append(0.0); r += 1
    A_ub = sparse.coo_matrix((vals, (rows, cols)), shape=(r, nvar)).tocsr()

    rows, cols, vals, b_eq = [], [], [], []
    for n in range(N):
        rows.append(n); cols.append(off_th + n); vals.append(1.0)
        for i in range(nI):
            rows.append(n); cols.append(off_x + i); vals.append(-NV_P)
            rows.append(n); cols.append(off_s + n * nI + i); vals.append(-1.0)
        b_eq.append(0.0)
    A_eq = sparse.coo_matrix((vals, (rows, cols)), shape=(N, nvar)).tocsr()

    bounds = [(0.0, x_ub)] * nI + [(0.0, None)] * (N * nI) + \
             [(None, None)] * N + [(None, None)] + [(0.0, None)] * M
    res = linprog(obj, A_ub=A_ub, b_ub=np.array(b_ub), A_eq=A_eq,
                  b_eq=np.array(b_eq), bounds=bounds, method="highs")
    if res.status != 0:
        raise RuntimeError(f"nv LP failed: {res.message}")
    return {"x": res.x[off_x:off_x + nI], "obj": float(res.fun) + const,
            "theta": res.x[off_th:off_th + N]}