Source code: AlphaEarth satellite embeddings in Python: few-shot habitat mapping, uncertainty and area estimates

"""Read AlphaEarth Foundations annual satellite embeddings for the series grid from Source Cooperative.

The public COG mirror (s3://us-west-2.opendata.source.coop/tge-labs/aef/, HTTPS at data.source.coop) needs no
login. Each COG is 8192 x 8192 pixels in a UTM zone, 64 int8 bands (A00-A63), stored "bottom-up" (positive
y pixel size). We read just the window over the study area, flip it to the usual top-down order, reproject
it onto the shared 10 m British National Grid with nearest-neighbour resampling and cache it as one GeoTIFF.
"""
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import rasterio
import requests
from affine import Affine
from pyproj import Transformer
from rasterio.warp import Resampling, reproject
from rasterio.windows import Window

BASE = "https://data.source.coop/tge-labs/aef/v1/annual"
INDEX_URL = f"{BASE}/aef_index.parquet"
NODATA = -128
GDAL_OPTS = {"GDAL_DISABLE_READDIR_ON_OPEN": "EMPTY_DIR", "GDAL_HTTP_MULTIRANGE": "YES",
             "GDAL_HTTP_MERGE_CONSECUTIVE_RANGES": "YES", "GDAL_HTTP_MAX_RETRY": "5", "GDAL_HTTP_RETRY_DELAY": "2",
             "VSI_CACHE": "TRUE", "CPL_VSIL_CURL_ALLOWED_EXTENSIONS": ".tiff"}


def dequantise(raw):
    """int8 -> float32 in [-1, 1]: (v / 127.5)^2 * sign(v); -128 (no data) -> NaN."""
    v = raw.astype(np.float32)
    out = (v / 127.5) ** 2 * np.sign(v)
    out[raw == NODATA] = np.nan
    return out


def find_tiles(bbox_wgs84, year, cache_dir):
    """Rows of the AEF index (one per COG) that intersect a lon/lat box for one year."""
    f = Path(cache_dir) / "aef_index.parquet"
    if not f.exists():
        f.parent.mkdir(parents=True, exist_ok=True)
        with requests.get(INDEX_URL, stream=True, timeout=300) as r:   # 78 MB, downloaded once
            r.raise_for_status()
            with open(f, "wb") as fh:
                for chunk in r.iter_content(1 << 20):
                    fh.write(chunk)
    cols = ["path", "crs", "year", "utm_zone", "utm_west", "utm_south", "utm_east", "utm_north",
            "wgs84_west", "wgs84_south", "wgs84_east", "wgs84_north"]
    df = pq.read_table(f, columns=cols, filters=[("year", "=", year)]).to_pandas()
    w, s, e, n = bbox_wgs84
    hit = (df.wgs84_west < e) & (df.wgs84_east > w) & (df.wgs84_south < n) & (df.wgs84_north > s)
    df = df[hit].copy()
    df["url"] = df.path.str.replace("s3://us-west-2.opendata.source.coop/tge-labs/aef/v1/annual",
                                    BASE, regex=False)
    return df.reset_index(drop=True)


