"""Rule-based habitat mapping with Sentinel-2 spectral indices, and an honest accuracy assessment.

Part 1 of the "AI habitat classification with Python" series (Dark Peak, Peak District, UK).
Builds 2023 seasonal Sentinel-2 composites on the shared 10 m British National Grid, computes spectral
indices, fits the thresholds of an ecologically ordered decision list on spatial training blocks,
maps nine habitat classes, and assesses the map against Living England 2022-23 on held-out blocks:
confusion matrix, user's and producer's accuracy, kappa (with its caveats), quantity and allocation
disagreement, block-bootstrap confidence intervals and why random splits overstate accuracy.
Run with:  python habitat_rules.py
"""
import platform
import sys
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Patch, Rectangle
from scipy.spatial import cKDTree

HERE = Path(__file__).resolve().parent
for p in (HERE, HERE.parent / "ai-habitat-classification-common"):   # habitat_common.py: shared series module
    sys.path.insert(0, str(p))
import habitat_common as hc  # noqa: E402

OUT = HERE / "output"
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 160)
pd.set_option("display.max_columns", 20)
KEYS = [c["key"] for c in hc.CLASSES]
NAMES = {c["key"]: c["name"] for c in hc.CLASSES}
COLOURS = {c["key"]: c["colour"] for c in hc.CLASSES}
KEY_ID = {c["key"]: c["id"] for c in hc.CLASSES}


def section(title):
    print(f"\n=== {title} ===")


print("=== versions ===")
print(f"python {platform.python_version()} | numpy {np.__version__} | pandas {pd.__version__} | "
      f"odc-stac {version('odc-stac')} | rasterio {version('rasterio')} | geopandas {version('geopandas')} | "
      f"matplotlib {matplotlib.__version__}")

# --- Step 1: study area, labels and the shared sample design -----------------------------------------
section("Step 1: study area, labels and samples")
print(f"AOI: {hc.AOI_NAME}")
print(f"grid: {hc.CRS}, {hc.RES} m, bounds {hc.AOI_BOUNDS}, {hc.shape()[1]} x {hc.shape()[0]} pixels")
labels = hc.label_raster()                       # Living England 2022-23 on the 10 m grid (0 = unlabelled)
samples = hc.load_samples()
blocks = hc.load_blocks()
print(f"{len(samples):,} samples in {samples.block_id.nunique()} blocks of {hc.BLOCK_SIZE / 1000:g} km; "
      f"{(samples.split == 'train').sum():,} train / {(samples.split == 'test').sum():,} test")

# --- Step 2: seasonal Sentinel-2 composites ----------------------------------------------------------
section("Step 2: seasonal Sentinel-2 L2A composites, 2023")
dates = hc.s2_dates()
print(f"{len(dates)} dates with at least 5% of the AOI clear; median clear fraction {dates.median():.0%}")
S = hc.s2_seasonal_composites()
for season, ds in S.items():
    n = ds["n_clear"].values
    print(f"{season:<7} median {np.median(n):.0f} clear views; none at {(n == 0).mean():.2%} of pixels")


# --- Step 3: spectral indices ------------------------------------------------------------------------
def nd(a, b):
    with np.errstate(divide="ignore", invalid="ignore"):
        return (a - b) / (a + b)


def features(S):
    """Index and band layers used by the rules: 2-D float32 arrays on the shared grid."""
    f = {}
    for s, ds in S.items():
        b = {k: ds[k].values for k in hc.S2_BANDS}
        f[f"ndvi_{s}"] = nd(b["B08"], b["B04"])     # greenness
        f[f"ndmi_{s}"] = nd(b["B08"], b["B11"])     # canopy and soil moisture
        f[f"ndwi_{s}"] = nd(b["B03"], b["B08"])     # open water (McFeeters)
        f[f"nir_{s}"], f[f"red_{s}"] = b["B08"], b["B04"]
        f[f"swir_{s}"], f[f"swir2_{s}"] = b["B11"], b["B12"]
    f["dndvi_spring_summer"] = f["ndvi_summer"] - f["ndvi_spring"]   # late green-up
    return f


