"""Land cover classification with Sentinel-2 and a scikit-learn random forest: Sheffield and the Peak District.

Builds seasonal Sentinel-2 L2A median composites for 2021 with odc-stac, samples training labels from
ESA WorldCover 2021 with a stratified random sample, trains a random forest on bands and spectral indices,
assesses it with a spatially blocked train/test split and blocked cross-validation (and shows how a random
split overstates accuracy), ranks the features with permutation importance and predicts a full land cover map.
Run with:  python land_cover_random_forest.py
"""
import platform
import time
import warnings
from importlib.metadata import version
from pathlib import Path

import dask
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import odc.geo.xr  # noqa: F401  (adds the .odc accessor to xarray objects)
import odc.stac
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import sklearn
import xarray as xr
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.patches import Patch
from scipy import ndimage
from sklearn.ensemble import RandomForestClassifier
from sklearn.inspection import permutation_importance
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
from sklearn.model_selection import GroupKFold, KFold, cross_val_predict

warnings.filterwarnings("ignore", category=RuntimeWarning)   # all-NaN medians where a season has no clear view
HERE = Path(__file__).parent
OUT, DATA = HERE / "output", HERE / "data"
for d in (OUT, DATA):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)
SEED = 42

print("=== versions ===")
print(f"python {platform.python_version()} | scikit-learn {sklearn.__version__} | odc-stac {odc.stac.__version__} | "
      f"xarray {xr.__version__} | dask {dask.__version__} | pystac-client {pystac_client.__version__} | "
      f"planetary-computer {version('planetary-computer')} | numpy {np.__version__}")


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: search Sentinel-2 for one growing season ---------------------------------------------
BBOX = (-1.75, 53.25, -1.05, 53.45)        # Peak District moors, Sheffield and the farmland to the east
CRS, RES = "EPSG:32630", 20                # UTM zone 30N at 20 m: the native grid of the red-edge and SWIR bands
YEAR = 2021                                # match the year of the ESA WorldCover map used for the labels
SEASONS = {"spring": [3, 4, 5], "summer": [6, 7, 8], "autumn": [9, 10]}
BANDS = ["B02", "B03", "B04", "B05", "B06", "B07", "B08", "B8A", "B11", "B12"]
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"

catalog = pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace)
items = catalog.search(collections=["sentinel-2-l2a"], bbox=BBOX, datetime=f"{YEAR}-03-01/{YEAR}-10-31",
                       query={"eo:cloud_cover": {"lt": 60}}).item_collection()
best = {}   # keep the newest processing baseline when a date and tile exist twice
for it in items:
    key = (it.datetime.date(), it.properties["s2:mgrs_tile"])
    if key not in best or it.properties["s2:processing_baseline"] > best[key].properties["s2:processing_baseline"]:
        best[key] = it
items = sorted(best.values(), key=lambda i: i.datetime)
section("STAC search")
print(f"{len(items)} Sentinel-2 L2A items (tiles {sorted({i.properties['s2:mgrs_tile'] for i in items})}) "
      f"with < 60% cloud, {items[0].datetime:%d %b} to {items[-1].datetime:%d %b %Y}")

# --- Step 2: lazy datacube, cloud mask and seasonal median composites -----------------------------
odc.stac.configure_rio(cloud_defaults=True)
ds = odc.stac.load(items, bands=BANDS + ["SCL"], bbox=BBOX, crs=CRS, resolution=RES,
                   chunks={"time": 1, "x": 2048, "y": 2048}, groupby="solar_day", dtype="uint16", nodata=0,
                   resampling={"*": "average", "SCL": "nearest"})
ds = ds.assign_coords(time=ds.time.dt.floor("D"))
gbox = ds.odc.geobox
print(f"lazy cube: {dict(ds.sizes)}, {ds.nbytes / 1e9:.1f} GB if loaded in full")
dask.config.set(scheduler="threads", num_workers=16)

