Stochastic Programming

# Imports
import sys

if "google.colab" in sys.modules:
    !wget "https://raw.githubusercontent.com/ndcbe/optimization/main/notebooks/helper.py"
    import helper

    helper.easy_install()
else:
    sys.path.insert(0, "../")
    import helper
helper.set_plotting_style()

import numpy as np
import pyomo.environ as pyo
from pyomo.environ import units as u
from pyomo.util.check_units import assert_units_consistent

u.load_definitions_from_strings(["USD = [currency]"])
USD_PER_ACRE = u.USD / u.acre
USD_PER_TON = u.USD / u.metric_ton
TON_PER_ACRE = u.metric_ton / u.acre
CROPS = ["WHEAT", "CORN", "BEETS"]
FEED_CROPS = ["WHEAT", "CORN"]
SALE_TYPES = ["WHEAT", "CORN", "BEETS_FAVORABLE", "BEETS_UNFAVORABLE"]
MEAN_YIELD = dict(zip(CROPS, [2.5, 3.0, 20.0]))  # T/acre
PLANT_COST = dict(zip(CROPS, [150.0, 230.0, 260.0]))  # $/acre
BUY_PRICE = dict(zip(FEED_CROPS, [238.0, 210.0]))  # $/T
SELL_PRICE = dict(zip(SALE_TYPES, [170.0, 150.0, 36.0, 10.0]))  # $/T
FEED_NEED = dict(zip(FEED_CROPS, [200.0, 240.0]))  # T
TOTAL_LAND, BEET_QUOTA = 500.0, 6000.0  # acres, T

def build_farmer(
    scenarios, weights=None, plant_cost=None, risk=None, check_units=False
):
    """Return a fresh farmer model; scenarios contain yields and optional prices.

    scenarios: nonempty list of dictionaries; 'yield' maps crops to T/acre.
    Optional 'buy' and 'sell' mappings give $/T. weights sum to one.
    plant_cost maps crops to $/acre. risk selects the optional Lecture 10 policy.
    """
    n = len(scenarios)
    if n == 0:
        raise ValueError("At least one scenario is required")
    w = np.full(n, 1 / n) if weights is None else np.asarray(weights, float)
    if w.shape != (n,) or not np.all(np.isfinite(w)) or np.any(w <= 0):
        raise ValueError("Retained scenarios need positive finite probabilities")
    if not np.isclose(w.sum(), 1.0):
        raise ValueError("Scenario probabilities must sum to one")
    costs = PLANT_COST if plant_cost is None else plant_cost
    m = pyo.ConcreteModel()
    m.CROPS = pyo.Set(initialize=CROPS)
    m.FEED_CROPS = pyo.Set(initialize=FEED_CROPS)
    m.SALE_TYPES = pyo.Set(initialize=SALE_TYPES)
    m.SCENARIOS = pyo.RangeSet(0, n - 1)
    m.weight = pyo.Param(m.SCENARIOS, initialize=dict(enumerate(w)))
    # Plant once, before learning the yields: no scenario index on acreage.
    m.acreage = pyo.Var(m.CROPS, domain=pyo.NonNegativeReals, units=u.acre)
    m.land_limit = pyo.Constraint(expr=sum(m.acreage.values()) <= TOTAL_LAND * u.acre)
    m.planting_cost = pyo.Expression(
        expr=sum(costs[c] * USD_PER_ACRE * m.acreage[c] for c in CROPS)
    )
    # The helper below attaches recourse to each Block, sharing root acreage.
    m.scenarios = pyo.Block(
        m.SCENARIOS, rule=lambda b, s: add_farmer_scenario(b, m, scenarios[s])
    )
    m.expected_profit = pyo.Expression(
        expr=sum(m.weight[s] * m.scenarios[s].profit for s in m.SCENARIOS)
    )
    m.objective = pyo.Objective(expr=m.expected_profit, sense=pyo.maximize)
    if risk:
        add_farmer_risk(m, risk)  # Optional policies introduced in Lecture 10.
    if check_units:
        assert_units_consistent(m)
    return m

def add_farmer_scenario(b, m, scenario):
    """Attach one scenario's recourse, using the root model's acreage."""
    crop_yield = scenario["yield"]
    buy = scenario.get("buy", BUY_PRICE)
    sell = scenario.get("sell", SELL_PRICE)
    b.purchases = pyo.Var(m.FEED_CROPS, domain=pyo.NonNegativeReals, units=u.metric_ton)
    b.sales = pyo.Var(m.SALE_TYPES, domain=pyo.NonNegativeReals, units=u.metric_ton)
    b.sales["BEETS_FAVORABLE"].setub(BEET_QUOTA)

    # Harvest + purchases - sales must cover animal feed in this scenario.
    @b.Constraint(m.FEED_CROPS)
    def feed_balance(block, c):
        return (
            crop_yield[c] * TON_PER_ACRE * m.acreage[c]
            + block.purchases[c]
            - block.sales[c]
            >= FEED_NEED[c] * u.metric_ton
        )

    b.beets_constraint = pyo.Constraint(
        expr=b.sales["BEETS_FAVORABLE"] + b.sales["BEETS_UNFAVORABLE"]
        <= crop_yield["BEETS"] * TON_PER_ACRE * m.acreage["BEETS"]
    )
    # Profit includes planting cost; weights summing to one count it once.
    b.profit = pyo.Expression(
        expr=sum(sell[k] * USD_PER_TON * b.sales[k] for k in SALE_TYPES)
        - sum(buy[c] * USD_PER_TON * b.purchases[c] for c in FEED_CROPS)
        - m.planting_cost
    )

