Source code: Habitat classification with LightGBM in Python: Sentinel-2, Sentinel-1 and terrain fusion

"""Habitat classification with LightGBM on multi-temporal Sentinel-2 and Sentinel-1 features (plus terrain).

Part 2 of the "AI habitat classification with Python" series (Dark Peak, Peak District, UK).
Builds a pixel feature stack from 2023 seasonal Sentinel-2 bands and indices (NDVI, NDWI, NDMI, red-edge),
seasonal Sentinel-1 RTC backscatter (VV, VH, VH/VV ratio) and Copernicus DEM elevation and slope, tunes a
LightGBM classifier with blocked (spatial-fold) cross-validation and class weights, runs an ablation
(S2 only, S2+S1, S2+S1+terrain), explains the model with TreeSHAP and grouped permutation importance, and
compares it with a random forest and the Part 1 rule-based baseline on the identical held-out test blocks.
Run with:  python habitat_lightgbm.py
"""
import itertools
import json
import platform
import subprocess
import sys
import time
import warnings
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import lightgbm as lgb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
from matplotlib.patches import Patch
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import KFold

T0 = time.time()
HERE = Path(__file__).resolve().parent
COMMON = HERE.parent / "ai-habitat-classification-common"
PART1 = HERE.parent / "habitat-mapping-sentinel-2-python"
sys.path.insert(0, str(COMMON))
import habitat_common as hc  # noqa: E402

OUT, DATA = HERE / "output", HERE / "data"
OUT.mkdir(exist_ok=True)
DATA.mkdir(exist_ok=True)
warnings.filterwarnings("ignore", category=RuntimeWarning)          # NaN where a season has no clear view
warnings.filterwarnings("ignore", message="X does not have valid feature names")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 170)
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}
SHORT = {k: k.split("_")[0][:5] for k in KEYS}


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


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

# --- Step 1: samples and the shared split ------------------------------------------------------------
section("Step 1: samples and the shared spatial split")
samples = hc.load_samples()
labels = hc.label_raster()
print(f"{len(samples):,} samples: {(samples.split == 'train').sum():,} train in "
      f"{samples[samples.split == 'train'].block_id.nunique()} blocks (5 folds), "
      f"{(samples.split == 'test').sum():,} test in {samples[samples.split == 'test'].block_id.nunique()} blocks")

# --- Step 2: the feature stack -----------------------------------------------------------------------
section("Step 2: Sentinel-2, Sentinel-1 and terrain features")


def nd(a, b):
    return (a - b) / (a + b)


S2 = hc.s2_seasonal_composites()                 # median reflectance per season (cached GeoTIFFs)
S1 = {s: hc.s1_composite(s, m) for s, m in hc.SEASONS.items()}   # median gamma0 per season


def copernicus_dem():
    """Copernicus DEM GLO-30 on the shared 10 m grid (bilinear), cached as data/cop_dem_10m.tif."""
    f = DATA / "cop_dem_10m.tif"
    if not f.exists():
        import odc.stac
        items = hc.stac_client().search(collections=["cop-dem-glo-30"], bbox=hc.aoi_bbox_wgs84()).item_collection()
        dem = odc.stac.load(items, bands=["data"], geobox=hc.geobox(), resampling="bilinear",
                            dtype="float32", nodata=np.nan)["data"].isel(time=0).values
        hc.write_geotiff(f, dem, nodata=np.nan, dtype="float32")
        print(f"Copernicus DEM: {len(items)} tiles reprojected to the 10 m grid")
    with rasterio.open(f) as src:
        return src.read(1)


FEAT, GROUP = {}, {}


def add(name, arr, group):
    a = np.asarray(arr, dtype="float32")
    a[~np.isfinite(a)] = np.nan                  # e.g. log10(0) or 0/0 at no-data edges
    FEAT[name], GROUP[name] = a, group


BAND_GROUP = {"B02": "S2 visible", "B03": "S2 visible", "B04": "S2 visible", "B05": "S2 red-edge bands",
              "B06": "S2 red-edge bands", "B07": "S2 red-edge bands", "B8A": "S2 red-edge bands", "B08": "S2 NIR",
              "B11": "S2 SWIR", "B12": "S2 SWIR"}
