The price signal and the substation: how many homes for how many kilowatts?

import os

for name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(name, "1")  # before NumPy: one BLAS thread per chain
os.environ.setdefault("JAX_PLATFORMS", "cpu")  # nothing here uses a GPU

import contextlib
import logging
import sys
import time

import arviz as az
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import nutpie
import pandas as pd
import pymc as pm
import scipy.sparse as sparse
from matplotlib.colors import BoundaryNorm, ListedColormap, LogNorm, TwoSlopeNorm
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch, Patch
from matplotlib.ticker import FuncFormatter
from scipy import stats

from pymc_challenges import data
from pymc_challenges.bigrun import PreflightError, Store, banner, estimated, git_hash, measured

# ArviZ leaves a lazy stub under "jax", and Matplotlib's first plot tests it, which would import all of
# JAX for nothing (see E89). Nothing here uses JAX, so the stub goes.
sys.modules.pop("jax", None)

RANDOM_SEED = 91
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)

LEVEL = os.environ.get("PYMC_CHALLENGES_RUN", "").strip().lower() or "small"
if LEVEL != "small":
    raise PreflightError(
        f"PYMC_CHALLENGES_RUN={LEVEL!r}: E91 has one level, small. The household fleet "
        "and SBC have no route under this machine's rules, so nothing can be replayed."
    )

CORES = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
CHAINS_AT_ONCE = min(4, CORES)
SF_C, CM_C, DUD_C, MUTED, INK = "#0072b2", "#d55e00", "#b0226f", "#8a8a8a", "#2b2b2b"  # randomised, placed, 28 Dec
SKY = "#3f8fc4"  # a second blue, at least 3:1 on white
GATE_BANDS = ("SF High evening", "SF Low evening", "CM High evening")
GROUPS = ("SF winter", "SF summer", "CM winter", "CM summer")  # winter: GMT months; summer: BST months. CM: one row only
ACCEPTANCE = 0.24  # an assumption: the slides' uptake has no traced denominator
DUD_DAY = "2013-12-28"
PRICE = {}


@contextlib.contextmanager
def stage(name):
    """Add the time spent in this block to the price table printed at the end."""
    start = time.perf_counter()
    yield
    PRICE[name] = PRICE.get(name, 0.0) + time.perf_counter() - start


print(f"{CORES} cores here: {CHAINS_AT_ONCE} chains at once")

def watts(array):
    """Whole watts as stored, with the missing marker turned into NaN."""
    out = array.astype(float)
    out[array == -32768] = np.nan
    return out


n_hours = hour_file["fleet_kw"].size
gmt = pd.date_range(str(hour_file["first_hour"]), periods=n_hours, freq="h")
wall = gmt.tz_localize("UTC").tz_convert("Europe/London").tz_localize(None)
hour_of_day = wall.hour.to_numpy()
band = np.select(
    [
        (hour_of_day >= 5) & (hour_of_day < 11),
        (hour_of_day >= 11) & (hour_of_day < 17),
        (hour_of_day >= 17) & (hour_of_day < 23),
    ],
    [0, 1, 2],
    3,
)  # morning, day, evening, night
level = hour_of_day * 2 + (wall.dayofweek.to_numpy() >= 5)  # s(t): wall hour x weekday or weekend
trial_day = (wall - pd.Timedelta(hours=5)).strftime("%Y-%m-%d").to_numpy().astype(str)  # days run 05:00-05:00
in_2013 = (gmt >= "2013-01-01") & (gmt < "2014-01-01")
pre_trial = gmt < "2013-01-01"
after_trial = gmt >= "2014-01-01"

fleet_kw, panel_kw = hour_file["fleet_kw"].astype(float), hour_file["panel_kw"].astype(float)
effect = hour_file["effect"].astype(int)
EFFECTS = [str(e) for e in hour_file["effects"]]
usable = ~hour_file["out_of_baseline"] & np.isfinite(panel_kw) & np.isfinite(fleet_kw)

cells = pd.DataFrame({c[5:]: table_file[c] for c in table_file.files if c.startswith("cell_")})
placebo_cells = pd.DataFrame(
    {
        c[8:]: table_file[c]
        for c in table_file.files
        if c.startswith("placebo_") and c not in ("placebo_homes", "placebo_w", "placebo_sigma_w")
    }
)
for frame in (cells, placebo_cells):
    frame["day"] = frame["day"].astype(object)
    frame["effect"] = frame["effect"].astype(object)
    frame["season"] = frame["season"].astype(object)