farmer_solver = pyo.SolverFactory("appsi_highs")
assert farmer_solver.available(), "HiGHS is required for the rest of this notebook"


def solve_farmer(m):
    """Solve, refusing anything but a proven optimum."""
    results = farmer_solver.solve(m)
    assert pyo.check_optimal_termination(
        results
    ), f"Solve failed: termination={results.solver.termination_condition}"
    return m


def plan_of(m):
    """Stage 1 acreage as a plain array, in the order of CROPS."""
    return np.array([pyo.value(m.acreage[c]) for c in CROPS])

def grade_plan(plan, scenarios, weights=None, plant_cost=None):
    """Fix acreage in CROPS order; optimize trades for the supplied scenarios.

    Return expected profit and an array of scenario profits, both in dollars.
    """
    m = build_farmer(scenarios, weights, plant_cost)
    for crop, acres in zip(CROPS, plan):
        m.acreage[crop].fix(acres)
    solve_farmer(m)
    return (
        pyo.value(m.expected_profit),
        np.array([pyo.value(m.scenarios[i].profit) for i in m.SCENARIOS]),
    )

def correlated_scenario(multiplier):
    """One scenario of Section 1.1b: every crop scaled by the same factor."""
    return {"yield": {c: MEAN_YIELD[c] * multiplier for c in CROPS}}


# The three equally likely scenarios of Section 1.1b
THREE = [correlated_scenario(f) for f in (1.2, 1.0, 0.8)]

mean_model = build_farmer([correlated_scenario(1.0)], check_units=True)
mean_model.dual = pyo.Suffix(direction=pyo.Suffix.IMPORT)
solve_farmer(mean_model)
mean_plan = plan_of(mean_model)
print("Mean-yield acreage [wheat, corn, beets]:", mean_plan)
print(f"Profit: ${pyo.value(mean_model.expected_profit):,.2f}")
print(f"Land multiplier: ${mean_model.dual[mean_model.land_limit]:.2f}/acre")
for c in FEED_CROPS:
    con = mean_model.scenarios[0].feed_balance[c]
    print(f"Feed minimum multiplier, {c}: ${mean_model.dual[con]:.2f}/T")
assert np.allclose(mean_plan, [120, 80, 300])
assert np.isclose(pyo.value(mean_model.expected_profit), 118600)

perfect_models = [solve_farmer(build_farmer([scenario])) for scenario in THREE]

print("Known yield     wheat     corn    beets       profit ($)")
for label, model in zip(["above", "average", "below"], perfect_models):
    a = plan_of(model)
    print(
        f"{label:12s} {a[0]:8.1f} {a[1]:8.1f} {a[2]:8.1f}"
        f" {pyo.value(model.expected_profit):16,.2f}"
    )

# Validation first: reproduce Birge & Louveaux Table 5, p. 8.
RP = solve_farmer(build_farmer(THREE))
rp_plan = plan_of(RP)
rp_expected = pyo.value(RP.expected_profit)
rp_profits = np.array([pyo.value(RP.scenarios[i].profit) for i in RP.SCENARIOS])

print(
    f"acres  wheat {rp_plan[0]:.1f}  corn {rp_plan[1]:.1f}  beets {rp_plan[2]:.1f}"
    "   (book: 170, 80, 250)"
)
print(f"expected profit = ${rp_expected:,.0f}   (book: $108,390)")
print(
    "scenario profits  above ${:,.0f} | average ${:,.0f} | below ${:,.0f}".format(
        *rp_profits
    )
)

assert np.allclose(rp_plan, [170.0, 80.0, 250.0], atol=1e-6)
assert abs(rp_expected - 108390.0) < 1.0

# Exercise the unit declarations once, on this small model.
assert_units_consistent(build_farmer(THREE, check_units=True))
print("units consistent")

print("Scenario    wheat sold   corn bought   corn sold   premium beets   profit ($)")
for label, s in zip(["above", "average", "below"], RP.SCENARIOS):
    b = RP.scenarios[s]
    print(
        f"{label:8s} {pyo.value(b.sales['WHEAT']):13.1f}"
        f" {pyo.value(b.purchases['CORN']):13.1f} {pyo.value(b.sales['CORN']):11.1f}"
        f" {pyo.value(b.sales['BEETS_FAVORABLE']):15.1f} {pyo.value(b.profit):12,.0f}"
    )