for s, ds in S2.items():
    b = {k: ds[k].values for k in hc.S2_BANDS}
    for k in hc.S2_BANDS:
        add(f"{k}_{s}", b[k], BAND_GROUP[k])
    add(f"ndvi_{s}", nd(b["B08"], b["B04"]), "NDVI")              # greenness
    add(f"ndwi_{s}", nd(b["B03"], b["B08"]), "NDWI")              # open water (McFeeters)
    add(f"ndmi_{s}", nd(b["B08"], b["B11"]), "NDMI")              # leaf and soil moisture
    add(f"ndre_{s}", nd(b["B8A"], b["B05"]), "Red-edge indices")  # normalised difference red-edge
    add(f"cire_{s}", b["B07"] / b["B05"] - 1, "Red-edge indices")  # red-edge chlorophyll index
ndvi = np.stack([FEAT[f"ndvi_{s}"] for s in hc.SEASONS])
add("ndvi_range", np.nanmax(ndvi, 0) - np.nanmin(ndvi, 0), "NDVI")
add("dndvi_spring_summer", FEAT["ndvi_summer"] - FEAT["ndvi_spring"], "NDVI")
S2_COLS = list(FEAT)

for s, ds in S1.items():
    add(f"vv_{s}", ds["vv_db"].values, "S1 VV")
    add(f"vh_{s}", ds["vh_db"].values, "S1 VH")
    add(f"ratio_{s}", ds["ratio_db"].values, "S1 VH/VV ratio")
vh = np.stack([FEAT[f"vh_{s}"] for s in hc.SEASONS])
add("vh_range", np.nanmax(vh, 0) - np.nanmin(vh, 0), "S1 VH")
S1_COLS = [c for c in FEAT if c not in S2_COLS]

dem = copernicus_dem()
gy, gx = np.gradient(dem, hc.RES)
add("elevation", dem, "Elevation")
add("slope", np.degrees(np.arctan(np.hypot(gx, gy))), "Slope")
TERRAIN_COLS = ["elevation", "slope"]

FEATURE_SETS = {"S2": S2_COLS, "S2+S1": S2_COLS + S1_COLS, "S2+S1+terrain": S2_COLS + S1_COLS + TERRAIN_COLS}
print(f"{len(S2_COLS)} Sentinel-2 features (4 seasons x 10 bands + 5 indices, plus 2 NDVI change features)")
print(f"{len(S1_COLS)} Sentinel-1 features (4 seasons x VV, VH, VH/VV in dB, plus the VH seasonal range)")
print(f"{len(TERRAIN_COLS)} terrain features; elevation {np.nanmin(dem):.0f} to {np.nanmax(dem):.0f} m, "
      f"slope median {np.nanmedian(FEAT['slope']):.1f} deg")

X_all = pd.DataFrame({k: hc.sample_values(v, samples) for k, v in FEAT.items()})
y_all = samples["class_id"].to_numpy() - 1      # LightGBM wants classes 0..8
is_train = (samples.split == "train").to_numpy()
Xtr, Xte, ytr, yte = X_all[is_train], X_all[~is_train], y_all[is_train], y_all[~is_train]
fold = samples.loc[is_train, "fold"].to_numpy()
test_blocks = samples.loc[~is_train, "block_id"].to_numpy()
print(f"sample feature table: {X_all.shape[0]:,} x {X_all.shape[1]}; share of missing values "
      f"{X_all.isna().to_numpy().mean():.2%} (mostly winter Sentinel-2: LightGBM handles NaN natively)")
show = ["ndre_summer", "ndmi_spring", "vh_winter", "ratio_summer", "vv_spring", "elevation", "slope"]
med = Xtr.groupby(samples.loc[is_train, "class_key"].to_numpy())[show].median().reindex(KEYS).astype(float).round(2).T
med.columns = [SHORT[k] for k in KEYS]
print("median of selected features by class (training samples only):")
print(med.to_string())

# --- Step 3: class weights and spatial-fold tuning ----------------------------------------------------
section("Step 3: class weights and LightGBM tuning with blocked 5-fold cross-validation")
counts = np.bincount(ytr, minlength=len(KEYS))
cw = len(ytr) / (len(KEYS) * counts)            # "balanced" weights: n / (classes x class count)
w_tr = cw[ytr]
print("class weights: " + ", ".join(f"{SHORT[k]} {w:.2f}" for k, w in zip(KEYS, cw)))
FOLDS = [(np.flatnonzero(fold != k), np.flatnonzero(fold == k)) for k in range(5)]
BASE = {"objective": "multiclass", "num_class": len(KEYS), "learning_rate": 0.05, "bagging_fraction": 0.8,
        "bagging_freq": 1, "lambda_l2": 1.0, "metric": "multi_logloss", "verbose": -1, "seed": hc.SEED,
        "deterministic": True, "force_col_wise": True, "num_threads": 8}