home_tables = {"local": watts(table_file["local_w"]), "plan": watts(table_file["plan_w"])}
placebo_table = watts(table_file["placebo_w"])
print(
    f"{n_hours:,} hours from {gmt[0]:%Y-%m-%d} GMT. Fleet {home_tables['local'].shape[0]:,} homes x "
    f"{len(cells)} event cells; placebo {placebo_table.shape[0]:,} homes x {len(placebo_cells)} cells; "
    f"flat panel {int(hour_file['n_panel']):,} homes."
)
print(cells.groupby("effect", sort=False).agg(cells=("hours", "size"), hours=("hours", "sum")).T.to_string())

fig, ax = plt.subplots(figsize=(12.5, 5.3), layout="constrained")
ax.set(xlim=(0, 12.5), ylim=(0, 6.1))
ax.axis("off")


def box(x, y, text, colour, width=2.45, height=1.05, tint=None):
    ax.add_patch(
        FancyBboxPatch(
            (x, y),
            width,
            height,
            boxstyle="round,pad=0.02,rounding_size=0.08",
            facecolor=tint or "#f5f5f2",
            edgecolor=colour,
            linewidth=1.6,
        )
    )
    ax.text(x + width / 2, y + height / 2, text, ha="center", va="center", fontsize=9.5, color=INK)


def arrow(start, end, label, colour, offset=(0, 0.14), style="-|>"):
    ax.add_patch(
        FancyArrowPatch(start, end, arrowstyle=style, mutation_scale=13, color=colour, lw=1.6, shrinkA=0, shrinkB=0)
    )
    ax.text(
        (start[0] + end[0]) / 2 + offset[0],
        (start[1] + end[1]) / 2 + offset[1],
        label,
        ha="center",
        fontsize=8.8,
        color=INK,
    )


CHOICE = "#7a7a7a"
ax.text(
    0.05, 5.8, "Who got the tariff: a chain of choices. No coin was tossed", fontsize=11, weight="bold", color=INK
)
ax.text(0.05, 2.45, "Which day had a price event: this is where the coin is", fontsize=11, weight="bold", color=INK)
box(0.05, 4.2, "EDF customers\nin London", CHOICE)
box(3.45, 4.2, "Joined the project\n(their own choice)", CHOICE)
box(
    6.85,
    4.75,
    f"Took the tariff, passed filters\n{tou_homes:,} homes in this release",
    CHOICE,
    width=2.9,
    height=0.95,
)
box(6.85, 3.55, f"Stayed on a flat price\n{std_homes:,} homes here", CHOICE, width=2.9, height=0.95)
box(10.3, 4.75, "Paid to sign and to stay,\nprotected against loss", CHOICE, width=2.1, height=0.95)
arrow((2.55, 4.72), (3.4, 4.72), "opt-in", CHOICE)
arrow((5.95, 4.85), (6.8, 5.2), "opt-in", CHOICE, offset=(-0.12, 0.17))
arrow((5.95, 4.6), (6.8, 4.05), "", CHOICE)
arrow((9.8, 5.2), (10.25, 5.2), "", CHOICE)
box(0.05, 0.75, "Every tariff home saw\nthe same calendar in 2013", INK)
box(
    3.65,
    1.35,
    "90 Supply-Following events\n+ 3 all-Low days\nRANDOMISED over the year",
    SF_C,
    width=3.25,
    height=1.0,
    tint="#e3f0f8",
)
box(
    3.65,
    0.1,
    "Constraint-Management days\nPLACED on peaks, after SF\nNOT randomised",
    CM_C,
    width=3.25,
    height=1.0,
    tint="#fbe9dd",
)
box(
    7.65,
    1.35,
    "15 evening High events:\n11 in summer, only 4 in winter",
    SF_C,
    width=2.45,
    height=1.0,
    tint="#e3f0f8",
)
box(10.5, 1.35, "4 winter evenings\nThe unit that\nlimits the answer", DUD_C, width=1.9, height=1.0, tint="#f8e4ef")
box(
    7.65,
    0.1,
    "17 evening High days,\nmostly winter, shown with\ntheir placebo beside them",
    CM_C,
    width=2.45,
    height=1.0,
    tint="#fbe9dd",
)
arrow((2.55, 1.45), (3.6, 1.8), "a draw", SF_C, offset=(-0.15, 0.2))
arrow((2.55, 1.05), (3.6, 0.65), "a choice", CM_C, offset=(-0.2, -0.33))
arrow((6.95, 1.85), (7.6, 1.85), "", SF_C)
arrow((6.95, 0.6), (7.6, 0.6), "", CM_C)
arrow((10.15, 1.85), (10.45, 1.85), "", DUD_C);