section("Step 3: spectral indices")
FEAT = features(S)
X = pd.DataFrame({k: hc.sample_values(v, samples) for k, v in FEAT.items()})
X["cls"], X["split"], X["fold"], X["block_id"] = samples.class_key, samples.split, samples.fold, samples.block_id
print(f"{len(FEAT)} feature layers; samples with any missing value: {X[list(FEAT)].isna().any(axis=1).sum()}")
SHOW = ["ndvi_spring", "ndvi_summer", "dndvi_spring_summer", "ndmi_spring", "nir_summer", "swir_summer",
        "red_summer", "red_spring"]
print("median of each feature by class (training samples only):")
med = X[X.split == "train"].groupby("cls")[SHOW].median().reindex(KEYS).round(3).T
med.columns = [k.split("_")[0][:5] for k in KEYS]
print(med.to_string())

# --- Step 4: the decision list -----------------------------------------------------------------------
# Order matters: the easiest, most distinct classes are peeled off first. Each node lists a few
# ecologically sensible features; the training data pick the feature and threshold with the best F1.
NODES = [
    ("water", ["nir_summer", "ndwi_summer"]),             # absorbs near-infrared
    ("built_bare", ["ndvi_summer", "ndvi_spring"]),       # little or no vegetation
    ("conifer", ["swir_summer", "swir2_summer"]),         # dark, shaded evergreen
    ("broadleaf", ["red_summer", "ndmi_summer", "ndvi_summer", "swir2_summer"]),
    ("bracken", ["dndvi_spring_summer", "ndvi_spring"]),  # brown spring, lush summer
    ("improved_grass", ["ndvi_spring", "nir_spring", "ndmi_spring", "ndvi_winter"]),
    ("seminat_grass", ["ndmi_spring", "nir_spring", "ndvi_spring"]),
    ("bog", ["red_spring", "ndvi_spring", "nir_summer"]), # pale cotton-grass, bare peat
]
DEFAULT = "heath"   # whatever is left: dark dwarf-shrub heath


def best_threshold(x, pos):
    """Threshold and direction that maximise F1 for 'pos' against everything else."""
    ok = ~np.isnan(x)
    x, pos = x[ok], pos[ok]
    best = (-1.0, np.nan, "<")
    for sign, op in ((1, "<"), (-1, ">")):
        v = sign * x
        o = np.argsort(v, kind="stable")
        vs, p = v[o], pos[o]
        tp, fp = np.cumsum(p), np.cumsum(~p)
        f1 = 2 * tp / (tp + fp + p.sum())
        f1[:-1][np.diff(vs) == 0] = -1          # only cut between distinct values
        i = int(np.argmax(f1[:-1]))
        if f1[i] > best[0]:
            best = (float(f1[i]), float(sign * (vs[i] + vs[i + 1]) / 2), op)
    return best


def fit_rules(train):
    rules, rest = [], train
    for cls, cands in NODES:
        pos = (rest["cls"] == cls).to_numpy()
        scored = [(best_threshold(rest[c].to_numpy(), pos), c) for c in cands]
        (f1, t, op), feat = max(scored, key=lambda s: s[0][0])
        rules.append({"cls": cls, "feature": feat, "op": op, "threshold": t, "train_f1": f1,
                      "n_reaching_node": len(rest)})
        hit = (rest[feat] < t) if op == "<" else (rest[feat] > t)
        rest = rest[~hit.fillna(False)]
    return pd.DataFrame(rules)


def apply_rules(rules, get):
    """Apply the decision list; get(feature) returns an array. Returns class keys (object array)."""
    first = get(rules.feature.iloc[0])
    out = np.full(first.shape, DEFAULT, dtype=object)
    done = np.zeros(first.shape, bool)
    for r in rules.itertuples():
        v = get(r.feature)
        with np.errstate(invalid="ignore"):
            hit = (v < r.threshold) if r.op == "<" else (v > r.threshold)
        hit &= ~done
        out[hit] = r.cls
        done |= hit
    return out


section("Step 4: fit the decision-list thresholds on the training blocks")
train, test = X[X.split == "train"], X[X.split == "test"]
rules = fit_rules(train)
for i, r in enumerate(rules.itertuples(), 1):
    rule = f"{r.feature} {r.op} {r.threshold:.3f}"
    print(f"{i}. {NAMES[r.cls]:<23} if {rule:<30} (train F1 {r.train_f1:.2f}, n = {r.n_reaching_node:,})")
print(f"9. otherwise              -> {NAMES[DEFAULT]}")