baseline = {}
for it in items:
    baseline.setdefault(pd.Timestamp(it.datetime.date()), set()).add(float(it.properties["s2:processing_baseline"]) >= 4.0)
# The two tiles of one date are mosaicked, so they must share the same offset convention.
mixed = [t for t, v in baseline.items() if len(v) > 1]
ds = ds.drop_sel(time=mixed)
print(f"dropped {len(mixed)} dates whose tiles mix processing baselines before and after 04.00")
offset = xr.DataArray([1000 if baseline[pd.Timestamp(t)] == {True} else 0 for t in ds.time.values],
                      dims="time", coords={"time": ds.time})

comp_file = DATA / f"composites_{YEAR}.nc"
if comp_file.exists():
    comp = xr.load_dataset(comp_file)
else:
    t0 = time.time()
    scl = ds["SCL"].compute()
    cloudy = np.isin(scl.values, [3, 8, 9, 10])                          # shadow, cloud (medium, high), cirrus
    grown = ndimage.binary_dilation(cloudy, iterations=2, axes=(1, 2))   # grow by 40 m
    clear = scl.isin([4, 5, 6, 7]) & ~xr.DataArray(grown, coords=scl.coords, dims=scl.dims)
    frac = clear.mean(("y", "x")).to_series()
    print(f"SCL read in {time.time() - t0:.0f} s; {(frac >= 0.1).sum()} of {len(frac)} dates are at least 10% clear")
    layers = {}
    for season, months in SEASONS.items():
        times = frac.index[(frac >= 0.1) & frac.index.month.isin(months)]
        sub = ds[BANDS].sel(time=times)
        refl = ((sub.astype("float32") - offset.sel(time=times)) / 10000).where(clear.sel(time=times))
        med = refl.chunk({"time": -1, "y": 512, "x": 512}).median("time").compute()
        n = clear.sel(time=times).sum("time")
        print(f"{season:<7} {len(times):>2} dates; pixels with no clear view: {float((n == 0).mean()):.2%}")
        for b in BANDS:
            layers[f"{season}_{b}"] = med[b]
    comp = xr.Dataset(layers).drop_vars("spatial_ref", errors="ignore")
    comp.to_netcdf(comp_file)
    print(f"composites built in {time.time() - t0:.0f} s")

# --- Step 3: features: bands plus indices for each season ----------------------------------------
def nd(a, b):
    return (a - b) / (a + b)


features = {}
for s in SEASONS:
    for b in BANDS:
        features[f"{s}_{b}"] = comp[f"{s}_{b}"]
    features[f"{s}_NDVI"] = nd(comp[f"{s}_B08"], comp[f"{s}_B04"])      # vegetation
    features[f"{s}_MNDWI"] = nd(comp[f"{s}_B03"], comp[f"{s}_B11"])     # open water
    features[f"{s}_NDBI"] = nd(comp[f"{s}_B11"], comp[f"{s}_B08"])      # built-up and bare surfaces
    features[f"{s}_NDRE"] = nd(comp[f"{s}_B8A"], comp[f"{s}_B05"])      # red-edge chlorophyll
FEATURES = list(features)
X_all = np.stack([features[f].values.astype("float32") for f in FEATURES], axis=-1)   # (y, x, feature)
H, W, F = X_all.shape
section("feature stack")
print(f"{H} x {W} pixels of {RES} m, {F} features ({len(SEASONS)} seasons x ({len(BANDS)} bands + 4 indices))")
print(f"pixels with at least one missing feature: {np.isnan(X_all).any(-1).mean():.2%}")

# --- Step 4: reference labels from ESA WorldCover 2021 -----------------------------------------------
CLASSES = {10: "Tree cover", 20: "Shrubland", 30: "Grassland", 40: "Cropland", 50: "Built-up",
           60: "Bare / sparse", 80: "Water", 90: "Herbaceous wetland", 100: "Moss and lichen"}