WS = np.mean([pyo.value(m.expected_profit) for m in perfect_models])
EEV, ev_profits = grade_plan(mean_plan, THREE)

print("Policy                         expected profit ($)")
print(f"Perfect information (WS)        {WS:17,.2f}")
print(f"Shared stochastic plan (RP)     {rp_expected:17,.2f}")
print(f"Mean-yield plan evaluated (EEV) {EEV:17,.2f}")
print(f"EVPI = WS - RP: ${WS - rp_expected:,.2f}")
print(f"VSS  = RP - EEV: ${rp_expected - EEV:,.2f}")
print("EEV scenario profits [above, average, below]:", ev_profits)
assert np.isclose(EEV, 107240)
assert WS >= rp_expected >= EEV

def sample_farmer_yields(n, seed):
    """Independent uniform yields, T/acre; Birge & Louveaux, Section 1.1d."""
    rng = np.random.default_rng(seed)
    return [
        {"yield": {c: rng.uniform(0.8, 1.2) * MEAN_YIELD[c] for c in CROPS}}
        for _ in range(n)
    ]


def build_sampled_farmer(n, seed):
    """Build an n-scenario SAA model with reproducible, equally weighted draws."""
    # Equal weights 1/n are the builder's default.
    return build_farmer(sample_farmer_yields(n, seed))

evaluation_sample = sample_farmer_yields(2000, seed=20260922)
trained = solve_farmer(build_sampled_farmer(200, seed=12))
evaluated, profits = grade_plan(plan_of(trained), evaluation_sample)
se = profits.std(ddof=1) / np.sqrt(len(profits))

print("N train   wheat    corn   beets   train mean   eval mean   eval SE ($)")
for n_train, seed in [(50, 11), (200, 12), (1000, 13)]:
    trained = solve_farmer(build_sampled_farmer(n_train, seed))
    plan = plan_of(trained)
    evaluated, profits = grade_plan(plan, evaluation_sample)
    se = profits.std(ddof=1) / np.sqrt(len(profits))
    print(
        f"{n_train:7d} {plan[0]:7.1f} {plan[1]:7.1f} {plan[2]:7.1f}"
        f" {pyo.value(trained.expected_profit):12,.0f} {evaluated:11,.0f} {se:13,.0f}"
    )

def add_farmer_cvar(m, risk):
    """Lower-tail profit CVaR: alpha is the confidence level."""
    w = m.weight
    n = len(m.SCENARIOS)
    # Rockafellar-Uryasev epigraph, B&L (9.8)-(9.10), p. 85, written for the
    # LOWER tail of profit: CVaR_a = max_nu { nu - E[(nu - profit)_+]/(1-a) }
    alpha = (risk.get("cvar") or risk.get("cvar_floor"))[0]
    m.nu = pyo.Var(domain=pyo.Reals, units=u.USD)
    m.shortfall = pyo.Var(m.SCENARIOS, domain=pyo.NonNegativeReals, units=u.USD)
    m.shortfall_con = pyo.Constraint(
        m.SCENARIOS,
        rule=lambda m, i: m.shortfall[i] >= m.nu - m.scenarios[i].profit,
    )
    m.cvar = m.nu - (1.0 / (1.0 - alpha)) * sum(w[i] * m.shortfall[i] for i in range(n))
    if "cvar_floor" in risk:
        m.cvar_con = pyo.Constraint(expr=m.cvar >= risk["cvar_floor"][1] * u.USD)

def add_farmer_risk(m, risk):
    """Attach the requested risk policy to a fresh farmer model."""
    w = m.weight
    n = len(m.SCENARIOS)
    m.del_component(m.objective)
    if "floor" in risk:
        # Exercise 7(b): a profit floor imposed in every scenario
        m.profit_floor = pyo.Constraint(
            m.SCENARIOS,
            rule=lambda m, i: m.scenarios[i].profit >= risk["floor"] * u.USD,
        )
    if "minimax" in risk:
        # B&L (9.15), p. 86, in profit form: maximize the worst scenario
        m.worst = pyo.Var(domain=pyo.Reals, units=u.USD)
        m.worst_con = pyo.Constraint(
            m.SCENARIOS, rule=lambda m, i: m.worst <= m.scenarios[i].profit
        )
    if "cvar" in risk or "cvar_floor" in risk:
        add_farmer_cvar(m, risk)
    if "downside" in risk:
        # Expected downside risk, B&L (5.1)-(5.3), p. 68
        target, level = risk["downside"]
        m.shortfall_below_target = pyo.Var(
            m.SCENARIOS, domain=pyo.NonNegativeReals, units=u.USD
        )
        m.downside_def = pyo.Constraint(
            m.SCENARIOS,
            rule=lambda m, i: m.shortfall_below_target[i]
            >= target * u.USD - m.scenarios[i].profit,
        )
        m.downside_limit = pyo.Constraint(
            expr=sum(w[i] * m.shortfall_below_target[i] for i in range(n))
            <= level * u.USD
        )

    if "minimax" in risk:
        m.objective = pyo.Objective(expr=m.worst, sense=pyo.maximize)
    elif "cvar" in risk:
        lam = risk["cvar"][1]
        m.objective = pyo.Objective(
            expr=(1 - lam) * m.expected_profit + lam * m.cvar, sense=pyo.maximize
        )
    else:
        m.objective = pyo.Objective(expr=m.expected_profit, sense=pyo.maximize)