calendar = store.load_json("calendar")
days_2013 = pd.date_range("2013-01-01", "2013-12-31", freq="D")
day_row = {d.strftime("%Y-%m-%d"): i for i, d in enumerate(days_2013)}
year = np.zeros((48, len(days_2013)), dtype=int)  # 0 default, 1-2 SF Low/High, 3-4 CM Low/High
OUT_DAYS = sorted(set(trial_day[hour_file["out_of_baseline"]]))
for e in calendar["events"]:
    if e["day"] in OUT_DAYS or e["day"] == "2013-03-29" or e["day"] not in day_row:
        continue
    start = pd.Timestamp(e["start_wall"])
    first = ((start.hour - 5) % 24) * 2 + start.minute // 30
    kind = (1 if e["band"] == "Low" else 2) + (0 if e["class"] == "single" else 2)
    year[first : first + e["slots"], day_row[e["day"]]] = kind
year[0:24, day_row["2013-03-29"]], year[24:36, day_row["2013-03-29"]], year[36:48, day_row["2013-03-29"]] = 3, 4, 3
for day in OUT_DAYS:
    if day in day_row:
        year[:, day_row[day]] = 5

colours = ["#f1f0ec", "#56b4e9", SF_C, "#e8963f", CM_C, "#7d7d7d"]
fig, ax = plt.subplots(figsize=(12.5, 4.6), layout="constrained")
ax.imshow(
    year,
    aspect="auto",
    interpolation="nearest",
    cmap=ListedColormap(colours),
    norm=BoundaryNorm(np.arange(-0.5, 6), 6),
    extent=(-0.5, 364.5, 29, 5),
)
ax.axhline(17, color=INK, lw=0.7, ls="--")
ax.axhline(23, color=INK, lw=0.7, ls="--")
sf_evening = cells[cells["effect"] == "SF High evening"]
winter = sf_evening[sf_evening["season"] == "gmt"]
for _, row in winter.iterrows():
    ax.annotate(
        pd.Timestamp(row["day"]).strftime("%-d %b"),
        (day_row[row["day"]], 17),
        xytext=(0, 18),
        textcoords="offset points",
        ha="center",
        fontsize=8.5,
        color=DUD_C if row["day"] == DUD_DAY else INK,
        weight="bold",
        arrowprops=dict(arrowstyle="-", color=DUD_C if row["day"] == DUD_DAY else INK, lw=0.9),
    )
months = pd.date_range("2013-01-01", "2013-12-01", freq="MS")
ax.set_xticks([day_row[m.strftime("%Y-%m-%d")] for m in months], [m.strftime("%b") for m in months])
ax.set_yticks([5, 11, 17, 23, 29], ["05:00", "11:00", "17:00", "23:00", "05:00"])
ax.set_ylabel("wall clock")
ax.legend(
    handles=[
        Patch(color=colours[2], label="SF High (randomised)"),
        Patch(color=colours[1], label="SF Low (randomised)"),
        Patch(color=colours[4], label="CM High (placed)"),
        Patch(color=colours[3], label="CM Low (placed)"),
        Patch(color=colours[5], label="days out of this analysis"),
    ],
    loc="lower center",
    ncols=5,
    bbox_to_anchor=(0.5, -0.27),
    fontsize=9,
)
ax.set_title(
    "2013 on one page. The four named evenings are all the randomised evidence for a winter peak",
    loc="left",
    fontsize=11,
);

gap = 1000 * (fleet_kw - panel_kw)
evening = band == 2
event_days = np.unique(trial_day[effect >= 0])
quiet = evening & ~np.isin(trial_day, event_days) & usable


def m0_row(label, mask):
    return {"where": label, "tariff minus flat (W)": np.nanmean(gap[mask]), "hours": int(mask.sum())}


m0 = pd.DataFrame(
    [
        m0_row("SF High evening event", effect == EFFECTS.index("SF High evening")),
        m0_row("CM High evening event", effect == EFFECTS.index("CM High evening")),
        m0_row("2013 evenings on days with no event", quiet & in_2013),
        m0_row("Jul-Dec 2012 evenings, no price existed", quiet & pre_trial),
        m0_row("Jan-Feb 2014 evenings, prices over", quiet & after_trial),
    ]
).set_index("where")
print(m0.round(1).to_string())

weeks = pd.Series(gap[quiet], index=gmt[quiet]).resample("W").mean()
fig, ax = plt.subplots(figsize=(11.5, 4.5), layout="constrained")
ax.axvspan(pd.Timestamp("2013-01-01"), pd.Timestamp("2014-01-01"), color="#f1f0ec", lw=0, zorder=0)
ax.axhline(0, color=INK, lw=0.8)
ax.plot(weeks.index, weeks.to_numpy(), color=MUTED, lw=1.6, label="evenings with no event, weekly mean")
for name, colour, marker in (("SF High evening", SF_C, "o"), ("CM High evening", CM_C, "s")):
    sel = cells[cells["effect"] == name]
    y = [np.nanmean(gap[(effect == EFFECTS.index(name)) & (trial_day == d)]) for d in sel["day"]]
    ax.scatter(
        pd.to_datetime(sel["day"]), y, s=46, color=colour, marker=marker, edgecolor="w", lw=1, zorder=3, label=name
    )