def load_embeddings(tiles, dst_shape, dst_transform, dst_crs, margin=200):
    """Mosaic the AEF tiles onto a target grid. Returns int8 (64, rows, cols), NODATA where not covered."""
    rows, cols = dst_shape
    out = np.full((64, rows, cols), NODATA, dtype=np.int8)
    x0, y1 = dst_transform * (0, 0)
    x1, y0 = dst_transform * (cols, rows)
    with rasterio.Env(**GDAL_OPTS):
        for t in tiles.itertuples():
            to_utm = Transformer.from_crs(dst_crs, t.crs, always_xy=True)
            xs, ys = to_utm.transform([x0, x1, x0, x1], [y0, y0, y1, y1])
            west, east = max(min(xs) - margin, t.utm_west), min(max(xs) + margin, t.utm_east)
            south, north = max(min(ys) - margin, t.utm_south), min(max(ys) + margin, t.utm_north)
            if west >= east or south >= north:
                continue
            with rasterio.open(f"/vsicurl/{t.url}") as src:
                tr = src.transform                       # bottom-up: origin at the south-west corner, +10 m in y
                c0, c1 = int((west - tr.c) // tr.a), int(np.ceil((east - tr.c) / tr.a))
                r0, r1 = int((south - tr.f) // tr.e), int(np.ceil((north - tr.f) / tr.e))
                win = Window(c0, r0, c1 - c0, r1 - r0)
                raw = src.read(window=win)[:, ::-1, :]   # flip rows so that row 0 is the northern edge
            src_tr = Affine(tr.a, 0, tr.c + c0 * tr.a, 0, -tr.e, tr.f + r1 * tr.e)
            part = np.full_like(out, NODATA)
            reproject(raw, part, src_transform=src_tr, src_crs=t.crs, src_nodata=NODATA,
                      dst_transform=dst_transform, dst_crs=dst_crs, dst_nodata=NODATA,
                      resampling=Resampling.nearest)
            fill = (out[0] == NODATA) & (part[0] != NODATA)
            out[:, fill] = part[:, fill]
    return out


def tile_table(tiles):
    return pd.DataFrame({"file": tiles.url.str.split("/").str[-1], "zone": tiles.utm_zone,
                         "utm_west": tiles.utm_west, "utm_east": tiles.utm_east})

"""Geospatial foundation model embeddings in practice: AlphaEarth Foundations for habitat mapping.

Part 5 of the "AI habitat classification with Python" series (Dark Peak, Peak District, UK).
Reads Google DeepMind's AlphaEarth Foundations 2023 annual satellite embeddings (64 dimensions, 10 m) from the
public Source Cooperative mirror, classifies nine habitats with a linear classifier, and compares them with
raw seasonal Sentinel-2 bands: a few-shot learning curve, uncertainty and calibration, simulated active
learning, and area-adjusted accuracy with class areas and 95% confidence intervals (Olofsson et al. 2014)
from a stratified random sample of the map.
Run with:  python alphaearth_embeddings.py
"""
import json
import platform
import sys
import time
from importlib.metadata import version
from pathlib import Path

T_START = time.perf_counter()

import matplotlib  # noqa: E402

matplotlib.use("Agg")
import matplotlib.patheffects  # noqa: E402
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
import rasterio  # noqa: E402
from lightgbm import LGBMClassifier  # noqa: E402
from matplotlib.patches import Patch  # noqa: E402
from sklearn.decomposition import PCA  # noqa: E402
from sklearn.impute import SimpleImputer  # noqa: E402
from sklearn.linear_model import LogisticRegression  # noqa: E402
from sklearn.neighbors import KNeighborsClassifier  # noqa: E402
from sklearn.pipeline import make_pipeline  # noqa: E402
from sklearn.preprocessing import StandardScaler  # noqa: E402

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 aef  # noqa: E402  (aef.py: AlphaEarth reader, next to this script)
import habitat_common as hc  # noqa: E402

OUT = HERE / "output"
DATA = HERE / "data"
OUT.mkdir(exist_ok=True)
DATA.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]
IDS = np.array([c["id"] 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}
RED, ORANGE, GREEN, BLUE, PURPLE = "#d7191c", "#fdae61", "#1a9641", "#2b83ba", "#5e3c99"
PART1_OA = 0.600


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


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

# --- Step 1: AlphaEarth embeddings on the series grid ------------------------------------------------
section("Step 1: AlphaEarth Foundations 2023 embeddings from Source Cooperative")
AEF_TIF = DATA / f"aef_{hc.YEAR}_bng.tif"
tiles = aef.find_tiles(hc.aoi_bbox_wgs84(), hc.YEAR, DATA)
print(f"{len(tiles)} COGs cover the study area:")
print(aef.tile_table(tiles).to_string(index=False))
if not AEF_TIF.exists():
    t0 = time.perf_counter()
    raw = aef.load_embeddings(tiles, hc.shape(), hc.transform(), hc.CRS)
    hc.write_geotiff(AEF_TIF, raw, nodata=aef.NODATA, dtype="int8")
    print(f"read the study-area windows and reprojected to EPSG:27700 in {time.perf_counter() - t0:.0f} s")
with rasterio.open(AEF_TIF) as src:
    raw = src.read()                                     # int8 (64, 2250, 2500)
valid = raw[0] != aef.NODATA
print(f"embedding cube {raw.shape} int8, {raw.nbytes / 1e6:.0f} MB; valid pixels {valid.mean():.2%}")
samples = hc.load_samples()
XE = aef.dequantise(hc.sample_values(raw, samples))     # (9000, 64) float32 in [-1, 1]
norms = np.linalg.norm(XE, axis=1)
print(f"sample embeddings {XE.shape}; vector length min {norms.min():.3f}, median {np.median(norms):.3f}, "
      f"max {norms.max():.3f} (unit vectors up to 8-bit quantisation)")

# The comparison: the same samples described by raw 2023 seasonal Sentinel-2 bands (4 seasons x 10 bands)
S = hc.s2_seasonal_composites()
S2_COLS = [f"{b}_{s}" for s in S for b in hc.S2_BANDS]
XS = np.column_stack([hc.sample_values(S[s][b].values, samples) for s in S for b in hc.S2_BANDS])
print(f"Sentinel-2 seasonal features {XS.shape}; missing values (no clear view) {np.isnan(XS).mean():.2%}")
y = samples.class_id.to_numpy()
is_tr = (samples.split == "train").to_numpy()
is_te = ~is_tr
fold = samples.fold.to_numpy()
FEATS = {"AlphaEarth embeddings (64)": XE, "Sentinel-2 seasonal bands (40)": XS}


def logreg(C=1.0):
    return make_pipeline(SimpleImputer(strategy="median"), StandardScaler(), LogisticRegression(C=C, max_iter=3000))


# --- Step 2: a linear classifier on all training samples ----------------------------------------------
section("Step 2: tune C with blocked 5-fold CV on the training blocks")
cv_rows = []
for fname, X in FEATS.items():
    for C in (0.01, 0.1, 1.0, 10.0):
        accs = []
        for k in range(5):
            trk, vak = is_tr & (fold != k), is_tr & (fold == k)
            accs.append((logreg(C).fit(X[trk], y[trk]).predict(X[vak]) == y[vak]).mean())
        cv_rows.append({"features": fname, "C": C, "cv_oa": np.mean(accs), "cv_sd": np.std(accs)})
cv_df = pd.DataFrame(cv_rows)
print(cv_df.pivot(index="C", columns="features", values="cv_oa").round(3).to_string())
BEST_C = {f: cv_df[cv_df.features == f].sort_values("cv_oa").iloc[-1].C for f in FEATS}
print("chosen C: " + "; ".join(f"{f}: {c:g}" for f, c in BEST_C.items()))

section("Step 2b: test-block accuracy, all 6,544 training samples")
rows = []
for fname, X in FEATS.items():
    models = {"logistic regression": logreg(BEST_C[fname]),
              "k-nearest neighbours (k=10)": make_pipeline(SimpleImputer(strategy="median"), StandardScaler(),
                                                           KNeighborsClassifier(10)),
              "LightGBM (300 trees)": LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31,
                                                     random_state=hc.SEED, verbose=-1)}
    for mname, m in models.items():
        t0 = time.perf_counter()
        m.fit(X[is_tr], y[is_tr])
        rows.append({"features": fname, "classifier": mname, "test_oa": (m.predict(X[is_te]) == y[is_te]).mean(),
                     "fit_s": time.perf_counter() - t0})
full_df = pd.DataFrame(rows)
print(full_df.pivot(index="classifier", columns="features", values="test_oa").round(3).to_string())

model = logreg(BEST_C["AlphaEarth embeddings (64)"]).fit(XE[is_tr], y[is_tr])
proba_te = model.predict_proba(XE[is_te])
yp = model.classes_[proba_te.argmax(1)]
yt = y[is_te]


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


def summarise(m):
    m = m.astype(float)
    n = m.sum()
    oa = np.trace(m) / n
    pe = (m.sum(0) * m.sum(1)).sum() / n ** 2
    pa = np.diag(m) / m.sum(1)
    ua = np.diag(m) / np.where(m.sum(0) == 0, np.nan, m.sum(0))
    return {"oa": oa, "kappa": (oa - pe) / (1 - pe), "pa": pa, "ua": ua, "f1": 2 * pa * ua / (pa + ua)}


cm = confusion(yt, yp)
res = summarise(cm)
rng = np.random.default_rng(hc.SEED)
tb = samples.block_id.to_numpy()[is_te]
ub = np.unique(tb)
correct = yt == yp
boot = []
for _ in range(2000):                                  # resample whole test blocks, as in Part 1
    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])
