"""Geodemographic classification in Python: a small Output Area Classification for Cardiff from
Census 2021 tables, with k-means, cluster profiles, and contiguous regionalisation (SKATER and max-p) in spopt.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/geodemographic-classification-python/
Data: ONS Census 2021 via Nomis bulk downloads (OGL v3), ONS Open Geography Output Area boundaries (OGL v3),
2021/2 Output Area Classification from the Geographic Data Service (OGL v3).
"""
import json
import random
import sys
import time
import warnings
import zipfile
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import libpysal
import matplotlib
import networkx as nx
import numpy as np
import pandas as pd
import requests

matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, TwoSlopeNorm
from matplotlib.patches import Patch
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components
from sklearn.cluster import KMeans
from sklearn.metrics import adjusted_rand_score, silhouette_samples, silhouette_score
from spopt.region import MaxPHeuristic, Skater

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
SEED = 42
CRS = "EPSG:27700"
LAD = "Cardiff"
NOMIS = "https://www.nomisweb.co.uk/output/census/2021/census2021-{}.zip"
ONS = "https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services"
OA_URL = f"{ONS}/Output_Areas_2021_EW_BGC_V2/FeatureServer/0/query"
GEODS = "https://data.geods.ac.uk/dataset/73a16168-f2b0-4732-b36a-3ec0253e34fe/resource"
OAC_URL = f"{GEODS}/4924ae88-d384-45f0-8b82-0a4a66e49e01/download/uk_oac_final.csv"
OAC_NAMES = f"{GEODS}/5ee35968-582e-4709-8433-90d9fd69d2a0/download/classification_codes_and_names_ukoac_lsoac_msoac.csv"
# Paul Tol's colour-blind-safe "muted" qualitative palette
TOL = ["#332288", "#88CCEE", "#44AA99", "#117733", "#999933", "#DDCC77", "#CC6677", "#882255", "#AA4499", "#DDDDDD"]


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


def download(url, path, **params):
    if not path.exists():
        r = requests.get(url, params=params or None, timeout=600, headers={"User-Agent": "TechnicalGIS-tutorial"})
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


section("versions")
print(f"python {sys.version.split()[0]} | scikit-learn {version('scikit-learn')} | spopt {version('spopt')} | "
      f"libpysal {version('libpysal')} | geopandas {version('geopandas')} | numpy {version('numpy')} | "
      f"pandas {version('pandas')}")

# ---------------------------------------------------------------- 1. Output Area boundaries for Cardiff
section("Output Areas")
oa_file = DATA / "cardiff_oa21.geojson"
if not oa_file.exists():
    feats, offset = [], 0
    while True:
        page = requests.get(OA_URL, params={"where": f"LSOA21NM LIKE '{LAD} %'", "outFields": "OA21CD,LSOA21CD,LSOA21NM",
                                            "outSR": 27700, "f": "geojson", "resultOffset": offset,
                                            "resultRecordCount": 1000}, timeout=300).json()
        feats += page["features"]; offset += len(page["features"])
        if len(page["features"]) < 1000:
            break
    oa_file.write_text(json.dumps({"type": "FeatureCollection", "features": feats}))
oa = gpd.read_file(oa_file).set_crs(CRS, allow_override=True).sort_values("OA21CD").reset_index(drop=True)
print(f"{len(oa):,} Output Areas in {oa.LSOA21CD.nunique()} LSOAs, {oa.area.sum() / 1e6:.1f} km2")
# map extent: the mainland only (Cardiff also includes the island of Flat Holm, 5 km offshore)
MAIN = max(oa.dissolve().explode().geometry, key=lambda p: p.area).bounds

# ---------------------------------------------------------------- 2. Census 2021 tables from Nomis
section("Census 2021 variables")


def table(code):
    z = zipfile.ZipFile(download(NOMIS.format(code), DATA / f"census2021-{code}.zip"))
    df = pd.read_csv(z.open(f"census2021-{code}-oa.csv")).set_index("geography code")
    df.columns = [c.split(": ", 1)[1].replace("; measures: Value", "") if ": " in c else c for c in df.columns]
    return df.loc[oa.OA21CD]


def pct(df, num, den):
    num = [num] if isinstance(num, str) else num
    return 100 * df[num].sum(axis=1) / df[den]