ax.annotate("2013: the prices run", (pd.Timestamp("2013-02-10"), 48), fontsize=9.5, color=INK)
ax.set(ylabel="tariff minus flat homes (W)", ylim=(-165, 70))
ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator()))
ax.legend(loc="lower left", ncols=3, fontsize=9)
ax.set_title(
    "The two groups differ on evenings with no event, so their gap cannot be read as an effect", loc="left", fontsize=11
);

clock = store.load_json("pilot_clock")
cost = store.load_json("clock_cost")
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.1), sharey=True, layout="constrained")
for ax, key, title, predicted in zip(axes, ("gmt_both", "bst_both"), ("GMT months", "BST months"), (1, -1)):
    t = clock["tables"][key]
    offset, mean, se = np.array(t["offset"]), np.array(t["mean_w"]), np.array(t["se_w"])
    ax.bar(
        offset,
        mean,
        yerr=1.96 * se,
        width=0.8,
        color=[SF_C if o == predicted else "#c7c7c7" for o in offset],
        error_kw=dict(lw=0.9, ecolor=INK),
    )
    ax.axvline(0, color=DUD_C, lw=1.4, ls="--")
    ax.annotate(
        "a naive join\nexpects it here",
        (0, 38),
        xytext=(-8, 0),
        textcoords="offset points",
        ha="right",
        fontsize=9,
        color=INK,
    )
    ax.axhline(0, color=INK, lw=0.7)
    ax.set_title(f"{title}, {t['n_events'][6]} edges", fontsize=10.5)
    ax.set(xticks=offset[::2])
    ax.tick_params(labelsize=10)
axes[0].set_ylabel("mean step, priced direction (W)", fontsize=10.5)
fig.supxlabel("offset in half-hours: stamp minus wall time of the price edge", fontsize=10.5)
fig.suptitle(
    "One rule fits both seasons: GMT all year, stamped at the end of the half-hour", x=0.01, ha="left", fontsize=11
)
naive, fixed = cost["wall_start"], cost["gmt_end"]
print("Mean tariff-minus-flat response inside SF runs, price direction (W). A timing check, never an effect:")
print(
    pd.DataFrame(
        {
            "naive join": [naive["gmt"]["mean_w"], naive["bst"]["mean_w"]],
            "gmt_end": [fixed["gmt"]["mean_w"], fixed["bst"]["mean_w"]],
        },
        index=["GMT", "BST"],
    )
    .assign(lost=lambda d: 1 - d["naive join"] / d["gmt_end"])
    .round(3)
    .to_string()
)
print(f"Clear of the worse neighbour by {clock['gate_margin_se']} se; the gate asked for 3");

def deviation(keep):
    """z[t]: the flat homes' load in kW, against its own mean at that wall hour and day type."""
    z = np.zeros(n_hours)
    for s in range(48):
        at = level == s
        z[at] = panel_kw[at] - panel_kw[at & keep].mean()
    return np.nan_to_num(z)


def month_band(keep):
    """One indicator per month x time band, for the post-hoc local baseline."""
    key = pd.factorize(pd.Series(gmt.strftime("%Y-%m")).astype(str) + band.astype(str))[0]
    used, column = np.unique(key[keep], return_inverse=True)
    row = np.flatnonzero(keep)
    return sparse.csr_matrix((np.ones(row.size), (row, column)), shape=(n_hours, used.size))


def design(keep, local=False):
    """Hours x columns for M1: 48 levels, z, 11 effects, and month x band terms if asked."""
    rows = np.arange(n_hours)
    on = effect >= 0
    parts = [
        sparse.csr_matrix((np.ones(n_hours), (rows, level)), shape=(n_hours, 48)),
        sparse.csr_matrix(deviation(keep)[:, None]),
        sparse.csr_matrix((np.ones(on.sum()), (rows[on], effect[on])), shape=(n_hours, 11)),
    ]
    if local:
        parts.append(month_band(keep))
    return sparse.hstack(parts, format="csr")