WC_COLOURS = {10: "#006400", 20: "#ffbb22", 30: "#ffff4c", 40: "#f096ff", 50: "#fa0000", 60: "#b4b4b4",
              80: "#0064c8", 90: "#0096a0", 100: "#fae6a0"}
wc_items = catalog.search(collections=["esa-worldcover"], bbox=BBOX,
                          query={"esa_worldcover:product_version": {"eq": "2.0.0"}}).item_collection()
wc = odc.stac.load(wc_items, bands=["map"], geobox=gbox, resampling="mode").squeeze("time")["map"].values
section("ESA WorldCover 2021 (resampled to 20 m with the mode)")
share = pd.Series(wc.ravel()).value_counts(normalize=True).sort_index()
print((share * 100).round(2).rename(lambda c: f"{c:>3} {CLASSES.get(c, '?')}").to_string())
MIN_SHARE = 0.005                         # drop classes covering less than 0.5% of the area
USE = [c for c in CLASSES if share.get(c, 0) >= MIN_SHARE]
print(f"classes used: {[CLASSES[c] for c in USE]}; dropped: "
      f"{[CLASSES[c] for c in CLASSES if c not in USE and share.get(c, 0) > 0]}")

# --- Step 5: stratified random sample of 'pure' pixels -----------------------------------------------
N_PER_CLASS = 2000
BLOCK = 2000                              # 2 km x 2 km spatial blocks for the train/test split
rng = np.random.default_rng(SEED)
valid = ~np.isnan(X_all).any(-1)
rows, cols, ys = [], [], []
for c in USE:
    # Keep pixels whose 3 x 3 neighbourhood is all the same class: fewer mixed pixels and edge errors.
    pure = ndimage.binary_erosion(wc == c, structure=np.ones((3, 3))) & valid
    r, cc = np.nonzero(pure)
    pick = rng.choice(len(r), size=min(N_PER_CLASS, len(r)), replace=False)
    rows.append(r[pick]); cols.append(cc[pick]); ys.append(np.full(len(pick), c))
