"""Object-based image analysis (OBIA) for habitat classification with Sentinel-2 and Python.

Part 3 of the "AI habitat classification with Python" series (Dark Peak, Peak District, UK).
Segments 2023 seasonal Sentinel-2 composites into image objects with scikit-image (SLIC and Felzenszwalb),
tunes the segmentation against Living England boundaries and by blocked cross-validation, computes per-object
spectral statistics, percentiles, shape metrics and GLCM texture with NumPy label aggregation, classifies the
objects with LightGBM, and compares the result with a pixel-wise LightGBM on the same seasonal features:
accuracy on the shared test samples, salt-and-pepper noise and patch counts. Exports a GeoPackage of habitat
polygons.
Run with:  python obia_habitat.py
"""
import json
import platform
import sys
import time
import warnings
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import geopandas as gpd
import lightgbm as lgb
import matplotlib.patheffects as pe
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import skimage
from matplotlib.patches import Patch
from rasterio.features import shapes
from scipy import ndimage
from shapely.geometry import shape as as_shape
from skimage.segmentation import felzenszwalb, find_boundaries, slic

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

T0 = time.perf_counter()
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}
ID_KEY = {c["id"]: c["key"] for c in hc.CLASSES}
ROWS, COLS = hc.shape()
PX_HA = hc.RES ** 2 / 1e4                        # one 10 m pixel = 0.01 ha


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


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


print("=== versions ===")
print(f"python {platform.python_version()} | numpy {np.__version__} | scikit-image {skimage.__version__} | "
      f"lightgbm {lgb.__version__} | geopandas {gpd.__version__} | rasterio {version('rasterio')}")