def m1(kw, keep, local=False):
    """Closed-form fleet: conjugate Gaussian with plug-in sigma, home by home. Effects in W."""
    x_all = design(keep, local)
    q = x_all.shape[1] - 60
    prior_mean = np.r_[np.full(48, 0.4), 1.0, np.zeros(11 + q)]
    prior_prec = 1 / np.r_[np.full(48, 0.5), 1.0, np.full(11, 0.3), np.full(q, 0.5)] ** 2
    out = {
        "mean": np.full((len(kw), 11), np.nan),
        "sd": np.full((len(kw), 11), np.nan),
        "sigma": np.full(len(kw), np.nan),
        "kurtosis": np.full(len(kw), np.nan),
    }
    for h, y in enumerate(kw):
        ok = np.isfinite(y) & keep
        x, target = x_all[ok], y[ok]
        gram, moment = (x.T @ x).toarray(), x.T @ target
        variance = target.var()
        for _ in range(2):  # plug-in sigma: two passes are enough at these lengths
            precision = gram / variance + np.diag(prior_prec)
            theta = np.linalg.solve(precision, moment / variance + prior_prec * prior_mean)
            residual = target - x @ theta
            variance = float(np.mean(residual**2))
        cov = np.linalg.inv(gram / variance + np.diag(prior_prec))
        out["mean"][h] = 1000 * theta[49:60]
        out["sd"][h] = 1000 * np.sqrt(np.diag(cov)[49:60])
        out["sigma"][h] = 1000 * np.sqrt(variance)
        out["kurtosis"][h] = stats.kurtosis(residual)
    return out


meter_wh = meter_file["wh"].astype(float)
meter_wh[meter_file["wh"] == 65535] = np.nan
meters = meter_wh / 1000  # Wh in an hour is mean W, so this is mean kW
started = time.perf_counter()
with stage("M1, 200 homes, three baselines"):
    live = {
        "whole window": m1(meters, usable),
        "2013 only": m1(meters, usable & in_2013),
        "2013, month x band": m1(meters, usable & in_2013, local=True),
    }
pooled = store.load_json("pilot_placebo_benchmark")
names = {"whole window": "p6", "2013 only": "p6_2013only", "2013, month x band": "p6_local_2013only"}
compare = {}
for name, fit in live.items():
    for band_name in GATE_BANDS:
        k = EFFECTS.index(band_name)
        compare[(name, band_name)] = {
            "200 homes, live": fit["mean"][:, k].mean(),
            "se between homes": fit["mean"][:, k].std(ddof=1) / np.sqrt(len(meters)),
            "1,088 homes, stored": pooled[names[name]]["bands"][band_name]["mean_w"],
        }
print(f"Three fleets of {len(meters)} homes in {time.perf_counter() - started:.1f} s, no sampler")
print(pd.DataFrame(compare).T.round(1).to_string())
print(
    f"\nResiduals: excess kurtosis across homes, median {np.median(live['whole window']['kurtosis']):.1f} "
    f"(10-90%: {np.quantile(live['whole window']['kurtosis'], 0.1):.1f} to "
    f"{np.quantile(live['whole window']['kurtosis'], 0.9):.1f}); a Gaussian has 0"
)

placebo = pooled["p5"]["bands"]
no_dud = pooled["p5_no28dec"]["bands"]
placebo_view = pd.DataFrame(
    {
        b: {
            "events": placebo[b]["events"],
            "mean (W)": placebo[b]["mean_w"],
            "se homes": placebo[b]["se_home_w"],
            "se events": placebo[b]["se_event_w"],
            "without the shifted 28 Dec": no_dud[b]["mean_w"],
        }
        for b in GATE_BANDS
    }
).T
print(
    "Placebo: the 2013 calendar moved back 364 days, Jul-Dec 2012, no prices existed. Gate: |mean| <= 10 W, se <= 5 W"
)
print(placebo_view.round(1).to_string())
print(
    f"Homes: {pooled['p5']['homes']}. Passes: {pooled['p5_passes']}. Over the 25 W line that ends the estimate: {pooled['p5_over25']}"
)
shifted = next(c for c in pooled["p5"]["cells"] if c["day"] == DUD_DAY and c["effect"] == "SF High evening")
dud_column = placebo_cells.index[(placebo_cells["day"] == DUD_DAY) & (placebo_cells["effect"] == "SF High evening")][0]
dud_se = np.nanstd(placebo_table[:, dud_column], ddof=1) / np.sqrt(np.isfinite(placebo_table[:, dud_column]).sum())
print(
    f"The placebo cell for 29 Dec 2012 alone, the day that lines up with 28 Dec 2013: {shifted['w']:+.1f} W (se {dud_se:.1f}), no price at all"
)
print(
    "All four SF High placebo cells:",
    [f"{c['w']:+.1f} W" for c in pooled["p5"]["cells"] if c["effect"] == "SF High evening"],
)
print(
    f"Post-hoc, not a rescue: with month x band terms CM's placebo is {pooled['p5_local']['bands']['CM High evening']['mean_w']:+.1f} W, "
    f"SF High evening {pooled['p5_local']['bands']['SF High evening']['mean_w']:+.1f} W"
)

baselines = {
    "Whole window (the plan)": "p6",
    "Same, without 28 Dec": "p6_no28dec",
    "2013 only": "p6_2013only",
    "2013, month x band (post-hoc)": "p6_local_2013only",
    "Same, without 28 Dec ": "p6_local_2013_no28dec",
}
window = pd.DataFrame({k: {b: pooled[v]["bands"][b]["mean_w"] for b in GATE_BANDS} for k, v in baselines.items()}).T
print(
    f"Benchmark on {pooled['p6']['homes']:,} homes (gate: CM -70 to -30 W, SF -50 to -15 W). Passes: {pooled['p6_passes']}"
)
print(window.round(1).to_string())

