Where should the next air monitors go?

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

import contextlib
import gc
import json
import logging
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import arviz as az
import matplotlib.patheffects as pe
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
from matplotlib.collections import LineCollection, PolyCollection
from matplotlib.colors import LinearSegmentedColormap, LogNorm, Normalize, TwoSlopeNorm
from matplotlib.ticker import NullFormatter
from scipy import linalg
from scipy.special import logsumexp, ndtr, ndtri

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

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

RANDOM_SEED = 89
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}: E89 has one level, small. No national fit has "
                         "been run and no script exists for one, so there is nothing to replay or scale up.")
assert pytensor.config.floatX == "float64", "the annual path is float64"

CORES = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
CHAINS_AT_ONCE = min(4, CORES)
WORKERS = min(4, CORES)
LINE = 9.05  # the annual standard: design values round to 0.1
UG = "µg/m³"
METHODS = ("T640", "BAM", "other")  # FRM is the reference, so predictions are FRM-equivalent
YEARS = (2023, 2024, 2025)
SCALES = ("NEIGHBORHOOD", "URBAN SCALE", "REGIONAL SCALE", "MIDDLE SCALE", "MICROSCALE", "unknown")
PRIOR_ELL, PRIOR_ETA, PRIOR_TAU = (np.log(400.0), 0.6), (np.log(0.2), 0.6), (np.log(0.1), 0.6)
ABOVE, BELOW, PICK, MUTED = "#b2182b", "#2166ac", "#e66101", "#8a8a8a"
CLOSE_CALL = LinearSegmentedColormap.from_list(  # darkest where the call is closest, pale where it is settled
    "close_call", [(0, "#e6edf2"), (0.25, "#6baed6"), (0.4999, "#08306b"), (0.5001, "#67000d"), (0.75, "#fb6a4a"),
                   (1, "#fde8e2")])
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, {WORKERS} scoring threads")

def albers(lon, lat, lon0=-96.0, lat0=37.5, p1=29.5, p2=45.5, radius=6371.0):
    """Albers equal-area conic in km, written out so that no GIS library is needed."""
    lon, lat, lon0, lat0, p1, p2 = map(np.radians, (lon, lat, lon0, lat0, p1, p2))
    n = (np.sin(p1) + np.sin(p2)) / 2
    c = np.cos(p1) ** 2 + 2 * n * np.sin(p1)
    rho = radius * np.sqrt(c - 2 * n * np.sin(lat)) / n
    rho0 = radius * np.sqrt(c - 2 * n * np.sin(lat0)) / n
    theta = n * (lon - lon0)
    return rho * np.sin(theta), rho0 - rho * np.cos(theta)


def county_shapes():
    """County outer rings in Albers km, with each ring's row in `counties`."""
    row = {fips: i for i, fips in enumerate(counties["fips"])}
    fixes = {"46113": "46102"}  # the geojson's Shannon County SD is now Oglala Lakota
    verts, owner = [], []
    for feature in json.loads(data.path("us_counties_geojson").read_text())["features"]:
        fips = fixes.get(feature["id"], feature["id"])
        if fips not in row:
            continue
        geometry = feature["geometry"]
        polygons = [geometry["coordinates"]] if geometry["type"] == "Polygon" else geometry["coordinates"]
        for polygon in polygons:
            ring = np.asarray(polygon[0], float)[:, :2]
            verts.append(np.column_stack(albers(ring[:, 0], ring[:, 1])))
            owner.append(row[fips])
    return verts, np.array(owner)


def outline(verts, owner, group):
    """Edges between two groups of counties, or with a county on one side only: no GIS dissolve."""
    seen = {}
    for ring, o in zip(verts, owner):
        rounded = np.round(ring, 3)
        for a, b in zip(map(tuple, rounded[:-1]), map(tuple, rounded[1:])):
            seen.setdefault((a, b) if a < b else (b, a), []).append(group[o])
    return np.array([edge for edge, groups in seen.items() if len(set(groups)) > 1 or len(groups) == 1])