GRID = {"num_leaves": [15, 63], "min_data_in_leaf": [10, 40], "feature_fraction": [0.3, 0.7]}


def tune(cols):
    """Grid search scored by blocked-CV log loss; early stopping picks the number of trees."""
    rows = []
    for vals in itertools.product(*GRID.values()):
        p = {**BASE, **dict(zip(GRID, vals))}
        ds = lgb.Dataset(Xtr[cols], ytr, weight=w_tr, free_raw_data=False)
        cv = lgb.cv(p, ds, num_boost_round=3000, folds=FOLDS,
                    callbacks=[lgb.early_stopping(100, verbose=False)])
        loss = cv["valid multi_logloss-mean"]
        rows.append({**dict(zip(GRID, vals)), "cv_logloss": loss[-1], "n_trees": len(loss)})
    return pd.DataFrame(rows).sort_values("cv_logloss").reset_index(drop=True)


def fit(cols, params, n_trees, X=None, y=None, w=None):
    X, y, w = (Xtr, ytr, w_tr) if X is None else (X, y, w)
    return lgb.train({**BASE, **params}, lgb.Dataset(X[cols], y, weight=w), num_boost_round=n_trees)


def oof_accuracy(cols, params, n_trees, splits):
    pred = np.empty(len(ytr), int)
    for tr, va in splits:
        m = fit(cols, params, n_trees, Xtr.iloc[tr], ytr[tr], w_tr[tr])
        pred[va] = m.predict(Xtr.iloc[va][cols]).argmax(1)
    return (pred == ytr).mean(), np.array([(pred[va] == ytr[va]).mean() for _, va in splits])


best = {}
for name, cols in FEATURE_SETS.items():
    t = time.time()
    res = tune(cols)
    top = res.iloc[0]
    params = {k: (int(top[k]) if k != "feature_fraction" else float(top[k])) for k in GRID}
    oa_cv, per_fold = oof_accuracy(cols, params, int(top.n_trees), FOLDS)
    best[name] = {"cols": cols, "params": params, "n_trees": int(top.n_trees), "cv_logloss": top.cv_logloss,
                  "cv_oa": oa_cv, "cv_oa_sd": per_fold.std()}
    print(f"\n{name} ({len(cols)} features): {len(res)} settings tuned in {time.time() - t:.0f} s")
    if name == "S2+S1+terrain":
        print(res.round(4).to_string(index=False))
    print(f"best: {params}, {int(top.n_trees)} trees; blocked-CV log loss {top.cv_logloss:.4f}, "
          f"OA {oa_cv:.3f} +/- {per_fold.std():.3f}")

FULL = "S2+S1+terrain"
b = best[FULL]
oa_rand, _ = oof_accuracy(b["cols"], b["params"], b["n_trees"],
                          list(KFold(5, shuffle=True, random_state=hc.SEED).split(Xtr)))
print(f"\nsame model, random 5-fold CV of the training samples: OA {oa_rand:.3f} "
      f"(blocked 5-fold: {b['cv_oa']:.3f})")

# --- Step 4: accuracy on the held-out test blocks ----------------------------------------------------
section("Step 4: test-block accuracy, ablation and comparison")


def summarise(y_true, y_pred):
    k = len(KEYS)                                # unclassified predictions (-1) go to an extra column k
    yp = np.where(np.asarray(y_pred) < 0, k, y_pred)
    full = np.bincount(np.asarray(y_true) * (k + 1) + yp, minlength=k * (k + 1)).reshape(k, k + 1)
    m, full = full[:, :k].astype(float), full.astype(float)
    n = full.sum()
    oa = np.trace(m) / n
    pe = (m.sum(0) * full.sum(1)).sum() / n ** 2
    pa = np.diag(m) / full.sum(1)
    ua = np.diag(m) / np.where(m.sum(0) == 0, np.nan, m.sum(0))
    f1 = np.nan_to_num(2 * pa * ua / (pa + ua))
    return {"cm": m, "oa": oa, "kappa": (oa - pe) / (1 - pe), "pa": pa, "ua": ua, "f1": f1, "macro_f1": f1.mean()}