sf_high = cells[cells["effect"] == "SF High evening"]
event_hours_w = sf_high["hours"].to_numpy(float)
picks = np.random.default_rng(RANDOM_SEED + 5).integers(0, len(sf_high), (20_000, len(sf_high)))
print("\nSF High evening, hour-weighted, resampling the 15 events (the calendar's sample, not the homes'):")
for name, key in (("2013, month x band", "local"), ("whole window", "plan")):
    event_w = np.nanmean(home_tables[key][:, sf_high.index], axis=0)
    resampled = (event_w[picks] * event_hours_w[picks]).sum(axis=1) / event_hours_w[picks].sum(axis=1)
    print(
        f"  {name}: {np.average(event_w, weights=event_hours_w):.1f} W, 90% interval "
        f"{np.quantile(resampled, 0.05):.1f} to {np.quantile(resampled, 0.95):.1f} W"
    )

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.5), width_ratios=(1, 1.35), layout="constrained")
ax = axes[0]
ax.axhspan(-10, 10, color="#f1f0ec", lw=0, label="gate: within 10 W")
for line in (-25, 25):
    ax.axhline(line, color=DUD_C, lw=1, ls="--")
ax.annotate(
    "past 25 W the estimate is abandoned",
    (2.42, 25),
    xytext=(0, 3),
    textcoords="offset points",
    ha="right",
    fontsize=8.5,
    color=INK,
)
for i, b in enumerate(GATE_BANDS):
    colour = CM_C if b.startswith("CM") else SF_C
    r = placebo[b]
    ax.errorbar(
        i, r["mean_w"], 1.96 * max(r["se_home_w"], r["se_event_w"]), fmt="o", ms=9, color=colour, capsize=0, lw=1.8
    )
    ax.annotate(
        f"{r['mean_w']:+.1f} W", (i, r["mean_w"]), xytext=(10, 2), textcoords="offset points", fontsize=9, color=INK
    )
ax.axhline(0, color=INK, lw=0.7)
ax.set_xticks(range(3), ["SF High\n4 events", "SF Low\n10 events", "CM High\n5 events"])
ax.set(xlim=(-0.45, 2.45), ylim=(-40, 40), ylabel="W, with no price in existence", title="Before prices: the placebo")
ax = axes[1]
y = np.arange(len(window))[::-1]
for b, colour, marker in zip(GATE_BANDS, (SF_C, SKY, CM_C), ("o", "D", "s")):
    ax.scatter(window[b], y, s=70, color=colour, marker=marker, edgecolor="w", lw=1, zorder=3, label=b)
ax.axvline(0, color=INK, lw=0.7)
ax.set_yticks(y, [k.strip() for k in window.index])
ax.set(xlabel="effect of the event (W)", title="During prices: the same fleet, five baselines")
ax.legend(loc="center", fontsize=9, bbox_to_anchor=(0.72, 0.55));

with pm.Model() as m2:
    y_data = pm.Data("y", np.zeros(200))
    x_data = pm.Data("X", np.zeros((200, 11)))
    z_data = pm.Data("z", np.zeros(200))
    s_data = pm.Data("s", np.zeros(200, dtype="int64"))
    b = pm.Normal("b", 0.4, 0.5, shape=48)
    beta = pm.Normal("beta", 1, 1)
    delta = pm.Normal("delta", 0, 0.3, shape=11)
    sigma = pm.LogNormal("sigma", np.log(0.3), 0.7)
    nu_extra = pm.Gamma("nu_extra", 2, 0.1)
    location = b[s_data] + beta * z_data + x_data @ delta
    pm.StudentT("obs", nu=2 + nu_extra, mu=location, sigma=sigma, observed=y_data, shape=location.shape)

with stage("M2, compile"):
    compiled = nutpie.compile_pymc_model(m2)