age, hh, ten = table("ts007a"), table("ts003"), table("ts054")
acc, car, sec = table("ts044"), table("ts045"), table("ts062")
qual, eth, hlth = table("ts067"), table("ts021"), table("ts037")
occ, welsh, cob = table("ts052"), table("ts032"), table("ts004")
eco, dens, pop = table("ts066"), table("ts006"), table("ts001")

V = pd.DataFrame(index=oa.OA21CD)
V["Aged 0-14"] = pct(age, ["Aged 4 years and under", "Aged 5 to 9 years", "Aged 10 to 14 years"], "Total")
V["Aged 20-24"] = pct(age, "Aged 20 to 24 years", "Total")
V["Aged 25-44"] = pct(age, [f"Aged {a} to {a + 4} years" for a in (25, 30, 35, 40)], "Total")
V["Aged 65+"] = pct(age, ["Aged 65 to 69 years", "Aged 70 to 74 years", "Aged 75 to 79 years",
                         "Aged 80 to 84 years", "Aged 85 years and over"], "Total")
V["Population density"] = dens.iloc[:, -1]
V["One-person households"] = pct(hh, "One person household", "Total")
V["Lone parents"] = pct(hh, "Single family household: Lone parent family: With dependent children", "Total")
V["Owns outright"] = pct(ten, "Owned: Owns outright", "Total: All households")
V["Owns with mortgage"] = pct(ten, "Owned: Owns with a mortgage or loan", "Total: All households")
V["Social rented"] = pct(ten, "Social rented", "Total: All households")
V["Private rented"] = pct(ten, "Private rented", "Total: All households")
V["Detached"] = pct(acc, "Detached", "Total: All households")
V["Terraced"] = pct(acc, "Terraced", "Total: All households")
V["Flats"] = pct(acc, [c for c in acc.columns if c.startswith(("In a purpose-built", "Part of", "In a commercial"))],
                 "Total: All households")
V["Overcrowded"] = pct(occ, ["Occupancy rating of bedrooms: -1", "Occupancy rating of bedrooms: -2 or less"],
                       "Total: All households")
V["No car"] = pct(car, "No cars or vans in household", "Total: All households")
sec_tot = "Total: All usual residents aged 16 years and over"
V["Higher professional"] = pct(sec, "L1, L2 and L3 Higher managerial, administrative and professional occupations",
                               sec_tot)
V["Routine occupations"] = pct(sec, ["L12 Semi-routine occupations", "L13 Routine occupations"], sec_tot)
V["Full-time students"] = pct(sec, "L15 Full-time students", sec_tot)
V["Degree level (L4+)"] = pct(qual, "Level 4 qualifications and above", sec_tot)
V["No qualifications"] = pct(qual, "No qualifications", sec_tot)
V["Unemployed"] = pct(eco, "Economically active (excluding full-time students): Unemployed", sec_tot)
V["Bad health"] = pct(hlth, ["Bad health", "Very bad health"], "Total: All usual residents")
V["Asian"] = pct(eth, "Asian, Asian British or Asian Welsh", "Total: All usual residents")
V["Black"] = pct(eth, "Black, Black British, Black Welsh, Caribbean or African", "Total: All usual residents")
V["Born outside UK"] = 100 - pct(cob, "Europe: United Kingdom", "Total")
V["Welsh speakers"] = pct(welsh, [c for c in welsh.columns if c.startswith("Can speak")],
                          "Total: All usual residents aged 3 years and over")
residents = pop["Total"].reindex(oa.OA21CD).to_numpy()
print(f"{V.shape[1]} candidate variables for {len(V):,} OAs; {V.isna().sum().sum()} missing values")
print(f"Cardiff: {residents.sum():,} usual residents; OA population median {int(np.median(residents))}, "
      f"range {residents.min()}-{residents.max()}")
print(V.describe().T[["mean", "50%", "min", "max"]].round(1).to_string())

# ---------------------------------------------------------------- 3. transform, standardise, check correlations
section("transformation, standardisation and correlation")
skew_raw = V.skew()
T = np.arcsinh(V)                 # inverse hyperbolic sine: log-like, but defined at zero
Z = (T - T.mean()) / T.std(ddof=0)
print(f"median |skewness|: raw {skew_raw.abs().median():.2f} -> IHS {T.skew().abs().median():.2f}; "
      f"most skewed raw: " + ", ".join(f"{k} {v:.1f}" for k, v in skew_raw.sort_values(ascending=False).head(3).items()))