rng = np.random.default_rng(hc.SEED)
UB = np.unique(test_blocks)
BOOT = [np.concatenate([np.flatnonzero(test_blocks == bb) for bb in rng.choice(UB, UB.size, replace=True)])
        for _ in range(2000)]                    # the same 2,000 block resamples for every model (paired)


def block_ci(correct):
    return np.percentile([correct[i].mean() for i in BOOT], [2.5, 97.5])


def paired_diff(correct_a, correct_b):
    d = np.array([correct_a[i].mean() - correct_b[i].mean() for i in BOOT])
    return correct_a.mean() - correct_b.mean(), np.percentile(d, [2.5, 97.5]), (d <= 0).mean()


models, preds, times = {}, {}, {}
for name, bb in best.items():
    t = time.time()
    models[name] = fit(bb["cols"], bb["params"], bb["n_trees"])
    times[name] = time.time() - t
    preds[name] = models[name].predict(Xte[bb["cols"]]).argmax(1)

t = time.time()
rf = RandomForestClassifier(n_estimators=500, class_weight="balanced", n_jobs=8, random_state=hc.SEED)
rf.fit(Xtr[b["cols"]], ytr)                      # scikit-learn >= 1.4 forests accept NaN too
times["Random forest"] = time.time() - t
preds["Random forest"] = rf.predict(Xte[b["cols"]])
print("fit times: " + ", ".join(f"{k} {v:.1f} s" for k, v in times.items()))

# Part 1's rule-based map, sampled at the same test pixels (runs Part 1 once if its map is missing)
p1_tif = PART1 / "output" / "habitat_rules_2023.tif"
if not p1_tif.exists():
    print("running Part 1 (habitat_rules.py) once to create its map ...", flush=True)
    subprocess.run([sys.executable, "habitat_rules.py"], cwd=PART1, check=True, capture_output=True)
with rasterio.open(p1_tif) as src:
    p1_map = src.read(1)
preds["Part 1 rules"] = hc.sample_values(p1_map, samples[~is_train]).astype(int) - 1
print(f"Part 1 map sampled at the test pixels; unclassified there: {(preds['Part 1 rules'] < 0).sum()} (scored as errors)")

ORDER = ["Part 1 rules", "Random forest", "S2", "S2+S1", "S2+S1+terrain"]
LABEL = {"Part 1 rules": "Part 1: spectral rules", "Random forest": "Random forest (S2+S1+terrain)",
         "S2": "LightGBM: S2", "S2+S1": "LightGBM: S2+S1", "S2+S1+terrain": "LightGBM: S2+S1+terrain"}
R = {k: summarise(yte, preds[k]) for k in ORDER}
CORRECT = {k: preds[k] == yte for k in ORDER}
I = {k: KEYS.index(k) for k in KEYS}
rows = []
for k in ORDER:
    r, m = R[k], R[k]["cm"]
    lo, hi = block_ci(CORRECT[k])
    rows.append({"model": LABEL[k], "n_feat": len(best[k]["cols"]) if k in best else
                 (len(b["cols"]) if k == "Random forest" else 8), "oa": r["oa"], "ci_lo": lo, "ci_hi": hi,
                 "macro_f1": r["macro_f1"], "kappa": r["kappa"],
                 "f1_bog": r["f1"][I["bog"]], "f1_heath": r["f1"][I["heath"]],
                 "f1_semin": r["f1"][I["seminat_grass"]], "f1_impro": r["f1"][I["improved_grass"]],
                 "bog<>heath": int(m[I["bog"], I["heath"]] + m[I["heath"], I["bog"]]),
                 "semin<>impro": int(m[I["seminat_grass"], I["improved_grass"]]
                                     + m[I["improved_grass"], I["seminat_grass"]])})
table = pd.DataFrame(rows).set_index("model")
print(f"test blocks: n = {len(yte):,} samples in {UB.size} blocks; 95% CI from 2,000 block-bootstrap resamples")
print(table.round(3).to_string())

print("\npaired block-bootstrap differences in overall accuracy (95% CI; share of resamples <= 0):")
for a, c in [("S2+S1", "S2"), ("S2+S1+terrain", "S2+S1"), ("S2+S1+terrain", "S2"),
             ("S2+S1+terrain", "Random forest"), ("S2", "Part 1 rules"), ("S2+S1+terrain", "Part 1 rules")]:
    d, (lo, hi), p = paired_diff(CORRECT[a], CORRECT[c])
    print(f"  {LABEL[a]:<30} - {LABEL[c]:<30} {d:+.3f} ({lo:+.3f} to {hi:+.3f}; {p:.3f})")