# Two samplers, seeded. The only difference is whether one draw is shared by all
# three crops (an instructor extension) or each crop is drawn on its own
# (Section 1.1d's independence). Both give each crop the SAME marginal:
# Uniform(0.8 * mean, 1.2 * mean).
YIELD_BOUNDS = {c: (0.8 * MEAN_YIELD[c], 1.2 * MEAN_YIELD[c]) for c in CROPS}


def sample_independent(n, seed):
    """Section 1.1d, p. 11: each crop's yield drawn independently."""
    rng = np.random.default_rng(seed)
    return [
        {"yield": {c: rng.uniform(*YIELD_BOUNDS[c]) for c in CROPS}} for _ in range(n)
    ]


def sample_correlated(n, seed):
    """Instructor extension: one common continuous weather factor."""
    rng = np.random.default_rng(seed)
    return [correlated_scenario(f) for f in rng.uniform(0.8, 1.2, size=n)]


# --- the solver-free reference: B&L's closed-form expected recourse, pp. 13-14
def Q_wheat(x):
    if x <= 200 / 3:
        return 47600 - 595 * x
    if x <= 100:
        return 119 * (200 - 2 * x) ** 2 / x - 85 * (200 - 3 * x) ** 2 / x
    return 34000 - 425 * x


def Q_corn(x):
    if x <= 200 / 3:
        return 50400 - 630 * x
    if x <= 100:
        return 87.5 * (240 - 2.4 * x) ** 2 / x - 62.5 * (240 - 3.6 * x) ** 2 / x
    return 36000 - 450 * x


def Q_beets(x):
    """Equations (1.10)-(1.12), pp. 12-13, with l3 = 16 and u3 = 24 T/acre."""
    lo, hi = 16.0, 24.0
    if x <= 0:
        return 0.0
    if lo * x > 6000:  # quota exceeded for every yield
        return -156000 - 10 * MEAN_YIELD["BEETS"] * x
    if hi * x < 6000:  # quota never reached
        return -36 * MEAN_YIELD["BEETS"] * x
    return -36 * MEAN_YIELD["BEETS"] * x + 13 * (hi * x - 6000) ** 2 / (x * (hi - lo))


def analytic_expected_profit(plan):
    """Exact E[profit] of a plan under continuous uniform yields -- no solver."""
    planting = sum(PLANT_COST[c] * a for c, a in zip(CROPS, plan))
    return -(planting + Q_wheat(plan[0]) + Q_corn(plan[1]) + Q_beets(plan[2]))


book_plan = np.array([135.83, 85.07, 279.10])
print(
    f"B&L KKT solution, p. 15:  ({book_plan[0]}, {book_plan[1]}, {book_plan[2]})"
    f"   analytic E[profit] = ${analytic_expected_profit(book_plan):,.2f}"
)
print()
print(
    f"{'sampler':>12s} {'N':>6s} {'wheat':>8s} {'corn':>8s} {'beets':>8s}"
    f" {'sampled E':>12s} {'analytic E':>12s}"
)
for sampler, label in [
    (sample_independent, "independent"),
    (sample_correlated, "correlated"),
]:
    for n, seed in [(200, 1), (1000, 2), (4000, 3)]:
        m = solve_farmer(build_farmer(sampler(n, seed)))
        p = plan_of(m)
        print(
            f"{label:>12s} {n:6d} {p[0]:8.2f} {p[1]:8.2f} {p[2]:8.2f}"
            f" {pyo.value(m.expected_profit):12,.0f} {analytic_expected_profit(p):12,.0f}"
        )

# Hold the plan FIXED at the book's solution and look at the distribution of
# profit it produces under each assumption. Same decision, same mean -- the only
# difference is whether the three crops fail together.
N_EVAL, EVAL_SEED = 20_000, 20260921


def tail_summary(profits, alpha):
    """Empirical mean, spread and lower-tail risk of a profit sample."""
    s = np.sort(profits)
    mass = (1 - alpha) * len(s)
    k = min(len(s), int(np.floor(mass + 1e-10)))
    fractional = max(0.0, mass - k)
    tail_mean = (s[:k].sum() + (fractional * s[k] if k < len(s) else 0)) / mass
    cutoff = max(0, int(np.ceil(mass - 1e-10)) - 1)
    return {
        "mean": s.mean(),
        "sd": s.std(ddof=1),
        "min": s[0],
        f"VaR_{alpha}": s[cutoff],
        f"CVaR_{alpha}": tail_mean,
    }