VERTS, OWNER = county_shapes()
BORDERS = outline(VERTS, OWNER, counties["fips"].str[:2].to_numpy())
CROP_EDGE = outline(VERTS, OWNER, counties["small_crop"].to_numpy())
in_crop = counties["small_crop"].to_numpy()
crop_row = np.cumsum(in_crop) - 1  # county row to row of `table`
CROP_VERTS = [v for v, o in zip(VERTS, OWNER) if in_crop[o]]
CROP_OWNER = crop_row[OWNER[in_crop[OWNER]]]
points = np.concatenate(CROP_VERTS)
XLIM = (points[:, 0].min() - 25, points[:, 0].max() + 25)
YLIM = (points[:, 1].min() - 25, points[:, 1].max() + 25)
del points
near = [v for v, o in zip(VERTS, OWNER) if not in_crop[o] and v[:, 0].max() > XLIM[0]
        and v[:, 0].min() < XLIM[1] and v[:, 1].max() > YLIM[0] and v[:, 1].min() < YLIM[1]]
print(f"{len(VERTS):,} county polygons, {sum(map(len, VERTS)):,} vertices, {len(BORDERS):,} border edges")

STATE_CODE = {"Colorado": "CO", "Illinois": "IL", "Indiana": "IN", "Iowa": "IA", "Kansas": "KS",
              "Kentucky": "KY", "Michigan": "MI", "Minnesota": "MN", "Missouri": "MO", "Montana": "MT",
              "Nebraska": "NE", "North Dakota": "ND", "Ohio": "OH", "Pennsylvania": "PA",
              "South Dakota": "SD", "Tennessee": "TN", "West Virginia": "WV", "Wisconsin": "WI",
              "Wyoming": "WY"}
short_name = (table["county"].str.replace(" County", "").astype(object) + " "
              + table["state"].map(STATE_CODE).astype(object)).to_numpy(dtype=object)


def paint(ax, values, cmap, norm, labels=()):
    """The crop as a county map. Other counties stay pale, so that the crop's edge can be seen."""
    ax.add_collection(PolyCollection(near, facecolor="#f2f1ec", edgecolor="#ffffff", linewidth=0.2))
    palette = plt.get_cmap(cmap).with_extremes(bad="#e4e2db")
    shapes = PolyCollection(CROP_VERTS, array=np.asarray(values, float)[CROP_OWNER], cmap=palette, norm=norm,
                            edgecolor="face", linewidth=0.1)
    ax.add_collection(shapes)
    ax.add_collection(LineCollection(BORDERS, colors="#4a4a4a", linewidths=0.35))
    ax.set(xlim=XLIM, ylim=YLIM, aspect="equal")
    ax.axis("off")
    for i, text in labels:
        ax.annotate(text, (table["x_km"].iloc[i], table["y_km"].iloc[i]), xytext=(4, 4), fontsize=8,
                    textcoords="offset points", path_effects=[pe.withStroke(linewidth=2.2, foreground="w")])
    return shapes


fig, ax = plt.subplots(figsize=(12.5, 7.4), layout="constrained")
shade = np.where(counties["monitored"], np.nan, counties["population"].to_numpy(float))
everywhere = PolyCollection(VERTS, array=shade[OWNER], cmap=plt.get_cmap("YlOrBr").with_extremes(bad="#c9d6df"),
                            norm=LogNorm(1e3, 1.3e6), edgecolor="face", linewidth=0.1)
ax.add_collection(everywhere)
ax.add_collection(LineCollection(BORDERS, colors="#4a4a4a", linewidths=0.35))
ax.add_collection(LineCollection(CROP_EDGE, colors="k", linewidths=1.3))
outside = ~sites["small_crop"]
ax.scatter(sites.loc[outside, "x_km"], sites.loc[outside, "y_km"], s=5, c="#5b6b78", lw=0, label="site outside the crop")
ax.scatter(crop_sites["x_km"], crop_sites["y_km"], s=9, c="#10222e", lw=0, label="site in the crop")
ax.autoscale_view()
ax.set(aspect="equal")
ax.set_title("Blue-grey counties have a monitor. Every other colour is an empty county. Black line: the crop",
             fontsize=11)
ax.axis("off")
ax.legend(loc="lower left", fontsize=9)
fig.colorbar(everywhere, ax=ax, shrink=0.55, label="people in a county with no monitor", pad=0.01);