R_full = R[FULL]
cm_df = pd.DataFrame(R_full["cm"].astype(int), index=[f"ref {k}" for k in KEYS], columns=[SHORT[k] for k in KEYS])
print(f"\nconfusion matrix, LightGBM {FULL} (rows = Living England, columns = model):")
print(cm_df.to_string())
per_class = pd.DataFrame({"n_test": R_full["cm"].sum(1).astype(int), "producers_acc": R_full["pa"],
                          "users_acc": R_full["ua"], "f1": R_full["f1"],
                          "f1_part1": R["Part 1 rules"]["f1"]}, index=[NAMES[k] for k in KEYS])
print(per_class.round(3).to_string())
lo, hi = block_ci(CORRECT[FULL])
print(f"overall accuracy {R_full['oa']:.3f} (block-bootstrap 95% CI {lo:.3f} to {hi:.3f}); "
      f"macro F1 {R_full['macro_f1']:.3f}; kappa {R_full['kappa']:.3f}")

# --- Step 5: what the model uses: TreeSHAP and grouped permutation importance -------------------------
section("Step 5: feature importance (gain, TreeSHAP, grouped permutation)")
model, cols = models[FULL], b["cols"]
gain = pd.Series(model.feature_importance("gain"), index=cols)
gain_grp = (gain.groupby(pd.Series(GROUP)[cols]).sum() / gain.sum()).sort_values(ascending=False)
contrib = model.predict(Xte[cols], pred_contrib=True).reshape(len(Xte), len(KEYS), len(cols) + 1)[:, :, :-1]
shap_abs = pd.DataFrame(np.abs(contrib).mean(0).T, index=cols, columns=KEYS)     # mean |SHAP| per class
shap_tot = shap_abs.sum(axis=1).sort_values(ascending=False)
shap_grp = shap_abs.sum(axis=1).groupby(pd.Series(GROUP)[cols]).sum()
shap_grp = (shap_grp / shap_grp.sum()).sort_values(ascending=False)
print("share of total gain and of total mean |SHAP| by feature group:")
print(pd.DataFrame({"gain": gain_grp, "shap": shap_grp}).sort_values("shap", ascending=False).round(3).to_string())
print("top 12 features by mean |SHAP| (summed over classes):")
print(shap_tot.head(12).round(3).to_string())
bh = np.isin(yte, [I["bog"], I["heath"]])
bog_vs_heath = pd.Series(np.abs(contrib[bh, I["bog"], :] - contrib[bh, I["heath"], :]).mean(0), index=cols)
print("top 8 features separating bog from heath (mean |SHAP(bog) - SHAP(heath)| on bog and heath test samples):")
print(bog_vs_heath.sort_values(ascending=False).head(8).round(3).to_string())

N_REP = 10
groups = list(dict.fromkeys(GROUP[c] for c in cols))
base_f1 = R_full["macro_f1"]
perm = {}
prng = np.random.default_rng(hc.SEED)
for g in groups:
    gcols = [c for c in cols if GROUP[c] == g]
    drops = []
    for _ in range(N_REP):
        Xp = Xte[cols].copy()
        Xp[gcols] = Xp[gcols].to_numpy()[prng.permutation(len(Xp))]   # shuffle the group's rows together
        drops.append(base_f1 - summarise(yte, model.predict(Xp).argmax(1))["macro_f1"])
    perm[g] = (np.mean(drops), np.std(drops), len(gcols))
perm = pd.DataFrame(perm, index=["drop_macro_f1", "sd", "n_features"]).T.sort_values("drop_macro_f1", ascending=False)
perm["n_features"] = perm["n_features"].astype(int)
print(f"grouped permutation importance on the test blocks (drop in macro F1, {N_REP} repeats):")
print(perm.round(3).to_string())

# --- Step 6: classify the whole study area -----------------------------------------------------------
section("Step 6: wall-to-wall habitat map")
t = time.time()
rows_, cols_ = hc.shape()
pred_map = np.zeros(hc.shape(), "uint8")
conf_map = np.full(hc.shape(), np.nan, "float32")
STEP = 250
for r0 in range(0, rows_, STEP):
    Xc = np.stack([FEAT[c][r0:r0 + STEP].ravel() for c in cols], axis=1)
    p = model.predict(Xc)
    pred_map[r0:r0 + STEP] = (p.argmax(1) + 1).reshape(-1, cols_)
    conf_map[r0:r0 + STEP] = p.max(1).reshape(-1, cols_)