C = Z.corr()
pairs = [(a, b, C.loc[a, b]) for i, a in enumerate(C) for b in C.columns[i + 1:] if abs(C.loc[a, b]) >= 0.8]
for a, b, r in sorted(pairs, key=lambda p: -abs(p[2])):
    print(f"|r| >= 0.8: {a} ~ {b}: r = {r:.2f}")
DROP = ["Degree level (L4+)", "Owns with mortgage"]
print(f"dropping {DROP}: degree level duplicates higher professional (r = {C.loc[DROP[0], 'Higher professional']:.2f}); "
      "the four tenure shares sum to ~100%, so one is redundant")
Zk = Z.drop(columns=DROP)
X = Zk.to_numpy()
print(f"{X.shape[1]} variables go into the clustering")

order = C.abs().mean().sort_values().index
fig, ax = plt.subplots(figsize=(8.6, 7.2))
im = ax.imshow(C.loc[order, order], cmap="RdBu_r", norm=TwoSlopeNorm(0, -1, 1))
ax.set_xticks(range(len(order)), order, rotation=90, fontsize=7); ax.set_yticks(range(len(order)), order, fontsize=7)
for a, b, r in pairs:
    for (p, q) in ((a, b), (b, a)):
        ax.add_patch(plt.Rectangle((list(order).index(q) - 0.5, list(order).index(p) - 0.5), 1, 1, fill=False,
                                   ec="black", lw=1.2))
fig.colorbar(im, ax=ax, shrink=0.75, label="Pearson r (IHS-transformed, standardised)")
ax.set_title(f"Correlation between {len(order)} Census 2021 variables across Cardiff's {len(V):,} Output Areas\n"
             "(boxes: |r| ≥ 0.8)")
fig.savefig(OUT / "01-correlation-matrix.png"); plt.close(fig)

# ---------------------------------------------------------------- 4. k-means: elbow, silhouette, stability
section("k-means: choosing k")
KS = range(2, 13)
inertia, sil = [], []
for k in KS:
    km = KMeans(k, n_init=50, random_state=SEED).fit(X)
    inertia.append(km.inertia_); sil.append(silhouette_score(X, km.labels_))
    print(f"k={k:>2}: within-cluster sum of squares {km.inertia_:8.0f} | silhouette {sil[-1]:.3f}")
K = 7
stab = {}
for k in (4, 5, 6, 7, 8, 9):
    runs = [KMeans(k, n_init=1, random_state=s).fit(X).labels_ for s in range(30)]
    best = KMeans(k, n_init=100, random_state=SEED).fit(X).labels_
    stab[k] = (np.median([adjusted_rand_score(best, r) for r in runs]),
               np.mean([adjusted_rand_score(runs[i], runs[j]) for i in range(30) for j in range(i + 1, 30)]))
    print(f"k={k}: 30 single-start runs vs the 100-start solution, median ARI {stab[k][0]:.3f}; "
          f"mean pairwise ARI {stab[k][1]:.3f}")
print(f"chosen k = {K}")

fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), layout="constrained")
axes[0].plot(list(KS), inertia, "o-", color="#332288"); axes[0].set_title("Elbow: within-cluster sum of squares")
axes[1].plot(list(KS), sil, "o-", color="#117733"); axes[1].set_title("Mean silhouette width")
ks = list(stab)
axes[2].plot(ks, [stab[k][0] for k in ks], "o-", color="#CC6677", label="single start vs best of 100")
axes[2].plot(ks, [stab[k][1] for k in ks], "s--", color="#882255", label="mean pairwise, 30 seeds")
axes[2].set_title("Stability across random seeds (ARI)"); axes[2].legend(fontsize=7); axes[2].set_ylim(0, 1.02)
for ax in axes:
    ax.axvline(K, color="#999999", ls=":", lw=1.2); ax.set_xlabel("number of clusters k"); ax.grid(alpha=0.3)