# --- Step 1: shared data -----------------------------------------------------------------------------
section("Step 1: labels, samples and seasonal composites (shared with Part 1)")
labels = hc.label_raster()                       # Living England 2022-23 on the 10 m grid (0 = unlabelled)
samples = hc.load_samples()
train_s, test_s = samples[samples.split == "train"], samples[samples.split == "test"]
S = hc.s2_seasonal_composites()                  # {"winter", "spring", "summer", "autumn"} -> Dataset
rr, cc = np.indices((ROWS, COLS))
BLK = (rr // 250) * 10 + cc // 250               # 2.5 km block of every pixel (same ids as hc.block_id)
blocks = hc.load_blocks()
TEST_BLOCKS = blocks.block_id[blocks.split == "test"].to_numpy()
in_test = np.isin(BLK, TEST_BLOCKS)
print(f"{len(train_s):,} training and {len(test_s):,} test samples; grid {COLS} x {ROWS} pixels at {hc.RES} m")

# Pixel features: 10 bands and 3 indices per season = 52 layers, float32, NaN where a season had no clear view
FEAT = {}
for s, ds in S.items():
    b = {k: ds[k].values.astype("float32") for k in hc.S2_BANDS}
    FEAT.update({f"{k}_{s}": b[k] for k in hc.S2_BANDS})
    FEAT[f"ndvi_{s}"] = nd(b["B08"], b["B04"])
    FEAT[f"ndmi_{s}"] = nd(b["B08"], b["B11"])
    FEAT[f"ndwi_{s}"] = nd(b["B03"], b["B08"])
FNAMES = list(FEAT)
F = np.stack([FEAT[k] for k in FNAMES])          # (52, rows, cols)
del FEAT
print(f"{len(FNAMES)} pixel feature layers ({F.nbytes / 1e9:.2f} GB in memory)")

# --- Step 2: the segmentation input ------------------------------------------------------------------
section("Step 2: segmentation input")
SEG_LAYERS = ["ndvi_spring", "ndvi_summer", "ndmi_spring", "B08_summer", "B11_summer", "B04_summer"]


def scaled(a):
    """Stretch a layer to 0-1 between its 2nd and 98th percentiles; fill gaps with the median."""
    lo, hi = np.nanpercentile(a, [2, 98])
    a = np.clip((a - lo) / (hi - lo), 0, 1)
    return np.where(np.isnan(a), np.nanmedian(a), a).astype("float32")


IMG = np.dstack([scaled(F[FNAMES.index(k)]) for k in SEG_LAYERS])
print(f"segmentation image: {IMG.shape[1]} x {IMG.shape[0]} pixels x {IMG.shape[2]} layers ({', '.join(SEG_LAYERS)})")


def cut_by_blocks(seg):
    """Split any object that crosses a 2.5 km block edge, so no object is shared by training and test blocks."""
    key = seg.astype(np.int64) * 100 + BLK
    u, inv = np.unique(key.ravel(), return_inverse=True)
    return inv.reshape(seg.shape).astype(np.int32), len(u)


def segment(method, param):
    if method == "slic":                         # param = target mean object size in hectares
        seg = slic(IMG, n_segments=int(ROWS * COLS / (param / PX_HA)), compactness=0.1, channel_axis=-1,
                   convert2lab=False, start_label=0)
    else:                                        # param = Felzenszwalb scale (larger = bigger objects)
        with warnings.catch_warnings():         # six layers is intended: a multichannel image
            warnings.simplefilter("ignore", RuntimeWarning)
            seg = felzenszwalb(IMG, scale=param, sigma=0.5, min_size=20, channel_axis=-1)
    return cut_by_blocks(seg)


# --- Step 3: per-object features with NumPy label aggregation -----------------------------------------
PCTL_LAYERS = [f"ndvi_{s}" for s in hc.SEASONS] + [f"ndmi_{s}" for s in hc.SEASONS] + \
              ["B08_summer", "B11_summer", "B04_spring"]
GLCM_LAYERS = ["B08_summer", "ndvi_spring"]


def object_stats(seg, n):
    """Mean and standard deviation of every layer, and 10th/50th/90th percentiles of a subset, per object."""
    flat = seg.ravel()
    cnt = np.bincount(flat, minlength=n)
    start = np.concatenate([[0], np.cumsum(cnt)[:-1]])
    out = {}
    for i, name in enumerate(FNAMES):
        v = F[i].ravel()
        ok = ~np.isnan(v)
        nv = np.bincount(flat[ok], minlength=n)
        s1 = np.bincount(flat[ok], weights=v[ok], minlength=n)
        s2 = np.bincount(flat[ok], weights=v[ok].astype(float) ** 2, minlength=n)
        with np.errstate(divide="ignore", invalid="ignore"):
            mean = s1 / nv
            out[f"{name}_mean"] = mean
            out[f"{name}_std"] = np.sqrt(np.maximum(s2 / nv - mean ** 2, 0))
        if name in PCTL_LAYERS:
            sv = v[np.lexsort((v, flat))]        # sorted by object, then value (NaN last within each object)
            for q in (10, 50, 90):
                k = start + np.floor(q / 100 * np.maximum(nv - 1, 0)).astype(int)
                out[f"{name}_p{q}"] = np.where(nv > 0, sv[np.minimum(k, flat.size - 1)], np.nan)
    return out


def object_shape(seg, n):
    """Area, perimeter, compactness (4 pi A / P^2) and elongation from second moments, in pixel units."""
    flat = seg.ravel()
    area = np.bincount(flat, minlength=n).astype(float)
    per = np.zeros(n)
    for a, b in ((seg[:, :-1], seg[:, 1:]), (seg[:-1, :], seg[1:, :])):
        d = a != b
        per += np.bincount(a[d], minlength=n) + np.bincount(b[d], minlength=n)
    for edge in (seg[0], seg[-1], seg[:, 0], seg[:, -1]):
        per += np.bincount(edge, minlength=n)
    r, c = rr.ravel().astype(float), cc.ravel().astype(float)
    mr, mc = np.bincount(flat, r, n) / area, np.bincount(flat, c, n) / area
    vr = np.bincount(flat, r * r, n) / area - mr ** 2 + 1 / 12
    vc = np.bincount(flat, c * c, n) / area - mc ** 2 + 1 / 12
    cv = np.bincount(flat, r * c, n) / area - mr * mc
    root = np.sqrt((vr - vc) ** 2 + 4 * cv ** 2)
    return {"area_ha": area * PX_HA, "perimeter_m": per * hc.RES, "compactness": 4 * np.pi * area / per ** 2,
            "elongation": np.sqrt((vr + vc + root) / np.maximum(vr + vc - root, 1 / 6))}


def object_glcm(seg, n, name, levels=16):
    """Grey-level co-occurrence (GLCM) texture inside each object, from horizontal and vertical neighbour pairs.

    All objects at once: quantise the layer to 16 levels, count (object, level_i, level_j) pairs with one
    bincount, make the matrix symmetric and normalise it, then compute Haralick measures per object.
    """
    a = F[FNAMES.index(name)]
    edges = np.nanpercentile(a, np.linspace(0, 100, levels + 1)[1:-1])
    q = np.digitize(a, edges).astype(np.int64)
    q[np.isnan(a)] = -1
    counts = np.zeros(n * levels * levels, np.int64)
    for (sa, sb), (qa, qb) in (((seg[:, :-1], seg[:, 1:]), (q[:, :-1], q[:, 1:])),
                               ((seg[:-1, :], seg[1:, :]), (q[:-1, :], q[1:, :]))):
        same = (sa == sb) & (qa >= 0) & (qb >= 0)
        counts += np.bincount((sa[same].astype(np.int64) * levels + qa[same]) * levels + qb[same],
                              minlength=n * levels * levels)
    P = counts.reshape(n, levels, levels).astype(float)
    P = P + P.transpose(0, 2, 1)
    tot = P.sum((1, 2), keepdims=True)
    P = np.divide(P, tot, out=np.zeros_like(P), where=tot > 0)
    i, j = np.meshgrid(np.arange(levels), np.arange(levels), indexing="ij")
    mu = (P * i).sum((1, 2))
    var = (P * (i - mu[:, None, None]) ** 2).sum((1, 2))
    with np.errstate(divide="ignore", invalid="ignore"):
        feats = {
            "contrast": (P * (i - j) ** 2).sum((1, 2)),
            "homogeneity": (P / (1 + (i - j) ** 2)).sum((1, 2)),
            "entropy": -(P * np.log2(np.where(P > 0, P, 1))).sum((1, 2)),
            "asm": (P ** 2).sum((1, 2)),
            "correlation": ((P * (i - mu[:, None, None]) * (j - mu[:, None, None])).sum((1, 2)) / var),
        }
    empty = tot.ravel() == 0
    return {f"glcm_{name}_{k}": np.where(empty, np.nan, v) for k, v in feats.items()}


def object_features(seg, n):
    f = object_stats(seg, n)
    f.update(object_shape(seg, n))
    for name in GLCM_LAYERS:
        f.update(object_glcm(seg, n, name))
    return pd.DataFrame(f)


def family(col):
    if col.startswith("glcm_"):
        return "Texture (GLCM)"
    if col in ("area_ha", "perimeter_m", "compactness", "elongation"):
        return "Shape"
    if col.endswith("_mean"):
        return "Spectral mean"
    return "Spread and percentiles"


# --- Step 4: classifiers and evaluation helpers -------------------------------------------------------
PARAMS = dict(n_estimators=300, learning_rate=0.05, num_leaves=15, min_child_samples=10, subsample=0.8,
              subsample_freq=1, colsample_bytree=0.5, reg_lambda=1.0, random_state=hc.SEED, verbose=-1)


def model():
    return lgb.LGBMClassifier(**PARAMS)


def training_objects(seg, smp):
    """One row per object that contains training samples: its majority sample class and its block fold."""
    df = pd.DataFrame({"obj": seg[smp.row, smp.col], "cls": smp.class_id.to_numpy(), "fold": smp.fold.to_numpy()})
    g = df.groupby("obj").agg(cls=("cls", lambda s: s.mode().iloc[0]), fold=("fold", "first"), n=("cls", "size"))
    return g


def cv_objects(seg, OF, cols=None):
    """Blocked 5-fold CV on the training blocks, scored on the training SAMPLES (each takes its object's class)."""
    g = training_objects(seg, train_s)
    X = OF.loc[g.index] if cols is None else OF.loc[g.index, cols]
    sample_obj = seg[train_s.row, train_s.col]
    acc = []
    for k in range(5):
        m = model().fit(X[g.fold != k], g.cls[g.fold != k])
        val = (train_s.fold == k).to_numpy()
        objs = np.unique(sample_obj[val])
        pred = pd.Series(m.predict(X.loc[objs]), index=objs)
        acc.append((pred.loc[sample_obj[val]].to_numpy() == train_s.class_id.to_numpy()[val]).mean())
    return np.mean(acc), np.std(acc)


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
    pa = np.diag(m) / m.sum(1)
    ua = np.diag(m) / np.where(m.sum(0) == 0, np.nan, m.sum(0))
    with np.errstate(invalid="ignore"):
        f1 = np.nan_to_num(2 * pa * ua / (pa + ua))
    return {"oa": oa, "kappa": (oa - pe_) / (1 - pe_), "pa": pa, "ua": ua, "f1": f1}


def block_bootstrap(correct, block, reps=2000):
    rng = np.random.default_rng(hc.SEED)
    ub = np.unique(block)
    idx_by_block = {b: np.flatnonzero(block == b) for b in ub}
    boot = [correct[np.concatenate([idx_by_block[b] for b in rng.choice(ub, ub.size)])].mean() for _ in range(reps)]
    return np.percentile(boot, [2.5, 97.5])


def evaluate(pred_map, name):
    """Accuracy at the shared test samples: each test pixel takes the class of the map at its location."""
    yt = test_s.class_key.to_numpy()
    yp = np.array([ID_KEY.get(v, "none") for v in pred_map[test_s.row, test_s.col]])
    cm = confusion(yt, yp)
    res = summarise(cm)
    res["ci"] = block_bootstrap(yt == yp, test_s.block_id.to_numpy())
    res["macro_f1"] = float(np.mean(res["f1"]))
    res["cm"] = cm
    valid = in_test & (labels > 0) & (pred_map > 0)
    res["w2w"] = (labels[valid] == pred_map[valid]).mean()
    print(f"{name:<34} OA {res['oa']:.3f} (95% CI {res['ci'][0]:.3f}-{res['ci'][1]:.3f})  macro F1 "
          f"{res['macro_f1']:.3f}  kappa {res['kappa']:.3f}  wall-to-wall (test blocks) {res['w2w']:.3f}")
    return res


def noise(m):
    """Salt-and-pepper metrics: share of pixels outvoted by their 3 x 3 neighbourhood, and patch counts."""
    counts = np.stack([ndimage.uniform_filter((m == c).astype("float32"), 3, mode="nearest") for c in hc.CLASS_IDS])
    own = np.take_along_axis(counts, (np.clip(m, 1, 9) - 1)[None].astype(np.int64), 0)[0]
    outvoted = (own < counts.max(0) - 1e-6) & (m > 0)
    eight = np.ones((3, 3), bool)
    patches = sum(ndimage.label(m == c, structure=eight)[1] for c in hc.CLASS_IDS)
    single = sum(int((np.bincount(ndimage.label(m == c, structure=eight)[0].ravel())[1:] == 1).sum())
                 for c in hc.CLASS_IDS)
    area_km2 = (m > 0).sum() * hc.RES ** 2 / 1e6
    return {"outvoted_3x3": float(outvoted.sum() / (m > 0).sum()), "patches": int(patches),
            "patches_per_km2": float(patches / area_km2), "single_pixel_patches": single,
            "mean_patch_ha": float((m > 0).sum() * PX_HA / patches)}


# --- Step 5: tune the segmentation ------------------------------------------------------------------
section("Step 5: segmentation candidates, scored on the training blocks only")
train_px = ~in_test & (labels > 0)
near_block_edge = np.isin(rr % 250, (0, 249)) | np.isin(cc % 250, (0, 249))


def edges(a):
    b = np.zeros(a.shape, bool)
    d = a[:, 1:] != a[:, :-1]
    b[:, 1:] |= d
    b[:, :-1] |= d
    d = a[1:, :] != a[:-1, :]
    b[1:, :] |= d
    b[:-1, :] |= d
    return b


LE_EDGE = edges(labels) & train_px & ~near_block_edge   # Living England class boundaries (both sides, 1 px)
print(f"Living England class-boundary pixels in training blocks: {LE_EDGE.sum():,}")


def segmentation_quality(seg, n):
    """Boundary recall (LE boundaries within 10 m of an object edge) and achievable segmentation accuracy (ASA)."""
    recall = (LE_EDGE & ndimage.binary_dilation(edges(seg), np.ones((3, 3), bool))).sum() / LE_EDGE.sum()
    cm = np.bincount(seg[train_px].astype(np.int64) * 10 + labels[train_px], minlength=n * 10).reshape(n, 10)
    asa = cm.max(1).sum() / cm.sum()             # best wall-to-wall agreement if every object took its majority class
    return recall, asa


CANDIDATES = [("slic", ha) for ha in (0.25, 0.5, 1, 2, 4, 8)] + [("felzenszwalb", s) for s in (10, 25, 50, 100, 200, 400)]
rows, SEGS = [], {}
for method, param in CANDIDATES:
    t = time.perf_counter()
    seg, n = segment(method, param)
    t_seg = time.perf_counter() - t
    recall, asa = segmentation_quality(seg, n)
    OF = object_features(seg, n)
    cv_mean, cv_sd = cv_objects(seg, OF)
    rows.append({"method": method, "param": param, "objects": n, "mean_ha": ROWS * COLS * PX_HA / n,
                 "boundary_recall": recall, "asa": asa, "cv_oa": cv_mean, "cv_sd": cv_sd, "seg_s": t_seg})
    SEGS[(method, param)] = (seg, n, OF)
    print(f"{method:<12} {param:>5g}  {n:>7,} objects, mean {ROWS * COLS * PX_HA / n:5.2f} ha  boundary recall "
          f"{recall:.3f}  ASA {asa:.3f}  blocked CV OA {cv_mean:.3f} +/- {cv_sd:.3f}  ({t_seg:.0f} s)")
tune = pd.DataFrame(rows)
best = tune.loc[tune.cv_oa.idxmax()]
BEST = (best.method, best.param)
seg, n_obj, OF = SEGS[BEST]
SHOW_SEGS = [k for k in (("slic", 0.5), ("felzenszwalb", 100)) if k != BEST][:1] + [BEST, ("slic", 4)]
for k in [k for k in SEGS if k not in SHOW_SEGS]:
    del SEGS[k]                                  # keep only what the figures need
print(f"chosen by blocked CV: {best.method} {best.param:g} ({n_obj:,} objects, mean {best.mean_ha:.2f} ha)")
sizes = np.bincount(seg.ravel()) * PX_HA
print(f"object size: median {np.median(sizes):.2f} ha, 5th-95th percentile {np.percentile(sizes, 5):.2f}-"
      f"{np.percentile(sizes, 95):.2f} ha; {OF.shape[1]} features per object")

# --- Step 6: which object features matter? ------------------------------------------------------------
section("Step 6: feature families (blocked CV on the training blocks)")
fam = pd.Series({c: family(c) for c in OF.columns})
print(fam.value_counts().to_string())
ablation = {}
for label, keep in (("Spectral mean only", ["Spectral mean"]),
                    ("+ spread and percentiles", ["Spectral mean", "Spread and percentiles"]),
                    ("+ shape", ["Spectral mean", "Spread and percentiles", "Shape"]),
                    ("+ texture (all features)", ["Spectral mean", "Spread and percentiles", "Shape", "Texture (GLCM)"])):
    cols = fam.index[fam.isin(keep)].tolist()
    ablation[label] = cv_objects(seg, OF, cols)
    print(f"{label:<28} {len(cols):>3} features  CV OA {ablation[label][0]:.3f} +/- {ablation[label][1]:.3f}")

# --- Step 7: train the object classifier and map every object -----------------------------------------
section("Step 7: object-based classification")
g = training_objects(seg, train_s)
print(f"{len(train_s):,} training samples fall in {len(g):,} objects "
      f"({(g.n > 1).sum()} objects hold more than one sample)")
t = time.perf_counter()
obj_model = model().fit(OF.loc[g.index], g.cls)
proba = obj_model.predict_proba(OF)
obj_class = obj_model.classes_[proba.argmax(1)].astype("uint8")
obj_conf = proba.max(1)
obia_map = obj_class[seg]
print(f"trained on {len(g):,} objects and classified {n_obj:,} objects in {time.perf_counter() - t:.1f} s")
hc.write_geotiff(OUT / "habitat_obia_2023.tif", obia_map, nodata=0)

# --- Step 8: the pixel-based comparison --------------------------------------------------------------
section("Step 8: pixel-based LightGBM on the same 52 seasonal features")
Xp = hc.sample_values(F, train_s)
t = time.perf_counter()
pix_model = model().fit(Xp, train_s.class_id)
pix_map = np.zeros((ROWS, COLS), "uint8")
for r0 in range(0, ROWS, 250):
    chunk = F[:, r0:r0 + 250].reshape(len(FNAMES), -1).T
    pix_map[r0:r0 + 250] = pix_model.predict(chunk).reshape(-1, COLS)
print(f"trained on {len(train_s):,} pixels and classified {ROWS * COLS:,} pixels in {time.perf_counter() - t:.1f} s")
pix_cv = []
for k in range(5):
    m = model().fit(Xp[train_s.fold != k], train_s.class_id[train_s.fold != k])
    pix_cv.append((m.predict(Xp[train_s.fold == k]) == train_s.class_id[train_s.fold == k]).mean())
print(f"pixel blocked CV OA {np.mean(pix_cv):.3f} +/- {np.std(pix_cv):.3f} (objects: {best.cv_oa:.3f} +/- {best.cv_sd:.3f})")
hc.write_geotiff(OUT / "habitat_pixel_lgbm_2023.tif", pix_map, nodata=0)
# A common shortcut: classify pixels, then give each object the majority vote of its pixels
votes = np.bincount(seg.ravel().astype(np.int64) * 10 + pix_map.ravel(), minlength=n_obj * 10).reshape(n_obj, 10)
vote_map = votes[:, 1:].argmax(1).astype("uint8")[seg] + 1

# --- Step 9: accuracy on the shared test samples -----------------------------------------------------
section("Step 9: accuracy on the 2,456 test samples (27 test blocks)")
R = {"pixel": evaluate(pix_map, "Pixel LightGBM"),
     "vote": evaluate(vote_map, "Pixel LightGBM + object majority"),
     "obia": evaluate(obia_map, "Object-based LightGBM (OBIA)")}
print("Part 1 baseline (spectral rules)   OA 0.600 (95% CI 0.527-0.664)  macro F1 0.630  kappa 0.547  "
      "wall-to-wall (test blocks) 0.483")
res = R["obia"]
cm_show = res["cm"].copy()
cm_show.index = [f"ref {k}" for k in KEYS]
cm_show.columns = [k.split("_")[0][:5] for k in KEYS]
print("\nconfusion matrix, object-based (rows = Living England reference, columns = OBIA map):")
print(cm_show.to_string())
per_class = pd.DataFrame({f"{m}_{s}": R[m][s] for m in ("pixel", "obia") for s in ("pa", "ua", "f1")},
                         index=[NAMES[k] for k in KEYS])
print("\nper-class producer's (pa) and user's (ua) accuracy and F1, pixel versus object:")
print(per_class.round(3).to_string())

# --- Step 10: salt-and-pepper noise ------------------------------------------------------------------
section("Step 10: map noise over the whole study area")
NOISE = {"Living England (reference)": noise(np.where(labels > 0, labels, 0).astype("uint8")),
         "Pixel LightGBM": noise(pix_map), "Pixel + object majority": noise(vote_map), "Object-based (OBIA)": noise(obia_map)}
noise_df = pd.DataFrame(NOISE).T
noise_df["outvoted_3x3"] = (100 * noise_df["outvoted_3x3"]).round(2)
print(noise_df.rename(columns={"outvoted_3x3": "outvoted_3x3_%"}).round(2).to_string())
area = pd.DataFrame({m: [(mp == c["id"]).sum() * hc.RES ** 2 / 1e6 for c in hc.CLASSES]
                     for m, mp in (("pixel_km2", pix_map), ("obia_km2", obia_map), ("living_england_km2", labels))},
                    index=[NAMES[k] for k in KEYS])
print("\nmapped area by class:")
print(area.round(1).to_string())

# --- Step 11: export habitat polygons ----------------------------------------------------------------
section("Step 11: export the objects as a GeoPackage of habitat polygons")
t = time.perf_counter()
geoms, ids = zip(*((as_shape(gm), int(v)) for gm, v in shapes(seg, transform=hc.transform())))
gdf = gpd.GeoDataFrame({"object_id": np.array(ids, "int32")}, geometry=list(geoms), crs=hc.CRS)
gdf = gdf.dissolve("object_id", as_index=False)  # an object cut by a block edge can have two parts
ids = gdf.object_id.to_numpy()
gdf["class_id"] = obj_class[ids]
gdf["class_key"] = [ID_KEY[v] for v in gdf.class_id]
gdf["habitat"] = [NAMES[k] for k in gdf.class_key]
gdf["confidence"] = obj_conf[ids].round(3)
gdf["area_ha"] = OF.area_ha.to_numpy()[ids].round(2)
gdf["ndvi_summer_mean"] = OF.ndvi_summer_mean.to_numpy()[ids].round(3)
gdf["block_split"] = np.where(np.isin(np.bincount(seg.ravel(), BLK.ravel(), n_obj)[ids] /
                                      np.bincount(seg.ravel(), minlength=n_obj)[ids], TEST_BLOCKS), "test", "train")
GPKG = OUT / "habitat_obia_2023.gpkg"
GPKG.unlink(missing_ok=True)
gdf.to_file(GPKG, layer="habitat_objects", driver="GPKG")
geoms, vals = zip(*((as_shape(gm), int(v)) for gm, v in shapes(obia_map, mask=obia_map > 0, transform=hc.transform())))
patches = gpd.GeoDataFrame({"class_id": np.array(vals, "uint8")}, geometry=list(geoms), crs=hc.CRS)
patches["habitat"] = [NAMES[ID_KEY[v]] for v in patches.class_id]
patches["area_ha"] = (patches.area / 1e4).round(2)
patches.to_file(GPKG, layer="habitat_patches", driver="GPKG")
print(f"{len(gdf):,} object polygons and {len(patches):,} dissolved habitat patches written to {GPKG.name} "
      f"({GPKG.stat().st_size / 1e6:.0f} MB) in {time.perf_counter() - t:.0f} s")
print(gdf.drop(columns="geometry").head(3).to_string(index=False))
RUNTIME = time.perf_counter() - T0
print(f"\nruntime up to the export: {RUNTIME / 60:.1f} min (composites were already cached by prepare_data.py)")

# --- metrics.json for the series comparison table -----------------------------------------------------
metrics = {
    "method": f"Object-based LightGBM on {best.method} segments ({best.param:g}, {n_obj} objects, mean {best.mean_ha:.2f} ha) "
              f"with per-object spectral statistics, percentiles, shape and GLCM texture from 2023 seasonal Sentinel-2",
    "oa": round(float(res["oa"]), 4), "oa_ci95": [round(float(v), 4) for v in res["ci"]],
    "macro_f1": round(res["macro_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(test_s)), "runtime_s": round(RUNTIME, 1),
    "wall_to_wall_test_blocks": round(float(res["w2w"]), 4),
    "cv_oa_train_blocks": [round(float(best.cv_oa), 4), round(float(best.cv_sd), 4)],
    "pixel_comparison": {m: {"oa": round(float(R[m]["oa"]), 4), "oa_ci95": [round(float(v), 4) for v in R[m]["ci"]],
                             "macro_f1": round(R[m]["macro_f1"], 4), "kappa": round(float(R[m]["kappa"]), 4),
                             "wall_to_wall_test_blocks": round(float(R[m]["w2w"]), 4)} for m in ("pixel", "vote")},
    "noise": {k: {kk: round(vv, 4) if isinstance(vv, float) else vv for kk, vv in v.items()} for k, v in NOISE.items()},
}
(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])
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]
le = hc.load_labels()
SHORT = {"slic": "SLIC", "felzenszwalb": "Felzenszwalb"}


def km_axes(ax, step=5000):
    x0, x1 = ax.get_xlim()
    y0, y1 = ax.get_ylim()
    ax.set_xticks(np.arange(np.ceil(x0 / step) * step, x1 + 1, step))
    ax.set_yticks(np.arange(np.ceil(y0 / step) * step, y1 + 1, step))
    ax.set_xticklabels([f"{v / 1000:g}" for v in ax.get_xticks()])
    ax.set_yticklabels([f"{v / 1000:g}" for v in ax.get_yticks()])
    ax.set_xlabel("Easting (km, British National Grid)")
    ax.set_ylabel("Northing (km)")


def window(x0, y0, x1, y1):
    c0, c1 = int((x0 - hc.AOI_BOUNDS[0]) / hc.RES), int((x1 - hc.AOI_BOUNDS[0]) / hc.RES)
    r0, r1 = int((hc.AOI_BOUNDS[3] - y1) / hc.RES), int((hc.AOI_BOUNDS[3] - y0) / hc.RES)
    return (slice(r0, r1), slice(c0, c1)), (x0, x1, y0, y1)


def overlay_edges(ax, lab, wext, colour):
    e = find_boundaries(lab, mode="inner")
    ax.imshow(np.ma.masked_where(~e, np.ones(e.shape)), cmap=matplotlib.colors.ListedColormap([colour]),
              extent=wext, interpolation="nearest", vmin=0, vmax=1)


# Figure 1: segmentations in a 2.3 x 1.7 km window in the Edale valley (inside one block: no block-edge cuts)
W1, wext1 = window(410100, 385100, 412400, 386800)
fig, axes = plt.subplots(2, 2, figsize=(11.5, 8.6), sharex=True, sharey=True)
axes = axes.ravel()
axes[0].imshow(rgb[W1], extent=wext1)
le.cx[wext1[0]:wext1[1], wext1[2]:wext1[3]].boundary.clip_by_rect(*[wext1[i] for i in (0, 2, 1, 3)]).plot(
    ax=axes[0], color="#fdae61", lw=0.7)
axes[0].set_title("Summer 2023 composite with Living England polygon boundaries")
panels = [(k, "#ffffbf") for k in SHOW_SEGS]
for ax, (key, col) in zip(axes[1:], panels):
    s_, n_, _ = SEGS[key]
    ax.imshow(rgb[W1], extent=wext1)
    overlay_edges(ax, s_[W1], wext1, col)
    row = tune[(tune.method == key[0]) & (tune.param == key[1])].iloc[0]
    tag = " (chosen)" if key == BEST else ""
    unit = " ha target" if key[0] == "slic" else " scale"
    ax.set_title(f"{SHORT[key[0]]} {key[1]:g}{unit}{tag}: mean {row.mean_ha:.2f} ha, boundary recall "
                 f"{row.boundary_recall:.0%}", fontsize=8.5)
for ax in axes:
    ax.set_xlim(wext1[0], wext1[1])
    ax.set_ylim(wext1[2], wext1[3])
    km_axes(ax, 500)
for ax in axes[:2]:
    ax.set_xlabel("")
for ax in axes[1::2]:
    ax.set_ylabel("")
fig.tight_layout()
fig.savefig(OUT / "01-segmentation-zoom.png")
plt.close(fig)

# Figure 2: segment tuning
fig, axes = plt.subplots(1, 3, figsize=(13, 4.2))
for method, col, mk in (("slic", "#2b83ba", "o"), ("felzenszwalb", "#d7191c", "s")):
    d = tune[tune.method == method].sort_values("mean_ha")
    for ax, ycol in zip(axes, ("boundary_recall", "asa", "cv_oa")):
        ax.plot(d.mean_ha, d[ycol], marker=mk, color=col, label=SHORT[method])
        if ycol == "cv_oa":
            ax.fill_between(d.mean_ha, d.cv_oa - d.cv_sd, d.cv_oa + d.cv_sd, color=col, alpha=0.15)
axes[2].axhline(np.mean(pix_cv), color="#555555", ls="--", lw=1)
axes[2].text(tune.mean_ha.max(), np.mean(pix_cv) + 0.003, f"pixel LightGBM {np.mean(pix_cv):.3f}", fontsize=7.5,
             color="#555555", ha="right", va="bottom")
axes[2].plot(best.mean_ha, best.cv_oa, "*", ms=14, color="#fdae61", mec="black", label="chosen")
titles = ("Boundary recall: Living England edges\nwithin 10 m of an object edge",
          "Achievable accuracy (ASA): agreement if each\nobject took its majority Living England class",
          "Blocked 5-fold CV accuracy of the\nobject classifier (training blocks)")
for ax, t_ in zip(axes, titles):
    ax.set_xscale("log")
    ax.set_xlabel("Mean object size (ha, log scale)")
    ax.set_title(t_, fontsize=8.5)
    ax.grid(alpha=0.3)
axes[0].legend(fontsize=8)
axes[2].legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(OUT / "02-segment-tuning.png")
plt.close(fig)

# Figure 3: feature importance
imp = pd.Series(obj_model.booster_.feature_importance("gain"), index=OF.columns)
imp = imp / imp.sum()
fam_share = imp.groupby(fam).sum().sort_values()
top = imp.sort_values().tail(15)
FAM_COL = {"Spectral mean": "#2b83ba", "Spread and percentiles": "#abdda4", "Shape": "#fdae61", "Texture (GLCM)": "#d7191c"}
fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.8), gridspec_kw={"width_ratios": [1.5, 1]})
axes[0].barh(top.index, top.values, color=[FAM_COL[fam[c]] for c in top.index], edgecolor="#333333", lw=0.4)
axes[0].set_xlabel("Share of total gain")
axes[0].set_title("The 15 most important object features", fontsize=9)
axes[0].tick_params(axis="y", labelsize=7.5)
labels_abl = list(ablation)
axes[1].barh(range(len(labels_abl)), [ablation[k][0] for k in labels_abl], xerr=[ablation[k][1] for k in labels_abl],
             color="#abdda4", edgecolor="#333333", lw=0.4, capsize=3)