def design(frame, scale=None):
    """What the model needs per site. Density is standardised here unless a scale is given."""
    density = frame["log_density"].to_numpy(float)
    centre, spread = scale or (density.mean(), density.std())
    has = frame[[f"has_{y}" for y in YEARS]].to_numpy(float)
    xy = frame[["x_km", "y_km"]].to_numpy(float)
    return {"site": frame["site"].to_numpy(str), "y": frame["log_mean"].to_numpy(float),
            "v": frame["var_log"].to_numpy(float), "density": (density - centre) / spread,
            "density_scale": (centre, spread),
            "methods": frame[[f"share_{m}" for m in METHODS]].to_numpy(float),
            "years": has / has.sum(axis=1, keepdims=True), "xy": xy, "days": frame["days"].to_numpy(),
            "scale": frame["scale"].map(SCALES.index).to_numpy(int), "distance": distance(xy, xy)}


def distance(a, b):
    return np.sqrt(((a[:, None] - b[None]) ** 2).sum(axis=-1))


def matern32(gap, ell):
    r = np.sqrt(3.0) * gap / ell
    return (1 + r) * np.exp(-r)


def m1(d, method_prior=(0.0, 0.2)):
    """The field integrated out: a marginal GP with 15 free parameters."""
    coords = {"method": METHODS, "year": YEARS, "scale": SCALES, "site": d["site"]}
    with pm.Model(coords=coords) as model:
        alpha = pm.Normal("alpha", 2.0, 0.5)
        beta_density = pm.Normal("beta_density", 0.0, 0.2)
        beta_method = pm.Normal("beta_method", *method_prior, dims="method")
        delta = pm.ZeroSumNormal("delta", sigma=0.1, dims="year")
        mu = (alpha + beta_density * d["density"] + pm.math.dot(d["methods"], beta_method)
              + pm.math.dot(d["years"], delta))
        eta = pm.LogNormal("eta", *PRIOR_ETA)
        ell = pm.LogNormal("ell", *PRIOR_ELL)
        tau = pm.LogNormal("tau", *PRIOR_TAU, dims="scale")
        r = np.sqrt(3.0) * pt.constant(d["distance"]) / ell
        cov = eta ** 2 * (1 + r) * pt.exp(-r) + pt.diag(tau[d["scale"]] ** 2 + d["v"])
        pm.MvNormal("y", mu=mu, cov=cov, observed=d["y"], dims="site")
    return model


d = design(crop_sites)
xy_county = table[["x_km", "y_km"]].to_numpy(float)
gap_county = distance(xy_county, xy_county)
density_county = (table["log_density"].to_numpy(float) - d["density_scale"][0]) / d["density_scale"][1]

with stage("prior check"):
    prior_rng = np.random.default_rng(RANDOM_SEED)
    rows = []
    for _ in range(200):
        eta, ell = prior_rng.lognormal(*PRIOR_ETA), prior_rng.lognormal(*PRIOR_ELL)
        tau = prior_rng.lognormal(*PRIOR_TAU, size=len(SCALES))
        mean = (prior_rng.normal(2.0, 0.5) + prior_rng.normal(0, 0.2) * d["density"]
                + d["methods"] @ prior_rng.normal(0, 0.2, size=3))
        cov = eta ** 2 * matern32(d["distance"], ell) + np.diag(tau[d["scale"]] ** 2 + d["v"])
        rows.append(prior_rng.multivariate_normal(mean, cov, method="cholesky"))
    prior_y = np.array(rows)
    print(pd.DataFrame({"prior 5%": [np.quantile(prior_y.std(axis=1), 0.05),
                                     np.quantile(np.exp(prior_y).mean(axis=1), 0.05)],
                        "prior 95%": [np.quantile(prior_y.std(axis=1), 0.95),
                                      np.quantile(np.exp(prior_y).mean(axis=1), 0.95)],
                        "observed": [d["y"].std(), np.exp(d["y"]).mean()]},
                       index=["sd of log site means", f"crop mean ({UG})"]).round(2).to_string())

    fig, axes = plt.subplots(2, 3, figsize=(12.5, 4.5), layout="constrained")
    for ax, q in zip(axes.flat, (0.05, 0.25, 0.45, 0.6, 0.8, 0.95)):  # ell across its prior, short to long
        ell = np.exp(PRIOR_ELL[0] + PRIOR_ELL[1] * ndtri(q))
        eta = prior_rng.lognormal(*PRIOR_ETA)
        lower = linalg.cholesky(eta ** 2 * matern32(gap_county, ell) + 1e-6 * np.eye(len(table)), lower=True)
        world = (prior_rng.normal(2.0, 0.5) + prior_rng.normal(0, 0.2) * density_county
                 + lower @ prior_rng.normal(size=len(table)))
        shapes = paint(ax, np.exp(world), "RdBu_r", TwoSlopeNorm(LINE, 3.0, 25.0))
        ax.set_title(f"ell = {ell:.0f} km, eta = {eta:.2f}", fontsize=9)
    fig.colorbar(shapes, ax=axes, shrink=0.7, label=f"annual PM2.5 ({UG})")
    fig.suptitle("Before any data: six worlds the priors allow, with ell stepped across its prior. "
                 "White is 9.05", fontsize=11)
    del lower, world, prior_y
    gc.collect();