z_all = deviation(usable)
effect_hours = np.zeros((n_hours, 11))
effect_hours[np.flatnonzero(effect >= 0), effect[effect >= 0]] = 1.0
student_rows = []
with stage("M2, 20 homes"):
    for i, h in enumerate(range(0, len(meters), 10)):
        ok = np.isfinite(meters[h]) & usable
        started = time.perf_counter()
        trace = nutpie.sample(
            compiled.with_data(y=meters[h, ok], X=effect_hours[ok], z=z_all[ok], s=level[ok].astype("int64")),
            chains=4,
            tune=500,
            draws=500,
            cores=CHAINS_AT_ONCE,
            seed=RANDOM_SEED + i,
            progress_bar=False,
            save_warmup=False,
            store_mass_matrix=False,
            store_gradient=False,
            store_unconstrained=False,
        )
        post = trace.posterior.to_dataset()
        hat = az.rhat(post)
        hat = hat.to_dataset() if hasattr(hat, "to_dataset") else hat
        ess = az.ess(post)
        ess = ess.to_dataset() if hasattr(ess, "to_dataset") else ess
        student_rows.append(
            {
                "home": int(meter_file["homes"][h]),
                "row": h,
                "seconds": time.perf_counter() - started,
                "max r_hat": max(float(np.nanmax(hat[v].values)) for v in hat.data_vars),
                "min ESS": min(float(np.nanmin(ess[v].values)) for v in ess.data_vars),
                "divergences": int(trace.sample_stats["diverging"].values.sum()),
                "nu": 2 + float(post["nu_extra"].median()),
                **{
                    b_name: 1000 * float(post["delta"].mean(("chain", "draw")).values[EFFECTS.index(b_name)])
                    for b_name in GATE_BANDS
                },
            }
        )
student = pd.DataFrame(student_rows)
print(
    student[["seconds", "max r_hat", "min ESS", "divergences", "nu"]]
    .describe()
    .loc[["min", "50%", "max"]]
    .round(3)
    .to_string()
)
print(
    f"Homes at max r_hat <= 1.01: {(student['max r_hat'] <= 1.01).sum()} of {len(student)}; <= 1.02: {(student['max r_hat'] <= 1.02).sum()}"
)
print(
    "Worst of the 20:",
    student.sort_values("max r_hat")
    .iloc[-1][["home", "max r_hat", "min ESS", "divergences", "seconds"]]
    .round(2)
    .to_dict(),
)
paired = {
    b_name: live["whole window"]["mean"][student["row"].to_numpy(), EFFECTS.index(b_name)]
    - student[b_name].to_numpy()
    for b_name in GATE_BANDS
}
print(
    pd.DataFrame(
        {
            b_name: {
                "M1 mean of 20 homes (W)": live["whole window"]["mean"][
                    student["row"].to_numpy(), EFFECTS.index(b_name)
                ].mean(),
                "M2 mean of 20 homes (W)": student[b_name].mean(),
                "paired se (W)": paired[b_name].std(ddof=1) / np.sqrt(len(student)),
            }
            for b_name in GATE_BANDS
        }
    )
    .round(1)
    .to_string()
)

fit_pilot, location_pilot = store.load_json("pilot_fit"), store.load_json("pilot_mean_location")
print(
    f"\nStored, 100 homes: median {fit_pilot['median_s']:.1f} s a fit, {fit_pilot['total_div']} divergences, "
    f"{fit_pilot['rhat_le_1.01']} homes at r_hat <= 1.01 and {fit_pilot['rhat_le_1.02']} at 1.02, "
    f"median nu {fit_pilot['nu_median']:.3f}. Worst homes:"
)
print(pd.DataFrame(fit_pilot["worst"])[["home", "rhat", "nu", "div"]].head(3).to_string(index=False))

QUANTILES = np.arange(0.05, 0.96, 0.05)


def shift_gaps(hours_on):
    """These hours minus matched default hours within 28 days, by quantile of each home's own hours, W."""
    clean = (effect < 0) & usable & in_2013
    matched = []
    for t in hours_on:
        window_hours = np.arange(max(t - 672, 0), min(t + 672, n_hours))
        matched.append(window_hours[clean[window_hours] & (level[window_hours] == level[t])])
    matched = np.concatenate(matched)
    usual = np.nanquantile(meters[:, matched], QUANTILES, axis=1)
    gaps = np.nanquantile(meters[:, hours_on], QUANTILES, axis=1) - usual
    mean_gap = np.nanmean(meters[:, hours_on], axis=1) - np.nanmean(meters[:, matched], axis=1)
    return 1000 * gaps, 1000 * usual, 1000 * mean_gap, 1000 * np.nanmean(meters[:, matched], axis=1)


def shift_function(which):
    """The shift across 200 homes: mean gap by quantile, its se, the mean hour, and the gap in percent."""
    gaps, usual, mean_gap, mean_usual = shift_gaps(np.flatnonzero((effect == EFFECTS.index(which)) & in_2013))
    return (
        gaps.mean(axis=1),
        gaps.std(axis=1, ddof=1) / np.sqrt(len(meters)),
        mean_gap.mean(),
        100 * gaps.mean(axis=1) / usual.mean(axis=1),
        100 * mean_gap.mean() / mean_usual.mean(),
    )