nodata = np.isnan(FEAT["ndvi_summer"]) & np.isnan(FEAT["ndvi_spring"])
pred_map[nodata], conf_map[nodata] = 0, np.nan
print(f"predicted {rows_ * cols_:,} pixels in {time.time() - t:.0f} s; unclassified (no clear spring or summer "
      f"view): {nodata.mean():.2%}")
hc.write_geotiff(OUT / "habitat_lightgbm_2023.tif", pred_map, nodata=0)
hc.write_geotiff(OUT / "habitat_lightgbm_confidence_2023.tif", conf_map, nodata=np.nan, dtype="float32")
area = pd.DataFrame({"lightgbm_km2": [(pred_map == KEY_ID[k]).sum() / 1e4 for k in KEYS],
                     "part1_km2": [(p1_map == KEY_ID[k]).sum() / 1e4 for k in KEYS],
                     "living_england_km2": [(labels == KEY_ID[k]).sum() / 1e4 for k in KEYS]},
                    index=[NAMES[k] for k in KEYS])
print(area.round(1).to_string())
ys, xs = np.mgrid[0:rows_, 0:cols_]
in_test = np.isin(hc.block_id(hc.AOI_BOUNDS[0] + (xs + 0.5) * hc.RES, hc.AOI_BOUNDS[3] - (ys + 0.5) * hc.RES), UB)
for nm, mp in (("LightGBM", pred_map), ("Part 1 rules", p1_map)):
    v = (labels > 0) & (mp > 0)
    print(f"wall-to-wall agreement with Living England, {nm}: test blocks {(labels == mp)[v & in_test].mean():.3f}, "
          f"whole area {(labels == mp)[v].mean():.3f}")
v = (labels > 0) & (pred_map > 0)
agree = labels == pred_map
for lo_, hi_ in ((0, 0.5), (0.5, 0.8), (0.8, 1.01)):
    sel = v & (conf_map >= lo_) & (conf_map < hi_)
    print(f"  pixels with top-class probability {lo_:.1f}-{min(hi_, 1):.1f}: {sel.sum() / v.sum():5.1%} of the map, "
          f"agreement {agree[sel].mean():.3f}")

# --- 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]


def km_axes(ax, ylabel=True):
    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, BNG)")
    if ylabel:
        ax.set_ylabel("Northing (km)")