# --- Step 5: classify the whole AOI -----------------------------------------------------------------
section("Step 5: classify the whole study area")
pred_key = apply_rules(rules, lambda f: FEAT[f])
pred = np.zeros(hc.shape(), "uint8")
for k in KEYS:
    pred[pred_key == k] = KEY_ID[k]
nodata = np.isnan(FEAT["ndvi_summer"]) | np.isnan(FEAT["ndvi_spring"])
pred[nodata] = 0
hc.write_geotiff(OUT / "habitat_rules_2023.tif", pred, nodata=0)
area = pd.Series({NAMES[k]: (pred == KEY_ID[k]).sum() * hc.RES ** 2 / 1e6 for k in KEYS})
ref_area = pd.Series({NAMES[k]: (labels == KEY_ID[k]).sum() * hc.RES ** 2 / 1e6 for k in KEYS})
print(pd.DataFrame({"rules_km2": area.round(1), "living_england_km2": ref_area.round(1)}).to_string())
print(f"unclassified (no clear spring or summer view): {nodata.mean():.2%} of pixels")

# --- Step 6: accuracy assessment on the test blocks --------------------------------------------------
section("Step 6: accuracy on the held-out test blocks")
yt = test["cls"].to_numpy()
yp = apply_rules(rules, lambda f: test[f].to_numpy())


def confusion(y_true, y_pred):
    return pd.crosstab(pd.Categorical(y_true, KEYS), pd.Categorical(y_pred, KEYS), dropna=False)


def summarise(cm):
    m = cm.to_numpy().astype(float)
    n = m.sum()
    oa = np.trace(m) / n
    pe = (m.sum(0) * m.sum(1)).sum() / n ** 2
    kappa = (oa - pe) / (1 - pe)
    pa = np.diag(m) / m.sum(1)                   # producer's accuracy (recall): rows are reference
    ua = np.diag(m) / np.where(m.sum(0) == 0, np.nan, m.sum(0))   # user's accuracy (precision)
    f1 = 2 * pa * ua / (pa + ua)
    p = m / n
    quantity = 0.5 * np.abs(p.sum(0) - p.sum(1)).sum()
    return {"oa": oa, "kappa": kappa, "pa": pa, "ua": ua, "f1": f1, "quantity": quantity,
            "allocation": (1 - oa) - quantity}


cm = confusion(yt, yp)
res = summarise(cm.astype(float))
cm_show = cm.copy()
cm_show.index = [f"ref {k}" for k in KEYS]
cm_show.columns = [k.split("_")[0][:5] for k in KEYS]   # short headers keep the table narrow
print("confusion matrix (rows = Living England reference, columns = rule-based map):")
print(cm_show.to_string())
per_class = pd.DataFrame({"n_test": cm.sum(axis=1).to_numpy(), "producers_acc": res["pa"], "users_acc": res["ua"],
                          "f1": res["f1"]}, index=[NAMES[k] for k in KEYS])
print(per_class.round(3).to_string())
n_test = len(yt)
naive = 1.96 * np.sqrt(res["oa"] * (1 - res["oa"]) / n_test)
rng = np.random.default_rng(hc.SEED)
tb = test["block_id"].to_numpy()
ub = np.unique(tb)
correct = (yt == yp)
boot = []
for _ in range(2000):                             # resample whole test blocks, not pixels
    pick = rng.choice(ub, ub.size, replace=True)
    idx = np.concatenate([np.flatnonzero(tb == b) for b in pick])
    boot.append(correct[idx].mean())
lo, hi = np.percentile(boot, [2.5, 97.5])
print(f"overall accuracy {res['oa']:.3f} (n = {n_test:,} samples in {ub.size} test blocks)")
print(f"  95% CI treating samples as independent: {res['oa'] - naive:.3f} to {res['oa'] + naive:.3f}")
print(f"  95% CI from a block bootstrap:          {lo:.3f} to {hi:.3f}")
print(f"macro F1 {np.nanmean(res['f1']):.3f}; Cohen's kappa {res['kappa']:.3f}")
print(f"disagreement {1 - res['oa']:.3f} = quantity {res['quantity']:.3f} + allocation {res['allocation']:.3f}")
# The sample is balanced (1,000 per class), so sample OA is not the map's accuracy. Weight each class's
# producer's accuracy by its share of the AOI in the reference map to estimate area-weighted agreement.
w = np.array([(labels == KEY_ID[k]).sum() for k in KEYS], float)
w /= w.sum()
print(f"area-weighted agreement (sum of reference-area share x producer's accuracy): {np.nansum(w * res['pa']):.3f}")
test_blocks = blocks[blocks.split == "test"]
in_test = np.isin(hc.block_id(*np.meshgrid(hc.AOI_BOUNDS[0] + (np.arange(hc.shape()[1]) + 0.5) * hc.RES,
                                            hc.AOI_BOUNDS[3] - (np.arange(hc.shape()[0]) + 0.5) * hc.RES)),
                  test_blocks.block_id.to_numpy())