print(
    f"Plan held fixed at ({book_plan[0]}, {book_plan[1]}, {book_plan[2]}) acres;"
    f" {N_EVAL:,} seeded draws each."
)
print(
    f"True E[profit] for BOTH, analytically: ${analytic_expected_profit(book_plan):,.2f}"
)
print()
print(
    f"{'assumption':>22s} {'E (sampled)':>13s} {'std dev':>10s} {'min':>10s}"
    f" {'VaR_0.95':>10s} {'CVaR_0.95':>10s}"
)
summaries = {}
for label, sampler in [
    ("independent (1.1d)", sample_independent),
    ("correlated uniform", sample_correlated),
]:
    _, profits = grade_plan(book_plan, sampler(N_EVAL, EVAL_SEED))
    s = tail_summary(profits, 0.95)
    summaries[label] = s
    print(
        f"{label:>22s} {s['mean']:13,.0f} {s['sd']:10,.0f} {s['min']:10,.0f}"
        f" {s['VaR_0.95']:10,.0f} {s['CVaR_0.95']:10,.0f}"
    )

gap = (
    summaries["independent (1.1d)"]["CVaR_0.95"]
    - summaries["correlated uniform"]["CVaR_0.95"]
)
print(f"\nSame plan, same true mean, but CVaR_0.95 differs by ${gap:,.0f}.")
print("The two sampled means differ only by Monte Carlo noise; the spread does not.")

# Exercise 7(a): grade the worst-case plan of Table 4, p. 7. Recourse is still
# optimized in every scenario -- only the acreage is fixed.
table_4_plan = np.array([100.0, 25.0, 375.0])
e_worst_case, profits_worst_case = grade_plan(table_4_plan, THREE)

print("Table 4 plan (plan for the worst case): 100 / 25 / 375 acres")
print(
    "  scenario profits  above ${:,.0f} | average ${:,.0f} | below ${:,.0f}".format(
        *profits_worst_case
    )
)
print(f"  expected profit = ${e_worst_case:,.0f}")
print(
    f"\nExercise 7(a): loss in expected profit = "
    f"${rp_expected:,.0f} - ${e_worst_case:,.0f} = ${rp_expected - e_worst_case:,.0f}"
)

# The below-average profit is printed in the book: Table 4, p. 7, gives $59,950.
assert abs(profits_worst_case.min() - 59950.0) < 1.0

protected = solve_farmer(build_farmer(THREE, risk={"floor": 56_000}))
protected_mean, protected_profits = grade_plan(plan_of(protected), THREE)

# Exercises 7(b) and 7(c): raise the floor and watch what the insurance costs.
print(
    f"{'floor':>10s} {'wheat':>8s} {'corn':>8s} {'beets':>8s}"
    f" {'E[profit]':>12s} {'loss vs RP':>12s}"
)
frontier = []
for floor in [48_000, 50_000, 52_000, 54_000, 56_000, 58_000]:
    m = solve_farmer(build_farmer(THREE, risk={"floor": floor}))
    p, e = plan_of(m), pyo.value(m.expected_profit)
    frontier.append((floor, p, e))
    print(
        f"{floor:10,d} {p[0]:8.1f} {p[1]:8.1f} {p[2]:8.1f} {e:12,.0f}"
        f" {rp_expected - e:12,.0f}"
    )

print("\nMarginal cost of each extra $2,000 of floor:")
for (f0, _, e0), (f1, _, e1) in zip(frontier, frontier[1:]):
    print(f"  ${f0:,} -> ${f1:,}:  ${e0 - e1:>8,.0f}")

# Verify the equivalence rather than asserting it: solve BOTH models and diff.
print(
    f"{'alpha':>8s} {'floor':>9s} {'E (floor model)':>17s} {'E (CVaR model)':>16s}"
    f" {'dE':>8s} {'max|d acres|':>13s}"
)
for alpha in [2 / 3, 0.75, 0.90, 0.95]:
    for floor, plan_floor, e_floor in frontier:
        if floor < 52_000:
            continue
        m = solve_farmer(build_farmer(THREE, risk={"cvar_floor": (alpha, floor)}))
        e_cvar, plan_cvar = pyo.value(m.expected_profit), plan_of(m)
        d_acres = np.abs(plan_cvar - plan_floor).max()
        print(
            f"{alpha:8.3f} {floor:9,d} {e_floor:17,.2f} {e_cvar:16,.2f}"
            f" {e_cvar - e_floor:8.2f} {d_acres:13.2e}"
        )
        assert abs(e_cvar - e_floor) < 1e-4 and d_acres < 1e-6

# Hand arithmetic first, then the LP epigraph, then check they agree.
sorted_profits = np.sort(rp_profits)
var_50 = sorted_profits[1]  # smallest t with P(profit <= t) >= 0.5
cvar_50 = (1 / 3 * sorted_profits[0] + 1 / 6 * sorted_profits[1]) / 0.5
conditional_mean = sorted_profits[sorted_profits <= var_50].mean()