axes[0].set_ylabel("sum of squared distances (z units²)"); axes[1].set_ylabel("silhouette"); axes[2].set_ylabel("ARI")
fig.savefig(OUT / "02-choosing-k.png"); plt.close(fig)

# ---------------------------------------------------------------- 5. final k-means and profiles
section("final k-means classification")
km = KMeans(K, n_init=100, random_state=SEED).fit(X)
sizes = pd.Series(km.labels_).value_counts()
relabel = {old: new for new, old in enumerate(sizes.index)}          # cluster 0 = largest
oa["kmeans"] = pd.Series(km.labels_).map(relabel).to_numpy()
s = silhouette_samples(X, oa.kmeans)
print(f"k = {K}: mean silhouette {s.mean():.3f}; OAs with negative silhouette {np.mean(s < 0):.1%}")

# index scores: cluster rate / Cardiff rate x 100, from raw percentages weighted by residents
city = (V.mul(residents, axis=0).sum() / residents.sum())
city["Population density"] = residents.sum() / (oa.area.sum() / 1e6)
idx = pd.DataFrame({c: V[oa.kmeans.to_numpy() == c].mul(residents[oa.kmeans == c], axis=0).sum()
                    / residents[oa.kmeans == c].sum() for c in range(K)}).T
idx["Population density"] = [residents[oa.kmeans == c].sum() / (oa[oa.kmeans == c].area.sum() / 1e6)
                             for c in range(K)]
index = (100 * idx / city).round(0)
for c in range(K):
    hi = index.loc[c].sort_values(ascending=False)
    lo = index.loc[c].sort_values()
    print(f"cluster {c} ({(oa.kmeans == c).sum()} OAs, {residents[oa.kmeans == c].sum():,} people): "
          f"high {', '.join(f'{k} {int(v)}' for k, v in hi.head(4).items())} | "
          f"low {', '.join(f'{k} {int(v)}' for k, v in lo.head(3).items())}")

# names chosen after reading the index scores printed above; treat them as labels for discussion, not facts
NAMES = {0: "Social-rented family estates", 1: "Detached owner-occupier suburbs", 2: "Older settled homeowners",
         3: "Dense multicultural inner city", 4: "Professional terraces", 5: "City-centre renters in flats",
         6: "Student quarter"}
label = lambda c: f"{c + 1}. {NAMES[c]}"
for c in range(K):
    print(f"  {label(c)}")
oa["kmeans_name"] = oa.kmeans.map(label)

fig, ax = plt.subplots(figsize=(10, 6))
cmap = ListedColormap(TOL[:K])
oa.plot(column="kmeans", cmap=cmap, ax=ax, edgecolor="white", linewidth=0.1, vmin=-0.5, vmax=K - 0.5)
oa.dissolve().boundary.plot(ax=ax, color="#333333", lw=0.6)
ax.legend(handles=[Patch(fc=TOL[c], label=f"{label(c)} ({(oa.kmeans == c).sum()} OAs)") for c in range(K)],
          loc="center left", bbox_to_anchor=(1.0, 0.5), fontsize=7.5, title="k-means cluster", title_fontsize=8)
ax.set_axis_off(); ax.set_title(f"A {K}-cluster geodemographic classification of Cardiff's {len(oa):,} Output Areas "
                                "(Census 2021)")
mx0, my0, mx1, my1 = MAIN
ax.set_xlim(mx0 - 300, mx1 + 300); ax.set_ylim(my0 - 900, my1 + 300)
ax.plot([mx0 + 300, mx0 + 2300], [my0 - 500] * 2, color="black", lw=2)
ax.text(mx0 + 1300, my0 - 380, "2 km", ha="center", fontsize=7)
fig.savefig(OUT / "03-kmeans-map.png"); plt.close(fig)

vars_order = list(V.columns)
fig, ax = plt.subplots(figsize=(10, 4.6))
L = np.log2(index[vars_order].clip(lower=12.5, upper=800) / 100)
im = ax.imshow(L, cmap="PuOr_r", norm=TwoSlopeNorm(0, -3, 3), aspect="auto")
for i in range(K):
    for j, v in enumerate(vars_order):
        val = int(index.loc[i, v])
        ax.text(j, i, f"{val}", ha="center", va="center", fontsize=6.3,
                color="white" if abs(L.iloc[i, j]) > 1.8 else "black")