valid = in_test & (labels > 0) & (pred > 0)
print(f"wall-to-wall agreement with Living England over all {valid.sum():,} labelled test-block pixels: "
      f"{(labels[valid] == pred[valid]).mean():.3f}")
both = (labels > 0) & (pred > 0)
print(f"wall-to-wall agreement over the whole study area (train and test blocks): {(labels[both] == pred[both]).mean():.3f}")

section("Step 6b: blocked 5-fold cross-validation on the training blocks")
cv = []
for k in range(5):
    r = fit_rules(train[train.fold != k])
    v = train[train.fold == k]
    cv.append((apply_rules(r, lambda f: v[f].to_numpy()) == v["cls"].to_numpy()).mean())
print("fold OA: " + ", ".join(f"{a:.3f}" for a in cv) + f"; mean {np.mean(cv):.3f} +/- {np.std(cv):.3f}")

# --- Step 7: why a random split overstates accuracy -------------------------------------------------
section("Step 7: random split versus spatial blocks")
xy = samples[["x", "y"]].to_numpy()
is_test = (samples.split == "test").to_numpy()
d_block = cKDTree(xy[~is_test]).query(xy[is_test])[0]
rand_test = np.random.default_rng(hc.SEED).random(len(samples)) < is_test.mean()
d_rand = cKDTree(xy[~rand_test]).query(xy[rand_test])[0]
print(f"median distance from test to nearest training sample: random split {np.median(d_rand):.0f} m, "
      f"block split {np.median(d_block):.0f} m")


def nn_coords(tr_idx, te_idx):
    """A 'model' that sees only location: the label of the nearest training sample."""
    j = cKDTree(xy[tr_idx]).query(xy[te_idx])[1]
    return (X["cls"].to_numpy()[tr_idx][j] == X["cls"].to_numpy()[te_idx]).mean()


REPEATS = 20
rows = []
ncols = int(np.ceil((hc.AOI_BOUNDS[2] - hc.AOI_BOUNDS[0]) / hc.BLOCK_SIZE))
all_blocks = blocks.block_id.to_numpy()
for rep in range(REPEATS):
    g = np.random.default_rng(1000 + rep)
    splits = {"random": g.random(len(samples)) < 0.3}
    tb_rep = g.choice(all_blocks, int(round(0.3 * all_blocks.size)), replace=False)
    splits["blocks"] = np.isin(samples.block_id.to_numpy(), tb_rep)
    for name, te_mask in splits.items():
        r = fit_rules(X[~te_mask])
        oa_rules = (apply_rules(r, lambda f: X.loc[te_mask, f].to_numpy()) == X.loc[te_mask, "cls"].to_numpy()).mean()
        rows.append({"split": name, "rep": rep, "rules": oa_rules,
                     "nearest_neighbour_xy": nn_coords(np.flatnonzero(~te_mask), np.flatnonzero(te_mask))})
rep_df = pd.DataFrame(rows)
summary = rep_df.groupby("split")[["rules", "nearest_neighbour_xy"]].agg(["mean", "std"])
print(f"overall accuracy over {REPEATS} repeated 70/30 splits:")
print(summary.round(3).to_string())

# --- Figures -----------------------------------------------------------------------------------------
ext = (hc.AOI_BOUNDS[0], hc.AOI_BOUNDS[2], hc.AOI_BOUNDS[1], hc.AOI_BOUNDS[3])
rgb = np.dstack([S["summer"][b].values for b in ("B04", "B03", "B02")])
rgb = np.nan_to_num(np.clip(rgb / np.nanpercentile(rgb, 99), 0, 1) ** 0.8, nan=1.0)
cmap, norm = hc.class_cmap()
legend = [Patch(color=COLOURS[k], label=NAMES[k]) for k in KEYS]
PLACES = {"Kinder Scout": (408600, 388400), "Bleaklow": (409800, 395800), "Edale": (412300, 385700),
          "Castleton": (415000, 382700), "Glossop": (403300, 394000), "Ladybower Reservoir": (420000, 386800),
          "Derwent Valley": (417700, 393000), "Hathersage": (423200, 381400)}