print(
    "profits at the risk-neutral plan: "
    + " | ".join(f"${p:,.0f}" for p in sorted_profits)
)
print(f"  VaR_0.5  (lower tail)          = ${var_50:,.2f}")
print(
    f"  CVaR_0.5 = (1/3 x {sorted_profits[0]:,.0f} + 1/6 x {sorted_profits[1]:,.0f}) / 0.5"
    f" = ${cvar_50:,.2f}"
)
print(
    f"  E[profit | profit <= VaR_0.5]  = ${conditional_mean:,.2f}"
    "   <- averages 2/3 of the mass, not 1/2"
)

# Second route: the Rockafellar-Uryasev epigraph solved at the same fixed plan.
m = build_farmer(THREE, risk={"cvar": (0.5, 1.0)})
for crop, acres in zip(CROPS, rp_plan):
    m.acreage[crop].fix(acres)
solve_farmer(m)
print(f"\n  LP epigraph at the same plan   = ${pyo.value(m.cvar):,.2f}")
assert abs(pyo.value(m.cvar) - cvar_50) < 1e-6

# The limits, verified rather than asserted. Read the PLAN off the risk solve,
# then grade it with a separate expectation solve (the next section explains why).
print(
    f"{'alpha':>8s} {'CVaR*':>12s} {'wheat':>7s} {'corn':>7s} {'beets':>7s}"
    f" {'E[profit]':>12s} {'worst':>10s}"
)
for alpha in [0.0, 0.1, 1 / 3, 0.5, 2 / 3, 0.9, 0.99]:
    m = solve_farmer(build_farmer(THREE, risk={"cvar": (alpha, 1.0)}))
    p = plan_of(m)
    e, scenario_profits = grade_plan(p, THREE)
    print(
        f"{alpha:8.3f} {pyo.value(m.cvar):12,.2f} {p[0]:7.1f} {p[1]:7.1f} {p[2]:7.1f}"
        f" {e:12,.0f} {scenario_profits.min():10,.0f}"
    )

print(
    f"\nalpha = 0    : CVaR* = ${rp_expected:,.2f} = the risk-neutral expected profit"
)
print(
    f"alpha >= 2/3 : CVaR* = ${profits_worst_case.min():,.2f} = the worst scenario profit"
)

# What the purpose-built models cost, against the general epigraph.
def size_of(m):
    """Active variables and constraints, scenario blocks included."""
    return (
        sum(1 for _ in m.component_data_objects(pyo.Var, active=True)),
        sum(1 for _ in m.component_data_objects(pyo.Constraint, active=True)),
    )


expectation_model = solve_farmer(build_farmer(THREE))
epigraph_at_0 = solve_farmer(build_farmer(THREE, risk={"cvar": (0.0, 1.0)}))
minimax_model = solve_farmer(build_farmer(THREE, risk={"minimax": True}))
epigraph_at_high = solve_farmer(build_farmer(THREE, risk={"cvar": (0.99, 1.0)}))

rows = [
    (
        "expectation (alpha -> 0 limit)",
        expectation_model,
        pyo.value(expectation_model.expected_profit),
    ),
    ("epigraph at alpha = 0", epigraph_at_0, pyo.value(epigraph_at_0.cvar)),
    ("minimax (alpha -> 1 limit)", minimax_model, pyo.value(minimax_model.worst)),
    ("epigraph at alpha = 0.99", epigraph_at_high, pyo.value(epigraph_at_high.cvar)),
]
print(f"{'model':>32s} {'objective':>12s} {'vars':>6s} {'cons':>6s}")
for label, m, value in rows:
    v, c = size_of(m)
    print(f"{label:>32s} {value:12,.2f} {v:6d} {c:6d}")

print("\nSame two answers, reached twice. The epigraph adds one threshold variable")
print("and one shortfall variable per scenario, plus one constraint per scenario,")
print("to compute numbers that a mean and a minimax deliver directly.")
print("\n1/(1-alpha), the epigraph coefficient:")
for alpha in [0.5, 0.9, 0.99, 0.999, 0.9999]:
    print(f"  alpha = {alpha:<8.4f}  1/(1-alpha) = {1 / (1 - alpha):>10,.1f}")

# Demonstrate it. Read the scenario profits straight off a pure-CVaR solve, then
# read them again after fixing the plan and re-optimizing recourse.
print(
    f"{'lambda':>7s} {'wheat':>7s} {'corn':>7s} {'beets':>7s}"
    f" {'E (read off)':>14s} {'E (re-optimized)':>18s} {'gap':>10s}"
)
for lam in [0.0, 0.25, 0.5, 0.75, 0.9, 1.0]:
    m = solve_farmer(build_farmer(THREE, risk={"cvar": (2 / 3, lam)}))
    p = plan_of(m)
    read_off = np.array([pyo.value(m.scenarios[i].profit) for i in m.SCENARIOS])
    e_true, _ = grade_plan(p, THREE)
    print(
        f"{lam:7.2f} {p[0]:7.1f} {p[1]:7.1f} {p[2]:7.1f}"
        f" {read_off.mean():14,.0f} {e_true:18,.0f} {e_true - read_off.mean():10,.0f}"
    )