cm_show = pd.DataFrame(cm, index=[f"ref {k}" for k in KEYS], columns=[k.split("_")[0][:5] for k in KEYS])
print("confusion matrix (rows = Living England reference, columns = AlphaEarth + logistic regression):")
print(cm_show.to_string())
per_class = pd.DataFrame({"n_test": cm.sum(1), "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())
print(f"overall accuracy {res['oa']:.3f} (n = {len(yt):,}); block-bootstrap 95% CI {lo:.3f} to {hi:.3f}")
print(f"macro F1 {np.nanmean(res['f1']):.3f}; Cohen's kappa {res['kappa']:.3f}; Part 1 rules: OA {PART1_OA:.3f}")

# --- Step 3: wall-to-wall map and per-pixel probabilities ------------------------------------------
section("Step 3: classify every pixel")
rows_, cols_ = hc.shape()
pred = np.zeros((rows_, cols_), np.uint8)
pmax = np.full((rows_, cols_), np.nan, np.float32)
margin = np.full((rows_, cols_), np.nan, np.float32)
entropy = np.full((rows_, cols_), np.nan, np.float32)
t0 = time.perf_counter()
for r0 in range(0, rows_, 250):
    blk = raw[:, r0:r0 + 250, :]
    X = aef.dequantise(blk.reshape(64, -1).T)
    ok = blk[0].ravel() != aef.NODATA
    P = np.zeros((X.shape[0], len(IDS)), np.float32)
    P[ok] = model.predict_proba(X[ok])
    srt = np.sort(P, axis=1)
    with np.errstate(divide="ignore", invalid="ignore"):
        H = -(P * np.log(np.where(P > 0, P, 1))).sum(1) / np.log(len(IDS))
    sl = slice(r0, r0 + blk.shape[1])
    pred[sl] = np.where(ok, model.classes_[P.argmax(1)], 0).reshape(-1, cols_)
    pmax[sl] = np.where(ok, srt[:, -1], np.nan).reshape(-1, cols_)
    margin[sl] = np.where(ok, srt[:, -1] - srt[:, -2], np.nan).reshape(-1, cols_)
    entropy[sl] = np.where(ok, H, np.nan).reshape(-1, cols_)
print(f"predicted {valid.sum():,} pixels in {time.perf_counter() - t0:.0f} s")
hc.write_geotiff(OUT / "habitat_alphaearth_2023.tif", pred, nodata=0)
hc.write_geotiff(OUT / "margin_alphaearth_2023.tif", margin.astype("float32"), nodata=np.nan, dtype="float32")
labels = hc.label_raster()
xc = hc.AOI_BOUNDS[0] + (np.arange(cols_) + 0.5) * hc.RES
yc = hc.AOI_BOUNDS[3] - (np.arange(rows_) + 0.5) * hc.RES
blk_split = hc.load_blocks().set_index("block_id").split
pix_block = hc.block_id(*np.meshgrid(xc, yc))
in_test = np.isin(pix_block, blk_split[blk_split == "test"].index.to_numpy())
both = (labels > 0) & (pred > 0)
agree_all = (labels[both] == pred[both]).mean()
agree_te = (labels[both & in_test] == pred[both & in_test]).mean()
agree_tr = (labels[both & ~in_test] == pred[both & ~in_test]).mean()
print(f"wall-to-wall agreement with Living England: whole area {agree_all:.3f}; "
      f"test blocks {agree_te:.3f}; training blocks {agree_tr:.3f} (Part 1 rules: 0.490 and 0.483)")
print(f"mean entropy {np.nanmean(entropy):.3f}; pixels with margin < 0.2: {np.nanmean(margin < 0.2):.1%}")

# --- Step 4: few-shot learning curve -----------------------------------------------------------------
section("Step 4: few-shot learning curve (labels per class drawn from the training blocks)")
SHOTS = [5, 10, 25, 50, 100, 250]
REPEATS = 20
tr_idx_by_class = {c: np.flatnonzero(is_tr & (y == c)) for c in IDS}
fs_rows = []
for fname, X in FEATS.items():
    for n in SHOTS + ["all"]:
        for rep in range(REPEATS if n != "all" else 1):
            g = np.random.default_rng(10_000 + rep)
            if n == "all":
                sel = np.flatnonzero(is_tr)
            else:
                sel = np.concatenate([g.choice(tr_idx_by_class[c], n, replace=False) for c in IDS])
            for mname, m in (("logistic", logreg(1.0)),
                             ("kNN", make_pipeline(SimpleImputer(strategy="median"), StandardScaler(),
                                                   KNeighborsClassifier(min(5, len(sel) // 9))))):
                m.fit(X[sel], y[sel])
                fs_rows.append({"features": fname, "classifier": mname, "per_class": n, "rep": rep,
                                "test_oa": (m.predict(X[is_te]) == y[is_te]).mean()})
fs = pd.DataFrame(fs_rows)
fs_mean = fs.groupby(["classifier", "per_class", "features"], sort=False).test_oa.mean().unstack(2)
fs_sd = fs[fs.per_class == 5].groupby(["classifier", "features"], sort=False).test_oa.std()
for mname, label in (("logistic", "logistic regression"), ("kNN", "k-nearest neighbours (k=5)")):
    print(f"test OA, {label}, mean of {REPEATS} random draws ('all' = every training sample, about 727 per class):")
    print(fs_mean.loc[mname].round(3).to_string())
print("s.d. across draws with 5 labels per class: " +
      "; ".join(f"{m} / {f.split(' (')[0]} {v:.3f}" for (m, f), v in fs_sd.items()))
lr = fs[fs.classifier == "logistic"].groupby(["features", "per_class"], sort=False).test_oa.mean().unstack(0)
e_col, s_col = list(FEATS)
match = next((n for n in SHOTS if lr.loc[n, e_col] >= lr.loc["all", s_col]), SHOTS[-1])
print(f"AlphaEarth + logistic regression with {match} labels per class ({lr.loc[match, e_col]:.3f}) already matches "
      f"Sentinel-2 bands with all {is_tr.sum():,} training labels ({lr.loc['all', s_col]:.3f})")

# --- Step 5: uncertainty and calibration -------------------------------------------------------------
section("Step 5: uncertainty and calibration")
conf_te = proba_te.max(1)
srt = np.sort(proba_te, 1)
margin_te = srt[:, -1] - srt[:, -2]
q = pd.qcut(margin_te, 5, labels=["lowest", "low", "middle", "high", "highest"])
print("test samples by margin quintile (margin = top probability - second):")
print(pd.DataFrame({"margin": margin_te, "error": ~correct}).groupby(q, observed=True)
      .agg(margin_from=("margin", "min"), margin_to=("margin", "max"), error_rate=("error", "mean"), n=("error", "size"))
      .round(3).to_string())


def reliability(conf, ok, w=None, bins=np.linspace(0, 1, 11), min_n=25):
    """Mean confidence and accuracy per confidence bin (optionally weighted); bins with < min_n samples dropped."""
    w = np.ones(len(conf)) if w is None else w
    b = np.clip(np.digitize(conf, bins) - 1, 0, len(bins) - 2)
    df = pd.DataFrame({"b": b, "cw": conf * w, "okw": ok * w, "w": w})
    df = df.groupby("b").agg(cw=("cw", "sum"), okw=("okw", "sum"), w=("w", "sum"), n=("w", "size"))
    df["conf"], df["acc"] = df.cw / df.w, df.okw / df.w
    ece = (df.w * (df.conf - df.acc).abs()).sum() / df.w.sum()
    return df[df.n >= min_n], ece


rel_te, ece_te = reliability(conf_te, correct)
print(f"expected calibration error on the pure test samples: {ece_te:.3f}")

# --- Step 6: simulated active learning ---------------------------------------------------------------
section("Step 6: active learning, uncertainty (margin) sampling vs random sampling")
pool = np.flatnonzero(is_tr)                      # the training blocks are the unlabelled pool
scaler = StandardScaler().fit(XE[pool])           # unlabelled pool: scaling needs no labels
Zp, Zte = scaler.transform(XE[pool]), scaler.transform(XE[is_te])
yp_pool = y[pool]
START, BATCH, ROUNDS, AL_REPEATS = 2, 18, 20, 10
al_rows = []
for rep in range(AL_REPEATS):
    g0 = np.random.default_rng(20_000 + rep)
    seed_idx = np.concatenate([g0.choice(np.flatnonzero(yp_pool == c), START, replace=False) for c in IDS])
    for strategy in ("random", "uncertainty"):
        g = np.random.default_rng(30_000 + rep)
        lab = list(seed_idx)
        for rnd in range(ROUNDS + 1):
            m = LogisticRegression(C=1.0, max_iter=3000).fit(Zp[lab], yp_pool[lab])
            al_rows.append({"strategy": strategy, "rep": rep, "n_labels": len(lab),
                            "test_oa": (m.predict(Zte) == y[is_te]).mean(),
                            "classes_seen": len(np.unique(yp_pool[lab]))})
            if rnd == ROUNDS:
                break
            rest = np.setdiff1d(np.arange(len(pool)), lab)
            if strategy == "random":
                new = g.choice(rest, BATCH, replace=False)
            else:
                P = np.sort(m.predict_proba(Zp[rest]), 1)
                new = rest[np.argsort(P[:, -1] - P[:, -2])[:BATCH]]     # smallest margin first
            lab.extend(new.tolist())
al = pd.DataFrame(al_rows)
al_sum = al.groupby(["n_labels", "strategy"]).test_oa.agg(["mean", "std"]).unstack(1)
print(f"test OA by number of labels (mean and s.d. over {AL_REPEATS} runs; start {START} per class, "
      f"batches of {BATCH}):")
print(al_sum.iloc[::4].round(3).to_string())
al_mean = al.groupby(["strategy", "n_labels"]).test_oa.mean().unstack(0)
target = al_mean.loc[al_mean.index.max(), "random"]
n_unc = al_mean.index[al_mean.uncertainty >= target].min()
print(f"random sampling reaches {target:.3f} with {al_mean.index.max()} labels; uncertainty sampling reaches it with "
      f"{n_unc} labels")

# --- Step 7: area-adjusted accuracy and class areas (Olofsson et al. 2014) --------------------------
section("Step 7: stratified random sample of the map, area-adjusted accuracy and areas")
popn = (labels > 0) & (pred > 0)                  # Living England labels every pixel except fen and the edges
mapped = np.array([(pred[popn] == c).sum() for c in IDS], float)
W = mapped / mapped.sum()
N_TOTAL, N_MIN = 2000, 100
n_h = np.maximum(N_MIN, np.round(W * N_TOTAL)).astype(int)
g = np.random.default_rng(hc.SEED)
flat_pred = np.where(popn, pred, 0).ravel()
pts = []
for c, n in zip(IDS, n_h):
    pts.append(g.choice(np.flatnonzero(flat_pred == c), n, replace=False))
pts = np.concatenate(pts)
s_map = flat_pred[pts]
s_ref = labels.ravel()[pts]
s_te = in_test.ravel()[pts]
print(f"population {popn.sum():,} pixels ({popn.sum() * 0.01:,.0f} ha); stratified sample n = {len(pts):,} "
      f"(at least {N_MIN} per mapped class, the rest proportional to mapped area; seed {hc.SEED})")
print("reference = Living England 2022-23 at each sampled pixel, any reliability, no purity filter")



def olofsson(map_cls, ref_cls, mapped_px, px_ha=0.01):
    """Stratified estimators of Olofsson et al. (2014): strata = map classes, weights = mapped area shares.

    Returns overall accuracy and its standard error, and per class: user's and producer's accuracy, estimated
    area (ha) and their standard errors. Valid only if the sample is a stratified random sample of the map.
    """
    nij = pd.crosstab(pd.Categorical(map_cls, IDS), pd.Categorical(ref_cls, IDS), dropna=False).to_numpy().astype(float)
    ni = nij.sum(1)
    W = mapped_px / mapped_px.sum()
    p = W[:, None] * nij / ni[:, None]               # estimated area proportions (map i, reference j)
    pj = p.sum(0)
    ua = np.diag(nij) / ni
    se_ua = np.sqrt(ua * (1 - ua) / (ni - 1))
    se_oa = np.sqrt((W ** 2 * ua * (1 - ua) / (ni - 1)).sum())
    se_pj = np.sqrt((W[:, None] ** 2 * (nij / ni[:, None]) * (1 - nij / ni[:, None]) / (ni[:, None] - 1)).sum(0))
    pa = np.diag(p) / pj
    Nj = (mapped_px[:, None] * nij / ni[:, None]).sum(0)      # estimated reference pixels of each class
    se_pa = np.empty(len(IDS))
    for j in range(len(IDS)):
        others = sum(mapped_px[i] ** 2 * nij[i, j] / ni[i] * (1 - nij[i, j] / ni[i]) / (ni[i] - 1)
                     for i in range(len(IDS)) if i != j)
        se_pa[j] = np.sqrt((mapped_px[j] ** 2 * (1 - pa[j]) ** 2 * ua[j] * (1 - ua[j]) / (ni[j] - 1)
                            + pa[j] ** 2 * others) / Nj[j] ** 2)
    A = mapped_px.sum() * px_ha
    return {"oa": np.trace(p), "se_oa": se_oa, "ua": ua, "se_ua": se_ua, "pa": pa, "se_pa": se_pa,
            "area_ha": pj * A, "se_area_ha": se_pj * A}


est = olofsson(s_map, s_ref, mapped)
oa, se_oa = est["oa"], est["se_oa"]
census = np.array([(labels[popn] == c).sum() for c in IDS]) * 0.01
area = pd.DataFrame({
    "n_sample": n_h, "mapped_ha": mapped * 0.01, "adjusted_ha": est["area_ha"], "ci95_ha": 1.96 * est["se_area_ha"],
    "living_england_ha": census, "users_acc": est["ua"], "ua_ci95": 1.96 * est["se_ua"],
    "producers_acc": est["pa"], "pa_ci95": 1.96 * est["se_pa"]}, index=[NAMES[k] for k in KEYS])
area["le_inside_ci"] = (area.living_england_ha - area.adjusted_ha).abs() <= area.ci95_ha
print("areas (ha) and the Living England census of the same population:")
print(area[["n_sample", "mapped_ha", "adjusted_ha", "ci95_ha", "living_england_ha", "le_inside_ci"]].round(0)
      .astype({"mapped_ha": int, "adjusted_ha": int, "ci95_ha": int, "living_england_ha": int})
      .rename(columns={"living_england_ha": "le_ha", "le_inside_ci": "le_in_ci"}).to_string())
print("area-adjusted user's and producer's accuracy with 95% CI half-widths:")
print(area[["users_acc", "ua_ci95", "producers_acc", "pa_ci95"]].rename(columns={"producers_acc": "prod_acc"})
      .round(3).to_string())
print(f"area-adjusted overall accuracy {oa:.3f} +/- {1.96 * se_oa:.3f} (95% CI); "
      f"compare {res['oa']:.3f} on the balanced, pure test samples")
print(f"sample points in test blocks: {s_te.sum()} (agreement {(s_map == s_ref)[s_te].mean():.3f}); "
      f"in training blocks: {(~s_te).sum()} (agreement {(s_map == s_ref)[~s_te].mean():.3f})")
bog = area.loc[NAMES["bog"]]
print(f"blanket bog: mapped {bog.mapped_ha:,.0f} ha; adjusted {bog.adjusted_ha:,.0f} ha "
      f"(95% CI {bog.adjusted_ha - bog.ci95_ha:,.0f} to {bog.adjusted_ha + bog.ci95_ha:,.0f}); "
      f"Living England census {bog.living_england_ha:,.0f} ha")
print(f"Living England census inside the 95% CI for {area.le_inside_ci.sum()} of {len(area)} classes")

# What goes wrong if the balanced, pure test samples are treated as if they were a stratified map sample
bad = olofsson(yp, yt, mapped)
bog_i = KEYS.index("bog")
print(f"misuse check, the 2,456 pure test samples plugged into the same estimator: 'area-adjusted' OA "
      f"{bad['oa']:.3f}; blanket bog {bad['area_ha'][bog_i]:,.0f} +/- {1.96 * bad['se_area_ha'][bog_i]:,.0f} ha")

# map-sample calibration: the same probabilities on the stratified sample, weighted by inclusion probability
conf_map = pmax.ravel()[pts]
w_map = (mapped / n_h)[np.searchsorted(IDS, s_map)]    # each sampled pixel stands for this many map pixels
rel_map, ece_map = reliability(conf_map, s_map == s_ref, w=w_map)
print(f"expected calibration error on the stratified map sample (area-weighted): {ece_map:.3f} "
      f"(pure test samples: {ece_te:.3f})")
print(f"area-weighted mean confidence on the map {np.average(conf_map, weights=w_map):.3f} vs area-adjusted "
      f"overall accuracy {oa:.3f}")

# --- metrics.json for the series comparison ----------------------------------------------------------
runtime = time.perf_counter() - T_START
metrics = {
    "oa": round(float(res["oa"]), 4), "oa_ci95": [round(float(lo), 4), round(float(hi), 4)],
    "macro_f1": round(float(np.nanmean(res["f1"])), 4), "kappa": round(float(res["kappa"]), 4),
    "per_class": {k: {"pa": round(float(res["pa"][i]), 4), "ua": round(float(res["ua"][i]), 4),
                      "f1": round(float(res["f1"][i]), 4)} for i, k in enumerate(KEYS)},
    "n_test": int(len(yt)), "runtime_s": round(runtime, 1),
    "method": f"AlphaEarth Foundations 2023 embeddings (64-d) + logistic regression (C={BEST_C[e_col]:g})",
    "wall_to_wall_agreement": round(float(agree_all), 4),
    "wall_to_wall_agreement_test_blocks": round(float(agree_te), 4),
    "area_adjusted_oa": round(float(oa), 4), "area_adjusted_oa_ci95": round(float(1.96 * se_oa), 4),
    "bog_area_ha": {"mapped": round(float(bog.mapped_ha)), "adjusted": round(float(bog.adjusted_ha)),
                    "ci95": round(float(bog.ci95_ha)), "living_england": round(float(bog.living_england_ha))},
}
(OUT / "metrics.json").write_text(json.dumps(metrics, indent=2))

# --- Figures -----------------------------------------------------------------------------------------
ext = (hc.AOI_BOUNDS[0], hc.AOI_BOUNDS[2], hc.AOI_BOUNDS[1], hc.AOI_BOUNDS[3])
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")])


# Figure 1: Sentinel-2 true colour and the first three principal components of the embeddings
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)
gpix = np.random.default_rng(hc.SEED).choice(np.flatnonzero(valid.ravel()), 200_000, replace=False)
pca = PCA(3, random_state=hc.SEED).fit(aef.dequantise(raw.reshape(64, -1)[:, gpix].T))
pcs = np.zeros((rows_, cols_, 3), np.float32)
for r0 in range(0, rows_, 250):
    blk = raw[:, r0:r0 + 250, :]
    pcs[r0:r0 + blk.shape[1]] = pca.transform(aef.dequantise(blk.reshape(64, -1).T)).reshape(blk.shape[1], cols_, 3)
lo_p, hi_p = np.percentile(pcs.reshape(-1, 3), [2, 98], axis=0)
pcs = np.clip((pcs - lo_p) / (hi_p - lo_p), 0, 1)
print(f"\nPCA of the embeddings: first three components explain "
      f"{pca.explained_variance_ratio_.sum():.0%} of the variance ({', '.join(f'{v:.0%}' for v in pca.explained_variance_ratio_)})")
fig, axes = plt.subplots(1, 2, figsize=(13, 5.6), sharey=True)
axes[0].imshow(rgb, extent=ext)
axes[0].set_title("Sentinel-2 summer 2023 median (true colour)")
axes[1].imshow(pcs, extent=ext, interpolation="nearest")
axes[1].set_title(f"AlphaEarth 2023 embeddings: principal components 1-3 as RGB "
                  f"({pca.explained_variance_ratio_.sum():.0%} of variance)")
for ax in axes:
    km_axes(ax)
    place_labels(ax)
axes[1].set_ylabel("")
fig.savefig(OUT / "01-embeddings-pca-rgb.png")
plt.close(fig)

# Figure 2: few-shot learning curve
fig, ax = plt.subplots(figsize=(8.6, 5))
xs = SHOTS + [int(is_tr.sum() / 9)]
styles = {(e_col, "logistic"): (BLUE, "-", "o"), (e_col, "kNN"): (BLUE, "--", "s"),
          (s_col, "logistic"): (RED, "-", "o"), (s_col, "kNN"): (RED, "--", "s")}
for (fname, mname), (col, ls, mk) in styles.items():
    d = fs[(fs.features == fname) & (fs.classifier == mname)].groupby("per_class", sort=False).test_oa
    mu, lo10, hi90 = d.mean().to_numpy(), d.quantile(0.1).to_numpy(), d.quantile(0.9).to_numpy()
    lbl = f"{fname.split(' (')[0]} + {'logistic regression' if mname == 'logistic' else 'kNN (k = 5)'}"
    ax.plot(xs, mu, ls=ls, marker=mk, color=col, label=lbl, lw=1.6, ms=4)
    ax.fill_between(xs, lo10, hi90, color=col, alpha=0.12, lw=0)
ax.axhline(PART1_OA, color="#555555", lw=1, ls=":")
ax.text(5, PART1_OA + 0.008, "Part 1 spectral rules (0.600)", fontsize=8, color="#555555")
ax.set_xscale("log")
ax.set_xticks(xs, [str(v) for v in SHOTS] + [f"all\n(~{xs[-1]})"])
ax.minorticks_off()
ax.set_xlabel("Training labels per class (drawn from the training blocks)")
ax.set_ylabel("Overall accuracy on the 2,456 test samples")
ax.set_ylim(0.3, 0.95)
ax.grid(alpha=0.3)
ax.legend(loc="lower right", fontsize=8)
ax.set_title(f"Few-shot learning curve: mean of {REPEATS} random draws, shading = 10th to 90th percentile")
fig.savefig(OUT / "02-few-shot-learning-curve.png")
plt.close(fig)

# Figure 3: wall-to-wall 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(f"AlphaEarth embeddings + logistic regression (test OA {res['oa']:.1%})")
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 {agree_all:.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 / "03-alphaearth-map-vs-living-england.png")
plt.close(fig)

# Figure 4: confusion matrix
rown = cm / cm.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 cm[i, j]:
            ax.text(j, i, cm[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("AlphaEarth + logistic regression (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 / "04-confusion-matrix.png")
plt.close(fig)

# Figure 5: uncertainty map and reliability diagram
fig, axes = plt.subplots(1, 2, figsize=(13, 5.4), gridspec_kw={"width_ratios": [1.25, 1]})
im = axes[0].imshow(margin, cmap="Spectral", vmin=0, vmax=1, extent=ext, interpolation="nearest")
km_axes(axes[0])
place_labels(axes[0], colour="white")
axes[0].set_title("Uncertainty: margin between the two most likely classes (red = unsure)")
fig.colorbar(im, ax=axes[0], shrink=0.8, label="Probability margin (top class - runner-up)")
ax = axes[1]
ax.plot([0, 1], [0, 1], color="#888888", lw=1, ls=":", label="Perfect calibration")
ax.plot(rel_te.conf, rel_te.acc, "o-", color=BLUE, label=f"Pure test samples (ECE {ece_te:.3f})")
ax.plot(rel_map.conf, rel_map.acc, "s-", color=RED, label=f"Stratified map sample, area-weighted (ECE {ece_map:.3f})")
ax.set_xlabel("Predicted probability of the chosen class (binned)")
ax.set_ylabel("Share of samples that agree with Living England")
ax.set_xlim(0, 1.02)
ax.set_ylim(0, 1.02)
ax.set_aspect("equal")
ax.grid(alpha=0.3)
ax.legend(loc="upper left", fontsize=8)
ax.set_title("Reliability diagram")
fig.tight_layout()
fig.savefig(OUT / "05-uncertainty-reliability.png")
plt.close(fig)

# Figure 6: active learning curve
fig, ax = plt.subplots(figsize=(8.6, 4.8))
for strategy, col, lbl in (("random", RED, "Random sampling"), ("uncertainty", BLUE, "Uncertainty (margin) sampling")):
    d = al[al.strategy == strategy].groupby("n_labels").test_oa
    ax.plot(d.mean().index, d.mean(), "o-", color=col, ms=3.5, lw=1.6, label=lbl)
    ax.fill_between(d.mean().index, d.min(), d.max(), color=col, alpha=0.12, lw=0)
ax.axhline(lr.loc["all", e_col], color=BLUE, lw=1, ls=":")
ax.text(al.n_labels.max(), lr.loc["all", e_col] - 0.014, f"all {is_tr.sum():,} training labels ({lr.loc['all', e_col]:.3f})",
        ha="right", fontsize=8, color=BLUE)
ax.set_xlabel("Labelled training samples (start: 2 per class, then batches of 18 from the training-block pool)")
ax.set_ylabel("Overall accuracy on the test samples")
ax.grid(alpha=0.3)
ax.legend(loc="lower right", fontsize=8)
ax.set_title(f"Simulated active learning with AlphaEarth + logistic regression: mean of {AL_REPEATS} runs, "
             "shading = min to max", fontsize=9)
fig.savefig(OUT / "06-active-learning.png")
plt.close(fig)

# Figure 7: class areas with 95% confidence intervals
fig, ax = plt.subplots(figsize=(9.4, 5))
pos = np.arange(len(KEYS))
ax.bar(pos - 0.27, area.mapped_ha / 1000, 0.27, color=ORANGE, label="Pixel counting (mapped area)")
ax.bar(pos, area.adjusted_ha / 1000, 0.27, color=BLUE, yerr=area.ci95_ha / 1000, capsize=3,
       label="Area-adjusted estimate with 95% CI")
ax.bar(pos + 0.27, area.living_england_ha / 1000, 0.27, color="#999999", label="Living England census")
ax.set_xticks(pos, [NAMES[k].replace(" woodland", "\nwoodland").replace(" grassland", "\ngrassland")
                    .replace(" and bare", "\nand bare") for k in KEYS], fontsize=8)
ax.set_ylabel("Area (thousand hectares)")
ax.grid(axis="y", alpha=0.3)
ax.legend(fontsize=8)
ax.set_title(f"Habitat areas from a stratified random sample of {len(pts):,} map pixels (Olofsson et al. 2014)",
             fontsize=9)
fig.savefig(OUT / "07-area-estimates.png")
plt.close(fig)

print(f"\nsaved 7 figures, habitat_alphaearth_2023.tif, margin_alphaearth_2023.tif and metrics.json to {OUT.name}/")
print(f"total runtime {time.perf_counter() - T_START:.0f} s")