def km_axes(ax):
    ax.set_xticks(np.arange(400000, 425001, 5000))
    ax.set_yticks(np.arange(380000, 400001, 5000))
    ax.set_xticklabels([f"{v / 1000:.0f}" for v in ax.get_xticks()])
    ax.set_yticklabels([f"{v / 1000:.0f}" for v in ax.get_yticks()])
    ax.set_xlabel("Easting (km, British National Grid)")
    ax.set_ylabel("Northing (km)")


def place_labels(ax, colour="white"):
    for name, (x, y) in PLACES.items():
        ax.plot(x, y, "o", ms=3, color=colour, mec="black", mew=0.5)
        ax.annotate(name, (x, y), xytext=(4, 3), textcoords="offset points", fontsize=7.5, color=colour,
                    fontweight="bold", path_effects=[matplotlib.patheffects.withStroke(linewidth=2, foreground="black")])


import matplotlib.patheffects  # noqa: E402

# Figure 1: true colour with the spatial blocks
fig, ax = plt.subplots(figsize=(8, 7.4))
ax.imshow(rgb, extent=ext)
for b in blocks.itertuples():
    x0, y0, x1, y1 = b.geometry.bounds
    if b.split == "test":
        ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, ec="#fdae61", lw=1.6, hatch="////", alpha=0.9))
    else:
        ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, ec="white", lw=0.4, alpha=0.6))
place_labels(ax)
km_axes(ax)
ax.legend(handles=[Patch(fill=False, ec="#fdae61", hatch="////", label="Test block (27)"),
                   Patch(fill=False, ec="white", label="Training block (63)")],
          loc="lower left", fontsize=8, framealpha=0.9)
ax.set_title("Dark Peak study area: Sentinel-2 summer 2023 median composite, 2.5 km spatial blocks")
fig.savefig(OUT / "01-study-area-blocks.png")
plt.close(fig)