def no_price_shifts(which, draws=40):
    """The same recipe on calendars with no price: each event moved whole weeks onto default hours."""
    hours_on = np.flatnonzero((effect == EFFECTS.index(which)) & in_2013)
    clean = (effect < 0) & usable & in_2013
    move_rng = np.random.default_rng(RANDOM_SEED + 4)
    options = []
    for day in np.unique(trial_day[hours_on]):
        block = hours_on[trial_day[hours_on] == day]
        moves = [block + 168 * w for w in range(-6, 7) if w and (block + 168 * w).max() < n_hours]
        options.append([m for m in moves if clean[m].all()])
    return np.array(
        [
            shift_gaps(np.concatenate([o[move_rng.integers(len(o))] for o in options]))[0].mean(axis=1)
            for _ in range(draws)
        ]
    )


q = QUANTILES
shift_sf, shift_se, shift_mean, percent_sf, percent_mean = shift_function("SF High evening")
shift_cm, shift_cm_se, cm_mean, percent_cm, _ = shift_function("CM High evening")
control_sf, control_cm = no_price_shifts("SF High evening"), no_price_shifts("CM High evening")
print(
    pd.DataFrame(
        {
            "SF High evening (W)": shift_sf,
            "no-price calendars, mean (W)": control_sf.mean(axis=0),
            "SF minus no-price (W)": shift_sf - control_sf.mean(axis=0),
            "SF, percent of that quantile": percent_sf,
            "CM High evening (W)": shift_cm,
            "CM no-price calendars, mean (W)": control_cm.mean(axis=0),
        },
        index=[f"{p:.0%}" for p in q],
    )
    .T.round(1)
    .iloc[:, [1, 5, 9, 13, 15, 16, 17, 18]]
    .to_string()
)
print(f"Mean hour: SF {shift_mean:.0f} W ({percent_mean:.1f}%), CM {cm_mean:.0f} W")
print(
    f"No-price calendars at the 95th percentile, 5-95%: {np.quantile(control_sf[:, -1], 0.05):.0f} to "
    f"{np.quantile(control_sf[:, -1], 0.95):.0f} W; at the median {np.quantile(control_sf[:, 9], 0.05):.0f} to "
    f"{np.quantile(control_sf[:, 9], 0.95):.0f} W"
)

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.6), width_ratios=(1.2, 1), layout="constrained")
ax = axes[0]
for values, se, colour, label, marker in (
    (shift_sf, shift_se, SF_C, "SF High evening", "o"),
    (shift_cm, shift_cm_se, CM_C, "CM High evening (not randomised)", "s"),
):
    ax.fill_between(q * 100, values - 1.96 * se, values + 1.96 * se, color=colour, alpha=0.14, lw=0)
    ax.plot(q * 100, values, color=colour, lw=2, marker=marker, ms=5, label=label)
ax.plot(q * 100, control_sf.mean(axis=0), color=MUTED, lw=1.8, ls="--", label="SF recipe on no-price calendars")
ax.axhline(0, color=INK, lw=0.8)
ax.axvspan(40, 60, color="#f1f0ec", lw=0, zorder=0)
ax.annotate("a robust fit\nlooks here", (50, -165), ha="center", fontsize=9.5, color=INK)
ax.annotate("the busiest\nhours", (86, -175), ha="center", fontsize=9.5, color=INK)
ax.set(
    xlabel="quantile of a home's hours, from quiet to busy",
    ylabel="event minus matched default hours (W)",
    ylim=(-200, 40),
)
ax.legend(loc="lower left", fontsize=9)
ax.set_title("The top quantiles move most, and part of that shift appears with no price", loc="left", fontsize=11)

ax = axes[1]
order = list(GATE_BANDS)
y = np.arange(3)[::-1]
for yi, b_name in zip(y, order):
    r = location_pilot["bands"][b_name]
    ax.plot([r["m1"], r["m2"]], [yi, yi], color="#c7c7c7", lw=3, zorder=1)
    ax.scatter(r["m1"], yi, s=95, color=SF_C, zorder=3, label="M1, the mean" if yi == 2 else None)
    ax.scatter(
        r["m2"],
        yi,
        s=95,
        facecolor="w",
        edgecolor=DUD_C,
        lw=2,
        zorder=3,
        label="M2, Student-t location" if yi == 2 else None,
    )
    ax.scatter(r["gd"], yi, s=80, marker="x", color=INK, zorder=4, label="group difference, a rough check" if yi == 2 else None)
ax.axvline(0, color=INK, lw=0.7)
ax.set_yticks(y, ["SF High\nevening", "SF Low\nevening", "CM High\nevening"])
ax.set(xlabel="what each method reads, same 100 homes (W)", ylim=(-0.6, 2.9))
ax.legend(loc="upper center", ncols=3, fontsize=8.2, bbox_to_anchor=(0.5, 1.0), handletextpad=0.1, columnspacing=0.7)
ax.set_title("The robust fit misses SF Low and CM. On SF High the gap is not clear", loc="left", fontsize=11);