lo_x = min(v[0] - v[1] for v in ablation.values()) - 0.03
for i, k in enumerate(labels_abl):
    axes[1].text(lo_x + 0.003, i, f"{ablation[k][0]:.3f}", va="center", ha="left", fontsize=8)
axes[1].set_yticks(range(len(labels_abl)), labels_abl, fontsize=8)
axes[1].invert_yaxis()
axes[1].set_xlim(lo_x, max(v[0] + v[1] for v in ablation.values()) + 0.01)
axes[1].set_xlabel("Blocked CV accuracy (mean and s.d. of 5 folds)")
axes[1].set_title("Adding feature families one at a time\n(share of gain: " +
                  ", ".join(f"{k.split(' (')[0].lower()} {v:.0%}" for k, v in fam_share.sort_values(ascending=False).items()) + ")",
                  fontsize=8.5)
axes[0].legend(handles=[Patch(color=c, label=k) for k, c in FAM_COL.items()], fontsize=7.5, loc="lower right")
fig.tight_layout()
fig.savefig(OUT / "03-object-features.png")
plt.close(fig)

# Figure 4: wall-to-wall OBIA map against Living England
fig, axes = plt.subplots(1, 2, figsize=(13, 5.6), sharey=True)
axes[0].imshow(np.ma.masked_equal(obia_map, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
axes[0].set_title(f"Object-based map ({n_obj:,} objects), test-sample OA {res['oa']:.1%}")
axes[1].imshow(np.ma.masked_equal(labels, 0), cmap=cmap, norm=norm, extent=ext, interpolation="nearest")
both = (labels > 0) & (obia_map > 0)
axes[1].set_title(f"Living England 2022-23 (pixel agreement with OBIA: {(labels[both] == obia_map[both]).mean():.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-obia-map-vs-living-england.png")
plt.close(fig)

# Figure 5: salt and pepper, zoomed: the Edale valley and the southern edge of Kinder Scout
W2, wext2 = window(409000, 384000, 413000, 387000)
fig, axes = plt.subplots(2, 2, figsize=(11.5, 9), sharex=True, sharey=True)
axes = axes.ravel()
axes[0].imshow(rgb[W2], extent=wext2)
axes[0].set_title("Summer 2023 Sentinel-2 composite")
for ax, (mp, t_) in zip(axes[1:], ((pix_map, "Pixel LightGBM"), (obia_map, "Object-based (OBIA)"),
                                   (labels, "Living England 2022-23"))):
    ax.imshow(np.ma.masked_equal(mp[W2], 0), cmap=cmap, norm=norm, extent=wext2, interpolation="nearest")
    if mp is not labels:
        nz = noise(mp[W2])
        t_ += f": {nz['outvoted_3x3']:.1%} of pixels outvoted, {nz['patches_per_km2']:.0f} patches per km²"
    ax.set_title(t_, fontsize=8.5)
for ax in axes:
    km_axes(ax, 1000)
for ax in axes[:2]:
    ax.set_xlabel("")
for ax in axes[1::2]:
    ax.set_ylabel("")
fig.legend(handles=legend, loc="lower center", ncol=5, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.06))
fig.tight_layout()
fig.savefig(OUT / "05-salt-and-pepper-zoom.png")
plt.close(fig)

# Figure 6: confusion matrix of the object-based map
m = res["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("Object-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 {res['ci'][0]:.1%} to "
             f"{res['ci'][1]:.1%}), kappa {res['kappa']:.2f}", fontsize=9)
fig.savefig(OUT / "06-confusion-matrix.png")
plt.close(fig)

# Figure 7: pixel versus object, accuracy and noise
fig, axes = plt.subplots(1, 3, figsize=(13, 4.3))
names = ["Part 1 rules", "Pixel LightGBM", "Pixel + object\nmajority", "Object-based\n(OBIA)"]
oas = [0.600, R["pixel"]["oa"], R["vote"]["oa"], R["obia"]["oa"]]
cis = [(0.527, 0.664), R["pixel"]["ci"], R["vote"]["ci"], R["obia"]["ci"]]
f1s = [0.630, R["pixel"]["macro_f1"], R["vote"]["macro_f1"], R["obia"]["macro_f1"]]
cols = ["#bababa", "#d7191c", "#fdae61", "#2b83ba"]
x = np.arange(4)
axes[0].bar(x - 0.18, oas, 0.36, yerr=np.array([[o - c[0], c[1] - o] for o, c in zip(oas, cis)]).T, capsize=3,
            color=cols, edgecolor="#333333", lw=0.4, label="Overall accuracy (block-bootstrap 95% CI)")
axes[0].bar(x + 0.18, f1s, 0.36, color=cols, alpha=0.55, edgecolor="#333333", lw=0.4, hatch="///", label="Macro F1")
for xi, v in zip(x, oas):
    axes[0].text(xi - 0.18, 0.05, f"{v:.3f}", ha="center", fontsize=7.5, rotation=90, color="black")
axes[0].set_xticks(x, names, fontsize=7.5)
axes[0].set_ylim(0, 1)
axes[0].set_title("Accuracy on the 2,456 test samples", fontsize=9)
axes[0].legend(fontsize=7, loc="upper left")
nk = ["Pixel LightGBM", "Pixel + object majority", "Object-based (OBIA)", "Living England (reference)"]
nnames = ["Pixel\nLightGBM", "Pixel + object\nmajority", "Object-based\n(OBIA)", "Living\nEngland"]
ncols = ["#d7191c", "#fdae61", "#2b83ba", "#bababa"]
for ax, key, t_, fmt in ((axes[1], "outvoted_3x3", "Pixels outvoted by their 3 x 3 neighbours", "{:.1%}"),
                         (axes[2], "patches_per_km2", "Habitat patches per km²", "{:,.0f}")):
    vals = [NOISE[k][key] for k in nk]
    ax.bar(range(4), vals, color=ncols, edgecolor="#333333", lw=0.4)
    for i, v in enumerate(vals):
        ax.text(i, v, fmt.format(v), ha="center", va="bottom", fontsize=8)
    ax.set_xticks(range(4), nnames, fontsize=7.5)
    ax.set_title(t_ + ", whole study area", fontsize=9)
    ax.set_ylim(0, max(vals) * 1.15)
axes[1].yaxis.set_major_formatter(matplotlib.ticker.PercentFormatter(1.0, decimals=0))
fig.tight_layout()
fig.savefig(OUT / "07-pixel-vs-object.png")
plt.close(fig)
print(f"saved 7 figures, 2 GeoTIFFs, {GPKG.name} and metrics.json to {OUT.name}/")