ax.set_xticks(range(len(vars_order)), vars_order, rotation=60, ha="right", fontsize=7)
ax.set_yticks(range(K), [label(c) for c in range(K)], fontsize=7.5)
cb = fig.colorbar(im, ax=ax, shrink=0.9, ticks=[-3, -2, -1, 0, 1, 2, 3])
cb.ax.set_yticklabels(["≤12.5", "25", "50", "100", "200", "400", "≥800"]); cb.set_label("index score (Cardiff = 100)")
ax.set_title("Cluster profiles: index scores against the Cardiff average (100 = same as the city; 200 = twice)")
fig.savefig(OUT / "04-index-scores.png"); plt.close(fig)

RADAR = ["Aged 0-14", "Aged 20-24", "Aged 65+", "Owns outright", "Social rented", "Private rented", "Flats",
         "No car", "Higher professional", "Routine occupations", "Full-time students", "Born outside UK",
         "Welsh speakers"]
means = Z.groupby(oa.kmeans.to_numpy())[RADAR].mean().clip(-2.5, 2.5)
ang = np.linspace(0, 2 * np.pi, len(RADAR), endpoint=False)
fig, axes = plt.subplots(2, 4, figsize=(13, 7), subplot_kw={"polar": True}, layout="constrained")
for c, ax in enumerate(axes.flat):
    if c >= K:
        ax.set_visible(False); continue
    vals = np.r_[means.loc[c].to_numpy(), means.loc[c].iloc[0]]
    ax.plot(np.r_[ang, ang[0]], vals, color=TOL[c] if c != 1 else "#1f78b4", lw=1.8)
    ax.fill(np.r_[ang, ang[0]], vals, color=TOL[c], alpha=0.35)
    ax.plot(np.linspace(0, 2 * np.pi, 100), np.zeros(100), color="#555555", lw=0.8, ls="--")
    ax.set_ylim(-2.5, 2.5); ax.set_yticks([-2, 0, 2], ["-2", "0", "+2"], fontsize=6)
    ax.set_rlabel_position(180 / len(RADAR))
    ax.set_xticks(ang, RADAR, fontsize=6.3)
    ax.set_title(label(c), fontsize=8.5, pad=14)
fig.suptitle("Cluster means in standard deviations from the Cardiff OA mean (dashed ring = 0, clipped at ±2.5)")
fig.savefig(OUT / "05-radar-profiles.png"); plt.close(fig)

# ---------------------------------------------------------------- 6. compare with the official 2021/2 OAC
section("comparison with the official 2021/2 OAC")
oac = pd.read_csv(download(OAC_URL, DATA / "uk_oac_final.csv"), dtype=str).set_index("GeographyCode")
nm = pd.read_csv(download(OAC_NAMES, DATA / "oac_names.csv"), dtype=str)
sg_names = nm[nm.Level == "Supergroup"].set_index("Classification Code")["Classification Name"]
oa["oac_sg"] = oac.Supergroup.reindex(oa.OA21CD).to_numpy()
oa["oac_group"] = oac.Group.reindex(oa.OA21CD).to_numpy()
print(f"matched {oa.oac_sg.notna().sum()} of {len(oa)} OAs to the official classification")
print("official supergroups in Cardiff: " + "; ".join(f"{k} {sg_names[k]} {v}" for k, v in
                                                     oa.oac_sg.value_counts().sort_index().items()))
print(f"ARI between our k-means and the official supergroups (8 classes, UK-wide): "
      f"{adjusted_rand_score(oa.oac_sg, oa.kmeans):.3f}")
print(f"ARI against the official groups ({oa.oac_group.nunique()} present): "
      f"{adjusted_rand_score(oa.oac_group, oa.kmeans):.3f}")
ct = pd.crosstab(oa.kmeans, oa.oac_sg)
print(ct.to_string())

fig, ax = plt.subplots(figsize=(9, 4.2))
share = ct.div(ct.sum(axis=1), axis=0)
im = ax.imshow(share, cmap="Blues", vmin=0, vmax=1, aspect="auto")
for i in range(share.shape[0]):
    for j in range(share.shape[1]):
        if ct.iloc[i, j]:
            ax.text(j, i, ct.iloc[i, j], ha="center", va="center", fontsize=7,
                    color="white" if share.iloc[i, j] > 0.55 else "black")