# Figure 1: Part 1 rules, LightGBM and Living England
fig, axes = plt.subplots(1, 3, figsize=(15, 5.0), sharey=True)
for ax, mp, title in zip(axes, (p1_map, pred_map, labels),
                         (f"Part 1: spectral rules (test OA {R['Part 1 rules']['oa']:.0%})",
                          f"Part 2: LightGBM, S2+S1+terrain (test OA {R_full['oa']:.0%})",
                          "Living England 2022-23 (reference)")):
    ax.imshow(np.ma.masked_equal(mp, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
    ax.set_title(title, fontsize=9.5)
    km_axes(ax, ylabel=ax is axes[0])
fig.legend(handles=legend, loc="lower center", ncol=9, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.05))
fig.tight_layout()
fig.savefig(OUT / "01-lightgbm-habitat-map.png")
plt.close(fig)

# Figure 2: confusion matrix
m = R_full["cm"].astype(int)
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, R_full["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, R_full["pa"])], fontsize=7.5)
ax.set_xlabel("LightGBM, S2+S1+terrain (UA = user's accuracy)")
ax.set_ylabel("Living England reference (PA = producer's accuracy)")
ax.set_title(f"Test blocks: overall accuracy {R_full['oa']:.1%} (block-bootstrap 95% CI {lo:.1%} to {hi:.1%}), "
             f"kappa {R_full['kappa']:.2f}", fontsize=9)
fig.savefig(OUT / "02-confusion-matrix.png")
plt.close(fig)

# Figure 3: ablation, per-class F1 for every model
MC = {"Part 1 rules": "#bababa", "Random forest": "#878787", "S2": "#abd9e9", "S2+S1": "#2c7bb6",
      "S2+S1+terrain": "#d7191c"}
fig, axes = plt.subplots(1, 2, figsize=(13, 4.6), gridspec_kw={"width_ratios": [4.2, 1]})
x = np.arange(len(KEYS))
wd = 0.16
for i, k in enumerate(ORDER):
    axes[0].bar(x + (i - 2) * wd, R[k]["f1"], wd, color=MC[k], label=LABEL[k])
axes[0].set_xticks(x, [NAMES[k].replace(" woodland", "").replace(" grassland", "\ngrassland") for k in KEYS],
                   fontsize=8)
axes[0].set_ylabel("F1 score on the test blocks")
axes[0].set_ylim(0, 1.05)
axes[0].legend(fontsize=7.5, ncol=5, loc="upper center", bbox_to_anchor=(0.5, 1.13), frameon=False)
for i, k in enumerate(ORDER):
    lo_, hi_ = block_ci(CORRECT[k])
    axes[1].errorbar(i, R[k]["oa"], yerr=[[R[k]["oa"] - lo_], [hi_ - R[k]["oa"]]], fmt="o", color=MC[k], ms=7,
                     capsize=3, mec="#333333", mew=0.6)
    axes[1].text(i + 0.12, R[k]["oa"], f"{R[k]['oa']:.1%}", fontsize=7.5, va="center")
axes[1].set_xticks(range(len(ORDER)), ["P1", "RF", "S2", "+S1", "+ter."], fontsize=8)
axes[1].set_xlim(-0.5, len(ORDER) + 0.1)
axes[1].set_ylabel("Overall accuracy (block-bootstrap 95% CI)")
fig.tight_layout()
fig.savefig(OUT / "03-ablation-f1-by-class.png")
plt.close(fig)

# Figure 4: TreeSHAP, top 20 features, stacked by class
top = shap_tot.head(20).index[::-1]
fig, ax = plt.subplots(figsize=(8.6, 6.4))
left = np.zeros(len(top))
for k in KEYS:
    ax.barh(range(len(top)), shap_abs.loc[top, k], left=left, color=COLOURS[k], label=NAMES[k], ec="white", lw=0.3)
    left += shap_abs.loc[top, k].to_numpy()
ax.set_yticks(range(len(top)), [f"{c}  ({GROUP[c]})" for c in top], fontsize=7.5)
ax.set_xlabel("Mean |SHAP value| on the test samples (log-odds), stacked by class")
ax.legend(fontsize=7.5, loc="lower right", framealpha=0.95)
ax.set_title("TreeSHAP: the 20 most influential features of the S2+S1+terrain model", fontsize=9.5)
fig.savefig(OUT / "04-shap-feature-importance.png")
plt.close(fig)

# Figure 5: grouped permutation importance and share of gain
GC = {"Elevation": "#8c510a", "Slope": "#bf812d", "S1 VV": "#2c7bb6", "S1 VH": "#4575b4", "S1 VH/VV ratio": "#74add1"}
fig, axes = plt.subplots(1, 2, figsize=(12, 4.4), sharey=True)
order_g = perm.index[::-1]
colours_g = [GC.get(g, "#66bd63") for g in order_g]
axes[0].barh(range(len(order_g)), perm.loc[order_g, "drop_macro_f1"], xerr=perm.loc[order_g, "sd"],
             color=colours_g, capsize=2)
axes[0].set_yticks(range(len(order_g)), [f"{g} ({perm.loc[g, 'n_features']})" for g in order_g], fontsize=8)
axes[0].set_xlabel(f"Drop in test macro F1 when the group is shuffled (mean and s.d. of {N_REP})")
axes[0].axvline(0, color="#333333", lw=0.6)
axes[1].barh(range(len(order_g)), gain_grp.reindex(order_g), color=colours_g)
axes[1].set_xlabel("Share of total split gain")
fig.legend(handles=[Patch(color="#66bd63", label="Sentinel-2"), Patch(color="#4575b4", label="Sentinel-1"),
                    Patch(color="#8c510a", label="Terrain")], loc="lower center", ncol=3, fontsize=8, frameon=False,
           bbox_to_anchor=(0.5, -0.06))
fig.tight_layout()
fig.savefig(OUT / "05-permutation-importance.png")
plt.close(fig)

# Figure 6: radar and terrain signatures of the confusable pairs
pairs = ["bog", "heath", "seminat_grass", "improved_grass"]
feats = ["vh_winter", "ratio_summer", "vv_spring", "ndmi_spring", "elevation", "slope"]
units = {"vh_winter": "dB", "ratio_summer": "dB", "vv_spring": "dB", "ndmi_spring": "", "elevation": "m",
         "slope": "degrees"}
fig, axes = plt.subplots(2, 3, figsize=(12, 6.2))
trk = samples.loc[is_train, "class_key"].to_numpy()
for ax, f in zip(axes.flat, feats):
    data = [Xtr.loc[trk == k, f].dropna() for k in pairs]
    bp = ax.boxplot(data, showfliers=False, patch_artist=True, widths=0.6, medianprops={"color": "black"})
    for patch, k in zip(bp["boxes"], pairs):
        patch.set_facecolor(COLOURS[k])
        patch.set_edgecolor("#333333")
    ax.set_xticks(range(1, 5), ["Blanket\nbog", "Heath", "Semi-natural\ngrass", "Improved\ngrass"], fontsize=7.5)
    ax.set_title(f"{f} ({GROUP[f]})", fontsize=9)
    ax.set_ylabel(units[f])
    ax.axvline(2.5, color="#999999", lw=0.6, ls=":")
fig.suptitle("Training-sample distributions for the two hardest pairs (box = interquartile range, whiskers = 1.5 IQR)",
             fontsize=9.5)
fig.tight_layout()
fig.savefig(OUT / "06-radar-terrain-signatures.png")
plt.close(fig)

# Figure 7: model confidence
fig, axes = plt.subplots(1, 2, figsize=(12.6, 5.0), gridspec_kw={"width_ratios": [1.15, 1]})
im = axes[0].imshow(conf_map, cmap="viridis", vmin=0.2, vmax=1, extent=ext, interpolation="nearest")
fig.colorbar(im, ax=axes[0], shrink=0.85, label="Top-class probability")
axes[0].set_title("LightGBM confidence: probability of the predicted class", fontsize=9.5)
km_axes(axes[0])
bins = np.linspace(0.1, 1, 19)
idx = np.digitize(conf_map[v], bins) - 1
acc = pd.Series(agree[v]).groupby(idx).mean()
share = pd.Series(idx).value_counts(normalize=True).sort_index()
centres = (bins[:-1] + bins[1:]) / 2
ok = acc.index[(acc.index >= 0) & (acc.index < len(centres))]
ok = ok[share.reindex(ok).to_numpy() >= 0.001]  # skip near-empty bins
axes[1].bar(centres[ok], share.reindex(ok), width=0.045, color="#abd9e9", label="Share of map pixels")
ax2 = axes[1].twinx()
ax2.plot(centres[ok], acc.reindex(ok), "o-", color="#d7191c", ms=4, label="Agreement with Living England")
ax2.plot([0, 1], [0, 1], ls="--", color="#999999", lw=0.8, label="Agreement = probability")
ax2.set_ylim(0, 1)
ax2.set_ylabel("Pixel agreement with Living England", color="#d7191c")
axes[1].set_xlabel("Top-class probability")
axes[1].set_ylabel("Share of map pixels")
axes[1].set_title("Every labelled pixel in the study area", fontsize=9.5)
h1, l1 = axes[1].get_legend_handles_labels()
h2, l2 = ax2.get_legend_handles_labels()
axes[1].legend(h1 + h2, l1 + l2, fontsize=7.5, loc="upper left")
fig.tight_layout()
fig.savefig(OUT / "07-confidence-map.png")
plt.close(fig)

# --- metrics.json for the series comparison ---------------------------------------------------------
runtime = time.time() - T0
lo, hi = block_ci(CORRECT[FULL])
metrics = {"oa": round(R_full["oa"], 4), "oa_ci95": [round(lo, 4), round(hi, 4)],
           "macro_f1": round(R_full["macro_f1"], 4), "kappa": round(R_full["kappa"], 4),
           "per_class": {k: {"pa": round(R_full["pa"][i], 4), "ua": round(R_full["ua"][i], 4),
                             "f1": round(R_full["f1"][i], 4)} for i, k in enumerate(KEYS)},
           "n_test": int(len(yte)), "runtime_s": round(runtime, 1),
           "method": f"LightGBM {lgb.__version__} on {len(cols)} pixel features (seasonal Sentinel-2 bands and "
                     f"indices, seasonal Sentinel-1 VV/VH/ratio, Copernicus DEM elevation and slope), tuned with "
                     f"blocked 5-fold CV, balanced class weights"}
(OUT / "metrics.json").write_text(json.dumps(metrics, indent=2))
print(f"\nsaved 7 figures, habitat_lightgbm_2023.tif, the confidence GeoTIFF and metrics.json to {OUT.name}/")
print(f"total runtime {runtime / 60:.1f} min (features, tuning three feature sets, random forest, importance, "
      f"map and figures)")