def site_mean(design_table, draw):
    return (draw["alpha"] + draw["beta_density"] * design_table["density"]
            + design_table["methods"] @ draw["beta_method"] + design_table["years"] @ draw["delta"])


def median_draw(chains):
    return {k: np.median(v.reshape(-1, *v.shape[2:]), axis=0) for k, v in chains.items()}


def site_covariance(design_table, draw):
    return (draw["eta"] ** 2 * matern32(design_table["distance"], draw["ell"])
            + np.diag(draw["tau"][design_table["scale"]] ** 2 + design_table["v"]))


def hold_out(precision_matrix, residual, groups):
    """Leave each group out, from the precision matrix: errors and sds of the held-out sites."""
    error, sd = np.empty_like(residual), np.empty_like(residual)
    scaled = precision_matrix @ residual
    for group in groups:
        cov = np.linalg.inv(precision_matrix[np.ix_(group, group)])
        error[group] = cov @ scaled[group]
        sd[group] = np.sqrt(np.diag(cov))
    return error, sd


def scorecard(error, sd):
    return {"sites": error.size, "bias": error.mean(), "se of bias": error.std() / np.sqrt(error.size),
            "rmse": np.sqrt((error ** 2).mean()), "mean predictive sd": sd.mean(),
            "90% coverage": (np.abs(error / sd) < 1.645).mean()}