ax.set_xticks(range(ct.shape[1]), [f"{k}. {sg_names[k]}" for k in ct.columns], rotation=30, ha="right", fontsize=7)
ax.set_yticks(range(K), [label(c) for c in range(K)], fontsize=7.5)
fig.colorbar(im, ax=ax, label="share of the cluster's OAs")
ax.set_title("Our Cardiff clusters (rows) against the official 2021/2 OAC supergroups (columns); numbers are OAs")
fig.savefig(OUT / "07-official-oac-crosstab.png"); plt.close(fig)

# ---------------------------------------------------------------- 7. MAUP: the same recipe on LSOAs
section("MAUP: repeat at LSOA level")
# population-weighted means of the OA percentages approximate the LSOA rates (denominators differ slightly)
lsoa_key = oa.LSOA21CD.to_numpy()
VL = V.drop(columns="Population density").mul(residents, axis=0).groupby(lsoa_key).sum().div(
    pd.Series(residents).groupby(lsoa_key).sum().to_numpy(), axis=0)
lsoa_geom = oa.dissolve("LSOA21CD")
VL["Population density"] = pd.Series(residents).groupby(lsoa_key).sum().to_numpy() / (lsoa_geom.area.to_numpy() / 1e6)
TL = np.arcsinh(VL[V.columns].drop(columns=DROP))
ZL = ((TL - TL.mean()) / TL.std(ddof=0)).to_numpy()
lab_l = KMeans(K, n_init=100, random_state=SEED).fit(ZL).labels_
oa_from_lsoa = pd.Series(lab_l, index=VL.index).reindex(lsoa_key).to_numpy()
mixed = oa.groupby("LSOA21CD").kmeans.nunique()
print(f"{len(VL)} LSOAs, k = {K}: silhouette {silhouette_score(ZL, lab_l):.3f}; "
      f"ARI with the OA classification (each OA takes its LSOA's cluster) {adjusted_rand_score(oa.kmeans, oa_from_lsoa):.3f}")
print(f"LSOAs whose OAs fall in more than one OA-level cluster: {np.mean(mixed > 1):.0%}; "
      f"max-min range of 'Full-time students' within LSOAs, median {V.groupby(lsoa_key)['Full-time students'].agg(np.ptp).median():.1f} points")

# ---------------------------------------------------------------- 8. contiguous regionalisation with spopt
section("contiguity: spatial weights")
w = libpysal.weights.Queen.from_dataframe(oa, use_index=False, silence_warnings=True)
print(f"queen contiguity: mean {w.mean_neighbors:.1f} neighbours, islands {len(w.islands)}, "
      f"connected components {w.n_components}")
if w.n_components > 1:
    w = libpysal.weights.fuse_islands(w, oa) if hasattr(libpysal.weights, "fuse_islands") else w
    print(f"after joining islands: components {w.n_components}")
attrs = list(Zk.columns)
reg = oa[["geometry"]].copy()
reg[attrs] = X
reg["residents"] = residents

within_ss = lambda lab: sum(((X[lab == c] - X[lab == c].mean(0)) ** 2).sum() for c in np.unique(lab))
total_ss = ((X - X.mean(0)) ** 2).sum()

section("SKATER")
t0 = time.perf_counter()
sk = Skater(reg, w, attrs, n_clusters=K, floor=20, trace=False, islands="increase")
sk.solve()
oa["skater"] = sk.labels_
print(f"SKATER with {K} regions, floor 20 OAs: {time.perf_counter() - t0:.1f} s; region sizes (OAs) "
      f"{sorted(pd.Series(sk.labels_).value_counts().tolist(), reverse=True)}")

section("max-p")
THRESH = 20_000
t0 = time.perf_counter()
np.random.seed(SEED); random.seed(SEED)       # max-p's heuristic is randomised
mp = MaxPHeuristic(reg, w, attrs, "residents", THRESH, top_n=2, max_iterations_construction=50)
mp.solve()
oa["maxp"] = mp.labels_
mp_pop = pd.Series(residents).groupby(mp.labels_).sum()
print(f"max-p, every region >= {THRESH:,} residents: p = {mp.p} regions in {time.perf_counter() - t0:.1f} s; "
      f"region populations {mp_pop.min():,}-{mp_pop.max():,}")