rows, cols, y = np.concatenate(rows), np.concatenate(cols), np.concatenate(ys)
X = X_all[rows, cols]
block_px = BLOCK // RES
block_id = (rows // block_px) * (W // block_px + 1) + cols // block_px
section("stratified random sample")
print(pd.Series(y).map(CLASSES).value_counts().reindex([CLASSES[c] for c in USE]).to_string())
print(f"{len(y)} samples in {len(np.unique(block_id))} blocks of {BLOCK / 1000:.0f} km")

# --- Step 6: spatially blocked train/test split, and a random split for comparison ------------------
blocks = np.unique(block_id)
fold_of_block = dict(zip(blocks, rng.permutation(len(blocks)) % 5))
fold = np.array([fold_of_block[b] for b in block_id])
test = fold == 0                                            # one fifth of the blocks are held out entirely
rf = RandomForestClassifier(n_estimators=300, max_features="sqrt", min_samples_leaf=1, n_jobs=-1,
                            random_state=SEED, oob_score=True)
t0 = time.time()
rf.fit(X[~test], y[~test])
pred = rf.predict(X[test])
section("spatial hold-out test (20% of blocks)")
print(f"trained on {(~test).sum()} samples in {time.time() - t0:.1f} s; tested on {test.sum()} samples")
print(f"overall accuracy {accuracy_score(y[test], pred):.3f} | macro F1 {f1_score(y[test], pred, average='macro'):.3f} | "
      f"out-of-bag accuracy {rf.oob_score_:.3f}")
names = [CLASSES[c] for c in USE]
print(classification_report(y[test], pred, labels=USE, target_names=names, digits=3))
recall = pd.Series(np.diag(confusion_matrix(y[test], pred, labels=USE, normalize="true")), index=USE)
w = share.reindex(USE) / share.reindex(USE).sum()
print(f"area-weighted accuracy (class recall weighted by WorldCover area share): {(w * recall).sum():.3f}")
cm = confusion_matrix(y[test], pred, labels=USE)

section("random vs spatially blocked 5-fold cross-validation")
cv = {}
for name, splitter, groups in [("random", KFold(5, shuffle=True, random_state=SEED), None),
                               ("spatial blocks", GroupKFold(5), fold)]:
    p = cross_val_predict(RandomForestClassifier(n_estimators=300, n_jobs=-1, random_state=SEED), X, y,
                          cv=splitter, groups=groups)
    cv[name] = p
    per_class = f1_score(y, p, labels=USE, average=None)
    print(f"{name:<15} accuracy {accuracy_score(y, p):.3f} | macro F1 {f1_score(y, p, average='macro'):.3f} | "
          + " ".join(f"{CLASSES[c].split()[0]} {v:.2f}" for c, v in zip(USE, per_class)))
gap = accuracy_score(y, cv["random"]) - accuracy_score(y, cv["spatial blocks"])
print(f"random CV overstates accuracy by {gap * 100:.1f} percentage points")

# --- Step 7: feature importance ---------------------------------------------------------------------
section("feature importance")
pi = permutation_importance(rf, X[test], y[test], n_repeats=5, random_state=SEED, n_jobs=-1, scoring="f1_macro")
imp = pd.DataFrame({"permutation": pi.importances_mean, "perm_sd": pi.importances_std,
                    "impurity": rf.feature_importances_}, index=FEATURES).sort_values("permutation", ascending=False)
print(imp.head(12).round(4).to_string())
by_season = imp["impurity"].groupby(imp.index.str.split("_").str[0]).sum()
print("impurity importance by season:", by_season.round(3).to_dict())

# --- Step 8: predict the full map ----------------------------------------------------------------------
final = RandomForestClassifier(n_estimators=300, max_features="sqrt", n_jobs=-1, random_state=SEED).fit(X, y)
flat = X_all.reshape(-1, F)
pred_map = np.zeros(H * W, dtype="uint8")
proba_max = np.zeros(H * W, dtype="float32")
t0 = time.time()
for i in range(0, H * W, 200_000):                          # predict in chunks to keep memory low
    chunk = flat[i:i + 200_000]
    proba = final.predict_proba(chunk)                      # scikit-learn >= 1.4 forests accept missing values
    pred_map[i:i + 200_000] = final.classes_[proba.argmax(1)]
    proba_max[i:i + 200_000] = proba.max(1)
pred_map, proba_max = pred_map.reshape(H, W), proba_max.reshape(H, W)
section("full map prediction")
print(f"predicted {H * W:,} pixels in {time.time() - t0:.0f} s")
in_use = np.isin(wc, USE)
agree = (pred_map == wc)[in_use].mean()
print(f"agreement with WorldCover over the classes used: {agree:.1%}; "
      f"median top-class probability {np.median(proba_max):.2f}")
f1_map = f1_score(wc[in_use], pred_map[in_use], labels=USE, average=None)
print("per-class F1 against WorldCover over every pixel: "
      + " ".join(f"{CLASSES[c].split()[0]} {v:.2f}" for c, v in zip(USE, f1_map)))
ha = RES * RES / 1e4
areas = pd.DataFrame({"WorldCover_ha": pd.Series(wc.ravel()).value_counts().reindex(USE).fillna(0) * ha,
                      "random_forest_ha": pd.Series(pred_map.ravel()).value_counts().reindex(USE).fillna(0) * ha})
areas.index = [CLASSES[c] for c in USE]
print(areas.round(0).astype(int).to_string())
with rasterio.open(OUT / "land_cover_rf_2021.tif", "w", driver="GTiff", height=H, width=W, count=1, dtype="uint8",
                   crs=CRS, transform=gbox.affine, compress="deflate", nodata=0) as dst:
    dst.write(pred_map, 1)
    dst.write_colormap(1, {c: tuple(int(WC_COLOURS[c][i:i + 2], 16) for i in (1, 3, 5)) + (255,) for c in USE})
pd.DataFrame(np.column_stack([rows, cols, y, fold]), columns=["row", "col", "label", "fold"]).to_csv(
    OUT / "training_samples.csv", index=False)

# --- Figures ------------------------------------------------------------------------------------------------
bb = gbox.boundingbox
ext = [bb.left, bb.right, bb.bottom, bb.top]
lc_cmap = ListedColormap([WC_COLOURS[c] for c in CLASSES])
lc_norm = BoundaryNorm([c - 0.5 for c in CLASSES] + [100.5], lc_cmap.N)
legend = [Patch(color=WC_COLOURS[c], label=CLASSES[c]) for c in USE]
section("figures")


def frame(ax, title):
    ax.set_title(title, loc="left")
    xt, yt = np.arange(0, ext[1] - ext[0], 10_000), np.arange(0, ext[3] - ext[2], 5_000)
    ax.set_xticks(ext[0] + xt, [f"{v / 1000:.0f}" for v in xt])
    ax.set_yticks(ext[2] + yt, [f"{v / 1000:.0f}" for v in yt])
    ax.set_xlim(ext[0], ext[1]); ax.set_ylim(ext[2], ext[3])
    ax.set_xlabel("km"); ax.set_ylabel("km")


rgb = np.dstack([comp[f"summer_{b}"].values for b in ("B04", "B03", "B02")])
rgb = np.clip(rgb / np.nanpercentile(rgb, 99), 0, 1) ** 0.75
rgb = np.nan_to_num(rgb, nan=0.85)

# 1. summer composite and WorldCover
fig, axes = plt.subplots(2, 1, figsize=(10, 10.2))
axes[0].imshow(rgb, extent=ext)
frame(axes[0], f"Sentinel-2 summer {YEAR} median composite (June–August, cloud-masked)")
axes[1].imshow(np.ma.masked_equal(wc, 0), cmap=lc_cmap, norm=lc_norm, extent=ext, interpolation="nearest")
frame(axes[1], "ESA WorldCover 2021, resampled to 20 m: the reference labels")
axes[1].legend(handles=legend, loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False)
fig.tight_layout()
fig.savefig(OUT / "01-composite-and-worldcover.png")
plt.close(fig)

# 2. sample points and the spatial blocks
fig, ax = plt.subplots(figsize=(10, 4.6))
ax.imshow(rgb.mean(axis=2), cmap="gray", extent=ext, vmin=0, vmax=1.3)
xs, ys_ = gbox.affine * (cols + 0.5, rows + 0.5)
for c in USE:
    m = y == c
    ax.scatter(xs[m], ys_[m], s=2, color=WC_COLOURS[c], linewidths=0)
for b in blocks[[fold_of_block[b] == 0 for b in blocks]]:
    r0, c0 = divmod(b, W // block_px + 1)
    x0, y0 = gbox.affine * (c0 * block_px, r0 * block_px)
    ax.add_patch(plt.Rectangle((x0, y0 - BLOCK), BLOCK, BLOCK, fill=False, ec="black", lw=1.1))
frame(ax, f"{len(y):,} stratified random samples; black squares = held-out 2 km test blocks")
ax.legend(handles=legend + [Patch(fill=False, ec="black", label="test block")], loc="upper left",
          bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False)
fig.savefig(OUT / "02-samples-and-blocks.png")
plt.close(fig)

# 3. confusion matrix (spatial hold-out)
cmn = cm / cm.sum(1, keepdims=True)
fig, ax = plt.subplots(figsize=(7.2, 6.2))
ax.imshow(cmn, cmap="Blues", vmin=0, vmax=1)
for i in range(len(USE)):
    for j in range(len(USE)):
        if cm[i, j]:
            ax.text(j, i, f"{cmn[i, j]:.0%}\n({cm[i, j]})", ha="center", va="center", fontsize=7,
                    color="white" if cmn[i, j] > 0.55 else "black")
ax.set_xticks(range(len(USE)), names, rotation=35, ha="right")
ax.set_yticks(range(len(USE)), names)
ax.set_xlabel("predicted by the random forest")
ax.set_ylabel("WorldCover label")
ax.set_title("Confusion matrix on the spatial hold-out blocks (row %, counts)")
fig.savefig(OUT / "03-confusion-matrix.png")
plt.close(fig)

# 4. F1 per class: random CV, spatial CV and the whole map
f1r = f1_score(y, cv["random"], labels=USE, average=None)
f1s = f1_score(y, cv["spatial blocks"], labels=USE, average=None)
fig, ax = plt.subplots(figsize=(9, 4.2))
xx = np.arange(len(USE))
ax.bar(xx - 0.27, f1r, 0.27, color="#b0b7b5", label=f"random 5-fold CV (accuracy {accuracy_score(y, cv['random']):.1%})")
ax.bar(xx, f1s, 0.27, color="#1d6fa3",
       label=f"spatially blocked 5-fold CV (accuracy {accuracy_score(y, cv['spatial blocks']):.1%})")
ax.bar(xx + 0.27, f1_map, 0.27, color="#c21f3a", label=f"every pixel of the final map (agreement {agree:.1%})")
for x0, v in zip(xx + 0.27, f1_map):
    ax.text(x0, v + 0.01, f"{v:.2f}", ha="center", fontsize=7)
ax.set_xticks(xx, names, rotation=20, ha="right")
ax.set_ylabel("F1 score against WorldCover")
ax.set_ylim(0, 1.05)
ax.legend(loc="lower left", fontsize=8)
ax.grid(axis="y", alpha=0.3)
ax.set_title("Per-class F1: pure-pixel samples look far better than the whole map")
fig.savefig(OUT / "04-f1-three-estimates.png")
plt.close(fig)

# 5. permutation importance
top = imp.head(15).iloc[::-1]
season_col = {"spring": "#2f9e44", "summer": "#d4a017", "autumn": "#c21f3a"}
fig, ax = plt.subplots(figsize=(8, 5.4))
ax.barh(top.index, top["permutation"], xerr=top["perm_sd"], color=[season_col[i.split("_")[0]] for i in top.index])
ax.set_xlabel("drop in macro F1 when the feature is shuffled (spatial hold-out)")
ax.set_title(f"Permutation importance: top 15 of {F} features")
ax.legend(handles=[Patch(color=v, label=k) for k, v in season_col.items()], loc="lower right", fontsize=8)
ax.grid(axis="x", alpha=0.3)
fig.savefig(OUT / "05-feature-importance.png")
plt.close(fig)

# 6. predicted map and disagreement with WorldCover
fig, axes = plt.subplots(2, 1, figsize=(10, 10.2))
axes[0].imshow(pred_map, cmap=lc_cmap, norm=lc_norm, extent=ext, interpolation="nearest")
frame(axes[0], f"Random forest land cover, {YEAR} (agreement with WorldCover {agree:.0%})")
axes[0].legend(handles=legend, loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=8, frameon=False)
dis = np.where(~in_use, np.nan, (pred_map != wc).astype(float))
axes[1].imshow(rgb.mean(axis=2), cmap="gray", extent=ext, vmin=0, vmax=1.3)
axes[1].imshow(np.ma.masked_where(dis != 1, dis), cmap=ListedColormap(["#c21f3a"]), extent=ext,
               interpolation="nearest")
frame(axes[1], "Where the random forest and WorldCover disagree (red)")
fig.tight_layout()
fig.savefig(OUT / "06-predicted-map.png")
plt.close(fig)
print(f"wrote {len(list(OUT.glob('*.png')))} figures, land_cover_rf_2021.tif and training_samples.csv")