# The same trap in a model that contains no CVaR at all.
mm = solve_farmer(build_farmer(THREE, risk={"minimax": True}))
mm_read_off = np.array([pyo.value(mm.scenarios[i].profit) for i in mm.SCENARIOS])
mm_true, mm_profits = grade_plan(plan_of(mm), THREE)
print(f"\nminimax (B&L 9.15): worst = ${pyo.value(mm.worst):,.0f}")
print(
    f"  scenario profits read off the solve : "
    + " | ".join(f"${p:,.0f}" for p in mm_read_off)
)
print(
    f"  same plan, recourse re-optimized    : "
    + " | ".join(f"${p:,.0f}" for p in mm_profits)
)
print(f"  error if the first line were reported: ${mm_true - mm_read_off.mean():,.0f}")

selected = solve_farmer(build_farmer(THREE, risk={"cvar": (2/3, 0.25)}))
plan = plan_of(selected)
mean, profits = grade_plan(plan, THREE)

print("weight   wheat   corn   beets   mean profit   worst profit ($)")
for weight in [0.0, 0.25, 0.5, 1.0]:
    selected = solve_farmer(build_farmer(THREE, risk={"cvar": (2 / 3, weight)}))
    plan = plan_of(selected)
    mean, profits = grade_plan(plan, THREE)
    print(
        f"{weight:6.2f} {plan[0]:7.1f} {plan[1]:6.1f} {plan[2]:7.1f}"
        f" {mean:13,.0f} {profits.min():16,.0f}"
    )

# Four plans, seven measures. Every plan is graded with recourse re-optimized.
def risk_profile(plan, target=100_000.0):
    _, p = grade_plan(plan, THREE)
    p = np.sort(p)
    out = {
        "E": p.mean(),
        "worst": p[0],
        "std dev": p.std(ddof=1),
        "MAD": np.abs(p - p.mean()).mean(),
    }
    for alpha in (2 / 3, 0.9):
        k = max(1, int(np.floor((1 - alpha) * len(p))))
        out[f"CVaR_{alpha:.2f}"] = p[:k].mean()
    out[f"EDR_{int(target / 1000)}k"] = np.maximum(target - p, 0).mean()
    return out


plans = [
    ("risk neutral (RP)", rp_plan),
    ("floor $56,000", [p for f, p, _ in frontier if f == 56_000][0]),
    ("Table 4 / worst case", table_4_plan),
    ("robust minimax (9.15)", plan_of(mm)),
]
header = None
for label, plan in plans:
    row = risk_profile(plan)
    if header is None:
        header = list(row)
        print(f"{'plan':>22s}" + "".join(f"{h:>12s}" for h in header))
    print(f"{label:>22s}" + "".join(f"{row[h]:12,.0f}" for h in header))

print("\nThe robust model returns exactly the Table 4 plan, which is what Birge and")
print("Louveaux assert on p. 19 without solving it.")

# The expected-downside-risk frontier, B&L (5.1)-(5.3), p. 68, at g = $100,000.
TARGET = 100_000.0

# What is the smallest achievable expected downside risk? Ask, rather than guess:
# a risk CONSTRAINT can be infeasible, where a risk OBJECTIVE never is.
m = build_farmer(THREE)
m.u = pyo.Var(m.SCENARIOS, domain=pyo.NonNegativeReals, units=u.USD)
m.u_def = pyo.Constraint(
    m.SCENARIOS, rule=lambda m, i: m.u[i] >= TARGET * u.USD - m.scenarios[i].profit
)
m.objective.deactivate()
m.min_edr = pyo.Objective(expr=sum(m.u[i] for i in m.SCENARIOS) / 3, sense=pyo.minimize)
solve_farmer(m)
edr_floor = pyo.value(m.min_edr)
print(f"smallest attainable expected downside risk = ${edr_floor:,.0f}")
print(f"any limit below that is infeasible\n")

print(
    f"{'limit l':>10s} {'wheat':>8s} {'corn':>8s} {'beets':>8s}"
    f" {'E[profit]':>12s} {'achieved EDR':>14s} {'worst':>10s}"
)
for level in [20_000, 17_000, 15_000, 14_000]:
    m = solve_farmer(build_farmer(THREE, risk={"downside": (TARGET, level)}))
    p = plan_of(m)
    e, scenario_profits = grade_plan(p, THREE)
    achieved = np.maximum(TARGET - scenario_profits, 0).mean()
    print(
        f"{level:10,d} {p[0]:8.1f} {p[1]:8.1f} {p[2]:8.1f} {e:12,.0f}"
        f" {achieved:14,.0f} {scenario_profits.min():10,.0f}"
    )