section("comparison: how much homogeneity does contiguity cost?")
res = {}
for name, lab in [("k-means (no contiguity)", oa.kmeans.to_numpy()), ("SKATER", oa.skater.to_numpy()),
                  ("max-p", oa.maxp.to_numpy())]:
    lab = np.asarray(lab)
    r2 = 1 - within_ss(lab) / total_ss
    # contiguous pieces = connected components of the queen graph once links between groups are cut
    A = w.sparse.tocoo()
    keep = lab[A.row] == lab[A.col]
    parts = connected_components(coo_matrix((A.data[keep], (A.row[keep], A.col[keep])), shape=A.shape))[0]
    res[name] = {"groups": int(len(np.unique(lab))), "r2": round(float(r2), 3), "pieces": int(parts)}
    print(f"{name:24s}: {len(np.unique(lab)):>2} groups | share of variance explained {r2:.3f} | "
          f"contiguous pieces on the map {parts}")
print(f"ARI k-means vs SKATER {adjusted_rand_score(oa.kmeans, oa.skater):.3f}; "
      f"k-means vs max-p {adjusted_rand_score(oa.kmeans, oa.maxp):.3f}")
lab16 = KMeans(mp.p, n_init=20, random_state=SEED).fit(X).labels_
print(f"for reference, aspatial k-means with k = {mp.p} explains {1 - within_ss(lab16) / total_ss:.3f}")

PAL = ["#88CCEE", "#DDCC77", "#CC6677", "#44AA99", "#AA4499", "#999933", "#332288"]


def region_colours(lab):
    """Greedy graph colouring so that neighbouring regions never share a colour."""
    A = w.sparse.tocoo()
    G = nx.Graph(); G.add_nodes_from(np.unique(lab))
    G.add_edges_from({(lab[i], lab[j]) for i, j in zip(A.row, A.col) if lab[i] != lab[j]})
    col = nx.greedy_color(G, strategy="largest_first")
    return [PAL[col[l] % len(PAL)] for l in lab]


fig, axes = plt.subplots(1, 3, figsize=(14, 4.6), layout="constrained")
for ax, col, title in [(axes[0], "kmeans", f"k-means, k = {K}: {res['k-means (no contiguity)']['pieces']} separate "
                                           f"pieces, R² {res['k-means (no contiguity)']['r2']:.2f}"),
                       (axes[1], "skater", f"SKATER, {K} regions: R² {res['SKATER']['r2']:.2f}"),
                       (axes[2], "maxp", f"max-p, every region ≥ {THRESH:,} residents: p = {mp.p}, "
                                         f"R² {res['max-p']['r2']:.2f}")]:
    if col == "kmeans":
        oa.plot(color=[TOL[c] for c in oa.kmeans], ax=ax, edgecolor="white", linewidth=0.05)
    else:
        oa.plot(color=region_colours(oa[col].to_numpy()), ax=ax, edgecolor="white", linewidth=0.05)
        oa.dissolve(col).boundary.plot(ax=ax, color="black", lw=0.7)
    ax.set_xlim(MAIN[0] - 300, MAIN[2] + 300); ax.set_ylim(MAIN[1] - 300, MAIN[3] + 300)
    ax.set_axis_off(); ax.set_title(title, fontsize=9)
fig.legend(handles=[Patch(fc=TOL[c], label=label(c)) for c in range(K)], loc="outside lower left", ncol=4,
           fontsize=7, title="left map: k-means clusters", title_fontsize=7.5)
fig.suptitle(f"Classification vs regionalisation: the same {len(attrs)} standardised variables, with and without a "
             "contiguity constraint\n(black lines: region boundaries; colours in the middle and right maps only "
             "separate neighbouring regions)")
fig.savefig(OUT / "06-skater-maxp.png"); plt.close(fig)

oa.to_file(OUT / "cardiff_oa_classification.gpkg", layer="oa", driver="GPKG")
json.dump({"k": K, "silhouette": round(float(s.mean()), 3), "regionalisation": res,
           "maxp_p": int(mp.p)}, open(OUT / "metrics.json", "w"), indent=2)
print("\nwrote output/cardiff_oa_classification.gpkg, output/metrics.json and 7 figures")