with stage("checks"):
    best = median_draw(fit)
    residual = d["y"] - site_mean(d, best)
    precision_sites = np.linalg.inv(site_covariance(d, best))
    n_sites = len(residual)

    # replicated networks, one per posterior draw
    stacked = {k: v.reshape(-1, *v.shape[2:]) for k, v in fit.items()}
    check_rng = np.random.default_rng(RANDOM_SEED)
    stats = []
    for i in check_rng.choice(stacked["ell"].size, 200, replace=False):
        draw = {k: v[i] for k, v in stacked.items()}
        replica = check_rng.multivariate_normal(site_mean(d, draw), site_covariance(d, draw),
                                                method="cholesky")
        stats.append((replica.std(), ((replica > np.log(8)) & (replica < np.log(10))).mean(), replica.max()))
    stats = np.array(stats)
    observed = (d["y"].std(), ((d["y"] > np.log(8)) & (d["y"] < np.log(10))).mean(), d["y"].max())
    print(pd.DataFrame(np.c_[np.quantile(stats, [0.05, 0.95], axis=0).T, observed],
                       columns=["replicated 5%", "replicated 95%", "observed"],
                       index=["sd of log means", "share of sites in 8-10", "largest log mean"]).round(3).to_string())

    # one site, 200 km blocks, and the new sites
    cells = pd.factorize(pd.Series(zip((d["xy"][:, 0] // 200).astype(int), (d["xy"][:, 1] // 200).astype(int))))[0]
    is_new = crop_sites["first_year"].to_numpy() >= 2023
    old_near = distance(d["xy"][is_new], d["xy"][~is_new]).min(axis=1)
    new_groups = [np.flatnonzero(is_new)]
    new_error, new_sd = hold_out(precision_sites, residual, new_groups)
    one_error, one_sd = hold_out(precision_sites, residual, [[i] for i in range(n_sites)])
    block_error, block_sd = hold_out(precision_sites, residual, [np.flatnonzero(cells == c) for c in np.unique(cells)])
    far = old_near > 25
    print(pd.DataFrame({"one site": scorecard(one_error, one_sd),
                        "200 km block": scorecard(block_error, block_sd),
                        "new sites": scorecard(new_error[is_new], new_sd[is_new]),
                        "new sites over 25 km": scorecard(new_error[is_new][far], new_sd[is_new][far])}
                       ).T.round(3).to_string())
    near_line = (np.exp(d["y"][is_new]) > 8) & (np.exp(d["y"][is_new]) < 10)
    print(f"New sites over 25 km with a window mean of 8-10: {(far & near_line).sum()}")
    by_scale = pd.DataFrame({"scale": crop_sites["scale"], "error": one_error}).groupby("scale")["error"]
    print(pd.DataFrame({"sites": by_scale.size(), "mean held-out error": by_scale.mean(),
                        "se": by_scale.std() / np.sqrt(by_scale.size())}).round(3).to_string())

    order = np.argsort(old_near)
    fig, ax = plt.subplots(figsize=(7.5, 5.6), layout="constrained")
    centre = (d["y"] - new_error)[is_new][order]
    half = 1.645 * new_sd[is_new][order]
    seen = d["y"][is_new][order]
    missed = np.abs(seen - centre) > half
    rows = np.arange(is_new.sum())
    ax.hlines(rows, np.exp(centre - half), np.exp(centre + half), color=MUTED, lw=2)
    ax.scatter(np.exp(seen)[~missed], rows[~missed], s=22, c="k", zorder=3, label="observed, inside 90%")
    ax.scatter(np.exp(seen)[missed], rows[missed], s=44, c=ABOVE, marker="X", zorder=3, label="observed, outside 90%")
    ax.axvline(LINE, color="k", lw=0.8, ls=":")
    ax.set_yticks(rows[::4], [f"{km_:.0f} km" for km_ in old_near[order][::4]])
    ax.set(xscale="log", xlabel=f"window mean ({UG})", ylabel="distance to the nearest older site",
           title="New sites, predicted from older ones")
    ax.set_xticks([3, 5, 7, 9, 12], ["3", "5", "7", "9", "12"])
    ax.xaxis.set_minor_formatter(NullFormatter())
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2, frameon=False);

def thin(chains, count, shift=0):
    """`count` draws evenly spaced through the chains; `shift` moves to a disjoint set."""
    total = chains["ell"].shape[0] * chains["ell"].shape[1]
    take = (np.linspace(0, total - 1, count).round().astype(int) + shift) % total
    return {k: v.reshape(total, *v.shape[2:])[take] for k, v in chains.items()}


def conditional(design_table, draw):
    """Exact mean of the latent field at the county centres, and H, for one draw."""
    lower = linalg.cholesky(site_covariance(design_table, draw), lower=True, check_finite=False)
    cross = draw["eta"] ** 2 * matern32(distance(xy_county, design_table["xy"]), draw["ell"])
    half = linalg.solve_triangular(lower, cross.T, lower=True, check_finite=False)
    residuals = linalg.solve_triangular(lower, design_table["y"] - site_mean(design_table, draw),
                                        lower=True, check_finite=False)
    density = (table["log_density"].to_numpy(float) - design_table["density_scale"][0]) / design_table["density_scale"][1]
    return draw["alpha"] + draw["beta_density"] * density + half.T @ residuals, half  # FRM, whole window


def make_field(design_table, chains, count=200, shift=0):
    """County means, H, and the noise of a new monitor, for `count` draws, as float32."""
    draws = thin(chains, count, shift)
    means = np.empty((count, len(table)), np.float32)
    half = np.empty((count, len(design_table["y"]), len(table)), np.float32)

    def fill(i):
        means[i], half[i] = conditional(design_table, {k: v[i] for k, v in draws.items()})

    with ThreadPoolExecutor(WORKERS) as pool:
        list(pool.map(fill, range(count)))
    noise = draws["tau"][:, SCALES.index("NEIGHBORHOOD")] ** 2 + np.median(design_table["v"][design_table["days"] >= 900])
    return {"means": means, "half": half, "noise": noise, "eta2": np.float32(draws["eta"] ** 2),
            "ells": np.float32(draws["ell"]), "ell": float(np.median(chains["ell"]))}


def covariance_rows(field_, rows_, keep=slice(None)):
    """Rows of S against every county: (draws, len(rows_), counties)."""
    eta2, ells = field_["eta2"][keep][:, None, None], field_["ells"][keep][:, None, None]
    half = field_["half"][keep]
    return eta2 * matern32(gap32[rows_], ells) - half[:, :, rows_].transpose(0, 2, 1) @ half


del idata
gc.collect()
with stage("county conditionals"):
    field = make_field(d, fit)
gap32 = np.float32(gap_county)
variance = field["eta2"][:, None] - (field["half"] ** 2).sum(axis=1)
p_now = ndtr((field["means"] - np.log(LINE)) / np.sqrt(variance)).mean(axis=0)
risk = population * np.minimum(p_now, 1 - p_now)
mean_log = field["means"].mean(axis=0)
sd_log = np.sqrt(variance.mean(axis=0) + field["means"].var(axis=0))
sd_ug = np.exp(mean_log) * sd_log
undecided = (p_now > 0.01) & (p_now < 0.99)
print(f"Expected wrong calls now: {risk.sum() / 1e6:.2f} million people of {population.sum() / 1e6:.1f} million "
      f"(modelled, not observed)")
print(f"Undecided counties (p between 0.01 and 0.99): {undecided.sum()} of {len(table)} ({undecided.mean():.0%})")
print(f"Counties called above 9.0: {(p_now > 0.5).sum()}, holding {population[p_now > 0.5].sum() / 1e6:.1f} million people")
monitored = table["monitored"].to_numpy()
print(f"Share of those wrong-call people in counties that already have a site: {risk[monitored].sum() / risk.sum():.0%}")

readings = pd.DataFrame({"fips": crop_sites["fips"], "reading": np.exp(d["y"]),
                         "settled": ndtr((d["y"] - np.log(LINE)) / np.sqrt(np.maximum(d["v"], 1e-8)))})
readings = readings.groupby("fips").agg(top=("reading", "max"), low=("reading", "min"), settled=("settled", "max"))
readings = table[["fips"]].join(readings, on="fips")
top_site, low_site = readings["top"].to_numpy(), readings["low"].to_numpy()
p_settled = np.where(monitored, readings["settled"].to_numpy(), p_now)
nugget_real = field["noise"][:, None] - np.median(d["v"][d["days"] >= 900])  # tau^2 of a neighbourhood site
p_local = ndtr((field["means"] - np.log(LINE)) / np.sqrt(variance + nugget_real)).mean(axis=0)
called = p_now > 0.5
print(f"Sites against calls: {(top_site >= LINE).sum()} counties have a site at or above {LINE}, and the model calls "
      f"{(called & (top_site >= LINE)).sum()} of them above; {(low_site >= LINE).sum()} have every site above it, and "
      f"{(called & (low_site >= LINE)).sum()} of those are called above")
conventions = pd.DataFrame({
    "people at risk (M)": [(population * np.minimum(p, 1 - p)).sum() / 1e6 for p in (p_settled, p_now, p_local)],
    "of which monitored (M)": [(population * np.minimum(p, 1 - p))[monitored].sum() / 1e6 for p in (p_settled, p_now, p_local)],
    "undecided counties": [((p > 0.01) & (p < 0.99)).sum() for p in (p_settled, p_now, p_local)]},
    index=["a site settles its county (sampling noise only)", "this notebook: truth is the smooth field",
           "the nugget is real air at the county centre"])
print(conventions.round(2).to_string())
worst = np.argsort(-risk)[:8]
print(pd.DataFrame({"people": population[worst].round(), "P(above)": p_now[worst].round(2),
                    "monitored": table["monitored"].to_numpy()[worst],
                    "people at risk of a wrong call": risk[worst].round()},
                   index=short_name[worst]).to_string())
allegheny = int(np.flatnonzero(short_name == "Allegheny PA")[0])
print(f"Allegheny PA: P(above) {p_now[allegheny]:.3f}, top site {top_site[allegheny]:.2f} {UG}, "
      f"{population[allegheny] / 1e6:.2f} million people")

fig, axes = plt.subplots(2, 1, figsize=(12.5, 11), layout="constrained")
shapes = paint(axes[0], p_now, CLOSE_CALL, Normalize(0, 1))
axes[0].scatter(crop_sites["x_km"], crop_sites["y_km"], s=4, c="#d95f02", lw=0)
axes[0].set_title("The call: chance that a county's air is above 9.0. The darkest counties are the closest "
                  "calls. Orange dots are monitors", loc="left", fontsize=11)
fig.colorbar(shapes, ax=axes[0], shrink=0.6, label="P(above 9.0)", pad=0.01)
paint(axes[1], np.where(undecided, 0.5, np.nan), "Greys", Normalize(0, 1))
for mask, colour, label in ((monitored, BELOW, "county with a monitor"), (~monitored, PICK, "empty county")):
    show = mask & (risk > 2e3)
    axes[1].scatter(xy_county[show, 0], xy_county[show, 1], s=risk[show] / 1.5e3, facecolor=colour, alpha=0.65,
                    edgecolor="w", lw=0.4, label=label)
for rank_, i in enumerate(np.argsort(-risk)[:6]):
    axes[1].annotate(short_name[i], xy_county[i], xytext=(0, 12 if rank_ % 2 else -12), textcoords="offset points",
                     fontsize=8, ha="center", va="center", path_effects=[pe.withStroke(linewidth=2.2, foreground="w")])
for size in (1e5, 5e5, 1e6):
    axes[1].scatter([], [], s=size / 1.5e3, facecolor=MUTED, alpha=0.65, edgecolor="w", lw=0.4,
                    label=f"{size / 1e6:g} million at risk")
axes[1].legend(loc="lower left", frameon=True, fontsize=9)
axes[1].set_title("The stakes: people in a county times its chance of a wrong call. "
                  "Grey counties are undecided", loc="left", fontsize=11);

def reading_nodes(mean, variance_, nodes):
    """Midpoints of `nodes` equal-probability bins of the mixture predictive of the reading."""
    sd = np.sqrt(variance_)
    grid = np.linspace((mean - 6 * sd).min(), (mean + 6 * sd).max(), 256)
    cdf = ndtr((grid[:, None] - mean) / sd).mean(axis=1)
    return np.interp((np.arange(nodes) + 0.5) / nodes, cdf, grid)


def update_weights(y, mean, variance_):
    """Each draw's posterior weight after reading `y` (axes: readings, draws)."""
    log = -0.5 * (y[:, None] - mean) ** 2 / variance_ - 0.5 * np.log(variance_)
    return np.exp(log - logsumexp(log, axis=1, keepdims=True))


def probability_above(y, weights, mu, var, cross, mu_c, q, line):
    """P(above the line) per county after each reading: rank-one update in each draw, then mix."""
    gain = cross / q[:, None]
    shift = (y[:, None] - mu_c)[:, :, None]
    sd = np.sqrt(np.maximum(var - gain * cross, 1e-10))
    return (weights[:, :, None] * ndtr((mu + shift * gain - line) / sd)).sum(axis=1)


def loss(p, weight, objective):
    """Expected loss, calling each county at p = 0.5."""
    if objective == "missed":  # called below, truly above
        return (weight * np.where(p < 0.5, p, 0.0)).sum(axis=-1)
    if objective == "false":   # called above, truly below
        return (weight * np.where(p >= 0.5, 1 - p, 0.0)).sum(axis=-1)
    return (weight * np.minimum(p, 1 - p)).sum(axis=-1)


def open_block(field_, line):
    """Covariance of every candidate with the undecided counties: all a score ever reads."""
    key = f"block_{line}"
    if key not in field_:
        variance_ = field_["eta2"][:, None] - (field_["half"] ** 2).sum(axis=1)
        p = ndtr((field_["means"] - np.float32(np.log(line))) / np.sqrt(variance_)).mean(axis=0)
        where = np.flatnonzero((p > 0.005) & (p < 0.995))  # a wider net than `open`, for subsets of draws
        block = np.empty((len(variance_), len(candidates), where.size), np.float32)
        gaps = gap32[np.ix_(candidates, where)]

        def fill(i):
            half = field_["half"][i]
            block[i] = field_["eta2"][i] * matern32(gaps, field_["ells"][i]) - half[:, candidates].T @ half[:, where]

        with ThreadPoolExecutor(WORKERS) as pool:
            list(pool.map(fill, range(len(block))))
        field_[key] = (block, where)
    return field_[key]


CANDIDATE_ROW = np.full(len(table), -1)
CANDIDATE_ROW[candidates] = np.arange(len(candidates))


class Network:
    """The county field given the monitors picked so far.

    The field is never written to. A pick appends one row of the low-rank factor `u`, and each
    draw's mean stays put, because the reading is unknown. So only the first round is an exact
    value: `joint_value` measures what a whole set is worth.
    """

    def __init__(self, field_, weight=None, line=LINE, objective="wrong", nodes=16, rounds=20,
                 noise_scale=1.0, keep=slice(None), reach=3.0):
        self.field, self.keep = field_, keep
        self.mu = field_["means"][keep]
        self.noise = np.float32(field_["noise"][keep] * noise_scale)
        self.weight = np.float32(population if weight is None else weight)
        self.line, self.objective, self.nodes = np.float32(np.log(line)), objective, nodes
        self.radius = reach * field_["ell"]
        self.u = np.zeros((self.mu.shape[0], rounds, self.mu.shape[1]), np.float32)
        self.k = 0
        block, self.where = open_block(field_, line)
        self.block = block[keep]
        self.var = field_["eta2"][keep][:, None] - (field_["half"][keep] ** 2).sum(axis=1)
        self.refresh()

    def refresh(self):
        self.p = ndtr((self.mu - self.line) / np.sqrt(np.maximum(self.var, 1e-10))).mean(axis=0)
        self.open = (self.p > 0.01) & (self.p < 0.99)

    def column(self, c):
        """Covariance of county `c` with every county, after the picks so far: slower, so for picks."""
        out = covariance_rows(self.field, [c], self.keep)[:, 0]
        if self.k:
            out -= (self.u[:, : self.k, c][:, None] @ self.u[:, : self.k])[:, 0]
        return out

    def value(self, c):
        """Expected fall in the loss from one more monitor at county `c` (EVSI)."""
        slots = np.flatnonzero((gap32[c, self.where] < self.radius) & self.open[self.where])
        if slots.size == 0:
            return 0.0
        local = self.where[slots]
        cross = self.block[:, CANDIDATE_ROW[c]][:, slots]
        if self.k:
            cross = cross - (self.u[:, : self.k, c][:, None] @ self.u[:, : self.k][:, :, local])[:, 0]
        var, mu, mu_c, q = self.var[:, local], self.mu[:, local], self.mu[:, c], self.var[:, c] + self.noise
        y = reading_nodes(mu_c, q, self.nodes)
        p = probability_above(y, update_weights(y, mu_c, q), mu, var, cross, mu_c, q, self.line)
        before = ndtr((mu - self.line) / np.sqrt(var)).mean(axis=0)
        return float(loss(before, self.weight[local], self.objective)
                     - loss(p, self.weight[local], self.objective).mean())

    def score(self, chosen):
        with ThreadPoolExecutor(WORKERS) as pool:
            return np.array(list(pool.map(self.value, chosen)))

    def pick(self, c):
        col = self.column(c)
        self.u[:, self.k] = col / np.sqrt(col[:, c] + self.noise)[:, None]
        self.var = self.var - self.u[:, self.k] ** 2
        self.k += 1
        self.refresh()


net = Network(field)
with stage("first round, every candidate"):
    first_value = np.full(len(table), np.nan)
    first_value[candidates] = net.score(candidates)
ranked = candidates[np.argsort(-first_value[candidates])]
print(f"Best single monitor: {short_name[ranked[0]]}, {first_value[ranked[0]] / 1e3:.0f} thousand people. "
      f"Median candidate: {np.median(first_value[candidates]) / 1e3:.1f} thousand")
print(f"Candidates with EVSI under 1,000 people: {(first_value[candidates] < 1e3).mean():.0%}")