# Figure 2: Living England labels and the samples
fig, axes = plt.subplots(1, 2, figsize=(13, 5.6), sharey=True)
axes[0].imshow(np.ma.masked_equal(labels, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[0].set_title("Living England 2022-23, mapped to nine classes")
axes[1].imshow(rgb.mean(axis=2), cmap="gray", extent=ext, vmin=0, vmax=1.4)
for k in KEYS:
    m = samples.class_key == k
    axes[1].scatter(samples.x[m], samples.y[m], s=2.5, color=COLOURS[k], linewidths=0)
axes[1].set_title("9,000 sample pixels (1,000 per class, high-reliability polygons)")
for ax in axes:
    km_axes(ax)
axes[1].set_ylabel("")
fig.legend(handles=legend, loc="lower center", ncol=9, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.04))
fig.savefig(OUT / "02-labels-and-samples.png")
plt.close(fig)

# Figure 3: class signatures for the features the rules chose
chosen = list(dict.fromkeys(rules.feature))
fig, axes = plt.subplots(2, 4, figsize=(13, 6.4))
for ax, r in zip(axes.flat, rules.itertuples()):
    data = [train.loc[train.cls == k, r.feature].dropna() for k in KEYS]
    bp = ax.boxplot(data, orientation="horizontal", showfliers=False, patch_artist=True, widths=0.6,
                    medianprops={"color": "black"})
    for patch, k in zip(bp["boxes"], KEYS):
        patch.set_facecolor(COLOURS[k])
        patch.set_edgecolor("#333333")
        patch.set_linewidth(2.2 if k == r.cls else 0.6)
    ax.axvline(r.threshold, color="#d7191c", lw=1.4, ls="--")
    ax.set_yticks(range(1, len(KEYS) + 1))
    ax.set_yticklabels([NAMES[k] for k in KEYS], fontsize=7)
    ax.invert_yaxis()
    ax.set_title(f"{NAMES[r.cls]}: {r.feature} {r.op} {r.threshold:.3f}", fontsize=8.5)
    ax.tick_params(axis="x", labelsize=7)
for ax in axes[:, 1:].flat:
    ax.set_yticklabels([])
fig.suptitle("Training-sample distributions of the feature chosen at each rule (box = interquartile range, "
             "whiskers = 1.5 IQR); dashed red line = fitted threshold", fontsize=9)
fig.tight_layout()
fig.savefig(OUT / "03-rule-features.png")
plt.close(fig)

# Figure 4: rule-based map against Living England
fig, axes = plt.subplots(1, 2, figsize=(13, 5.6), sharey=True)
axes[0].imshow(np.ma.masked_equal(pred, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[0].set_title("Rule-based map from 2023 Sentinel-2 indices")
agree = np.where((labels > 0) & (pred > 0), (labels == pred).astype(float), np.nan)
axes[1].imshow(np.ma.masked_equal(labels, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[1].set_title(f"Living England 2022-23 (pixel agreement with the rules: {np.nanmean(agree):.0%})")
for ax in axes:
    km_axes(ax)
axes[1].set_ylabel("")
fig.legend(handles=legend, loc="lower center", ncol=9, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.04))
fig.savefig(OUT / "04-rule-map-vs-living-england.png")
plt.close(fig)

# Figure 5: confusion matrix with user's and producer's accuracy
m = cm.to_numpy()
rown = m / m.sum(1, keepdims=True)
fig, ax = plt.subplots(figsize=(8.6, 7.4))
ax.imshow(rown, cmap="YlGnBu", vmin=0, vmax=1)
for i in range(len(KEYS)):
    for j in range(len(KEYS)):
        if m[i, j]:
            ax.text(j, i, m[i, j], ha="center", va="center", fontsize=8, color="white" if rown[i, j] > 0.55 else "black")
short = [NAMES[k].replace(" woodland", "").replace(" grassland", " grass") for k in KEYS]
ax.set_xticks(range(len(KEYS)), [f"{s}\nUA {u:.0%}" for s, u in zip(short, res["ua"])], rotation=45, ha="right",
              fontsize=7.5)
ax.set_yticks(range(len(KEYS)), [f"{s}  PA {p:.0%}" for s, p in zip(short, res["pa"])], fontsize=7.5)
ax.set_xlabel("Rule-based map (UA = user's accuracy)")
ax.set_ylabel("Living England reference (PA = producer's accuracy)")
ax.set_title(f"Test blocks: overall accuracy {res['oa']:.1%} (block-bootstrap 95% CI {lo:.1%} to {hi:.1%}), "
             f"kappa {res['kappa']:.2f}", fontsize=9)
fig.savefig(OUT / "05-confusion-matrix.png")
plt.close(fig)

# Figure 6: random versus spatial split
fig, axes = plt.subplots(1, 2, figsize=(12, 4.4))
bins = np.linspace(0, 1500, 61)
axes[0].hist(d_rand, bins=bins, color="#d7191c", alpha=0.75, label=f"Random split (median {np.median(d_rand):.0f} m)")
axes[0].hist(np.clip(d_block, 0, 1500), bins=bins, color="#2b83ba", alpha=0.75,
             label=f"Block split (median {np.median(d_block):.0f} m)")
axes[0].set_xlabel("Distance from each test sample to the nearest training sample (m; last bin = 1,500+)")
axes[0].set_ylabel("Test samples")
axes[0].legend(fontsize=8)
pos = np.arange(2)
for i, (name, col) in enumerate((("random", "#d7191c"), ("blocks", "#2b83ba"))):
    vals = rep_df[rep_df.split == name]
    means = [vals.rules.mean(), vals.nearest_neighbour_xy.mean()]
    sds = [vals.rules.std(), vals.nearest_neighbour_xy.std()]
    bars = axes[1].bar(pos + (i - 0.5) * 0.36, means, 0.36, yerr=sds, capsize=3, color=col,
                       label="Random split" if name == "random" else "Spatial block split")
    for bx, v in zip(bars, means):
        axes[1].text(bx.get_x() + bx.get_width() / 2, v + 0.03, f"{v:.0%}", ha="center", fontsize=8)
axes[1].set_xticks(pos, ["Spectral rules\n(this tutorial)", "Nearest training sample\n(coordinates only, no imagery)"])
axes[1].set_ylabel(f"Overall accuracy (mean and s.d. of {REPEATS} splits)")
axes[1].set_ylim(0, 1)
axes[1].legend(fontsize=8, loc="upper left")
fig.tight_layout()
fig.savefig(OUT / "06-random-vs-block-split.png")
plt.close(fig)
print(f"\nsaved 6 figures and habitat_rules_2023.tif to {OUT.name}/")