# Exercise 8 data, p. 19. Purchase prices are 40% above selling prices.
EX8_PLANT_COST = {"WHEAT": 180.0, "CORN": 280.0, "BEETS": 310.0}
EX8_SELL_PRICE = {
    "WHEAT": 300.0,
    "CORN": 170.0,
    "BEETS_FAVORABLE": 41.0,
    "BEETS_UNFAVORABLE": 11.0,
}


def price_scenario(wheat_price, yield_multiplier=1.0):
    """One scenario of Exercise 8: a wheat selling price, optionally with a yield."""
    sell = dict(EX8_SELL_PRICE, WHEAT=wheat_price)
    return {
        "yield": {c: MEAN_YIELD[c] * yield_multiplier for c in CROPS},
        "sell": sell,
        "buy": {c: 1.4 * sell[c] for c in FEED_CROPS},
    }


# Wait-and-see: what the farmer would do knowing next year's wheat price.
wait_and_see = {}
for price in (300.0, 220.0):
    m = solve_farmer(build_farmer([price_scenario(price)], plant_cost=EX8_PLANT_COST))
    wait_and_see[price] = pyo.value(m.expected_profit)
    p = plan_of(m) + 0.0  # normalize any -0.0 the LP returns
    print(
        f"perfect forecast, wheat at ${price:.0f}/T:"
        f"  acres {p[0]:5.1f} / {p[1]:5.1f} / {p[2]:5.1f}"
        f"   profit ${wait_and_see[price]:,.0f}"
    )

# Here-and-now: one plan, both prices, recourse after the price is revealed.
price_sp = solve_farmer(
    build_farmer(
        [price_scenario(300.0), price_scenario(220.0)], plant_cost=EX8_PLANT_COST
    )
)
price_plan = plan_of(price_sp)
rp_price = pyo.value(price_sp.expected_profit)
ws_price = np.mean(list(wait_and_see.values()))
print(
    f"\nstochastic plan: acres {price_plan[0]:5.1f} / {price_plan[1]:5.1f}"
    f" / {price_plan[2]:5.1f}   E[profit] ${rp_price:,.0f}"
)
print(f"WS   = ${ws_price:,.0f}")
print(f"RP   = ${rp_price:,.0f}")
print(f"EVPI = ${ws_price - rp_price:,.0f}   <- Exercise 8's question")

# EEV: plan at the mean price, then live with both.
mean_price_model = solve_farmer(
    build_farmer([price_scenario(260.0)], plant_cost=EX8_PLANT_COST)
)
eev, _ = grade_plan(
    plan_of(mean_price_model),
    [price_scenario(300.0), price_scenario(220.0)],
    plant_cost=EX8_PLANT_COST,
)
print(f"EEV  = ${eev:,.0f}")
print(f"VSS  = ${rp_price - eev:,.0f}")

# Why VSS = 0 here: does the recourse actually move with the price?
for i in price_sp.SCENARIOS:
    b = price_sp.scenarios[i]
    sales = {s: pyo.value(b.sales[s]) for s in SALE_TYPES}
    buys = {c: pyo.value(b.purchases[c]) for c in FEED_CROPS}
    label = "wheat $300" if i == 0 else "wheat $220"
    print(
        f"{label}:  sales "
        + ", ".join(f"{k} {v:,.0f}" for k, v in sales.items())
        + "   purchases "
        + ", ".join(f"{k} {v:,.0f}" for k, v in buys.items())
    )

identical = all(
    abs(
        pyo.value(price_sp.scenarios[0].sales[s])
        - pyo.value(price_sp.scenarios[1].sales[s])
    )
    < 1e-6
    for s in SALE_TYPES
) and all(
    abs(
        pyo.value(price_sp.scenarios[0].purchases[c])
        - pyo.value(price_sp.scenarios[1].purchases[c])
    )
    < 1e-6
    for c in FEED_CROPS
)
print(f"\nrecourse identical in both price scenarios: {identical}")

# Prices AND yields: two prices x three yields = six scenarios.
joint = [
    price_scenario(price, factor)
    for price in (300.0, 220.0)
    for factor in (1.2, 1.0, 0.8)
]
joint_model = solve_farmer(build_farmer(joint, plant_cost=EX8_PLANT_COST))
p = plan_of(joint_model)
print(
    f"acres {p[0]:.1f} / {p[1]:.1f} / {p[2]:.1f}"
    f"   E[profit] ${pyo.value(joint_model.expected_profit):,.0f}"
)
print(
    "scenario profits  "
    + " | ".join(
        f"${pyo.value(joint_model.scenarios[i].profit):,.0f}"
        for i in joint_model.SCENARIOS
    )
)
print(
    f"\nobjective polynomial degree = "
    f"{joint_model.objective.expr.polynomial_degree()}   (1 means linear)"
)