"""Location-allocation in Python with spopt: p-median, p-center and the maximal covering location
problem (MCLP) for siting urgent-care or mobile clinic bases in Cornwall, on OSMnx drive times.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/location-allocation-python/
Part 2 of "Healthcare accessibility analysis in Python". From Part 1 (spatial-accessibility-2sfca-python)
it reuses the cached OSMnx drive graph, data/cornwall_drive.graphml, and the Geofabrik extract. If
Part 1 has not been run, this script downloads the data and builds the same graph itself.

Data: ONS Census 2021 TS001 via Nomis and ONS Open Geography boundaries and centroids (OGL v3);
OpenStreetMap via Geofabrik and OSMnx (ODbL).
"""
import json
import sys
import time
import warnings
import zipfile
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import geopandas as gpd
import requests
import osmnx as ox
import osmium
import shapely
import pulp
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import dijkstra
from spopt.locate import PMedian, PCenter, MCLP, LSCP
import matplotlib
matplotlib.use("Agg")   # render to files only; no GUI backend
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.collections import LineCollection

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

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)
PART1 = Path("../spatial-accessibility-2sfca-python/data")   # Part 1's cache, used if present
CRS = "EPSG:27700"
LAD = "E06000052"            # Cornwall (not the Isles of Scilly)
BUFFER_M = 25_000            # same graph extent as Part 1
SNAP_KMH = 20
P = 6                        # number of new sites in the main comparison
RADIUS = 15                  # MCLP service standard, minutes (free-flow drive time)
P_RANGE = range(1, 13)
THIN_M = 2_000               # keep at most one candidate building per 2 km
UA = {"User-Agent": "TechnicalGIS-tutorial (https://www.technicalgis.com)"}
CBC = lambda: pulp.PULP_CBC_CMD(msg=False, timeLimit=600)

ONS = "https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services"
LAD_URL = f"{ONS}/Local_Authority_Districts_May_2026_Boundaries_UK_BGC/FeatureServer/0/query"
OA_URL = f"{ONS}/Output_Areas_2021_EW_BGC_V2/FeatureServer/0/query"
OA_PWC_URL = f"{ONS}/OA_December_2021_EW_PWC_V4/FeatureServer/0/query"
LSOA_PWC_URL = f"{ONS}/LSOA_PopCentroids_EW_2021_V4/FeatureServer/0/query"
TS001 = "https://www.nomisweb.co.uk/output/census/2021/census2021-ts001.zip"
GEOFABRIK = [f"https://download.geofabrik.de/europe/united-kingdom/england/{c}-latest.osm.pbf"
             for c in ("cornwall", "devon")]
NOT_DRIVE = {"abandoned", "bridleway", "bus_guideway", "construction", "corridor", "cycleway", "elevator",
             "escalator", "footway", "no", "path", "pedestrian", "planned", "platform", "proposed", "raceway",
             "razed", "rest_area", "service", "services", "steps", "track"}
NOT_SERVICE = {"alley", "driveway", "emergency_access", "parking", "parking_aisle", "private"}
CANDIDATE_AMENITY = {"community_centre", "townhall", "village_hall"}


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


def cached(name):
    """A file from Part 1's data folder if it exists there, otherwise this tutorial's own data folder."""
    return PART1 / name if (PART1 / name).exists() else DATA / name


def download(url, path):
    if not path.exists():
        r = requests.get(url, timeout=600, headers=UA)
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


def arcgis(url, path, **params):
    """Query an ArcGIS FeatureServer layer page by page and cache the result as GeoJSON."""
    if not path.exists():
        feats, offset = [], 0
        while True:
            q = {"outFields": "*", "outSR": 27700, "f": "geojson", "resultOffset": offset,
                 "resultRecordCount": 2000, "where": "1=1", **params}
            r = requests.post(url, data=q, timeout=300, headers=UA)
            r.raise_for_status()
            page = r.json()["features"]
            feats += page
            if len(page) < 2000:
                break
            offset += len(page)
        path.write_text(json.dumps({"type": "FeatureCollection", "features": feats}))
    return gpd.read_file(path).set_crs(CRS, allow_override=True)


def envelope(x0, y0, x1, y1):
    return {"geometry": f"{x0},{y0},{x1},{y1}", "geometryType": "esriGeometryEnvelope", "inSR": 27700,
            "spatialRel": "esriSpatialRelIntersects"}


def drivable(tags):
    """The same rules as OSMnx's network_type="drive" filter, applied to a way's tags."""
    hw = tags.get("highway")
    return (hw is not None and hw not in NOT_DRIVE and tags.get("area") != "yes"
            and tags.get("motor_vehicle") != "no" and tags.get("motorcar") != "no"
            and tags.get("access") != "private" and tags.get("service") not in NOT_SERVICE)


def build_drive_graph(area, folder):
    """Part 1's graph builder: drive network for `area` from Geofabrik extracts, read by OSMnx."""
    poly = gpd.GeoSeries([area], crs=CRS).to_crs(4326).iloc[0]
    osm_file = folder / "drive_network.osm"
    keep_nodes, ways = set(), {}
    for url in GEOFABRIK:
        pbf = download(url, folder / url.rsplit("/", 1)[1])
        for w in osmium.FileProcessor(str(pbf), osmium.osm.NODE | osmium.osm.WAY).with_locations():
            if not w.is_way() or w.id in ways or not drivable(w.tags):
                continue
            xy = np.array([(n.lon, n.lat) for n in w.nodes])
            if shapely.contains_xy(poly, xy[:, 0], xy[:, 1]).any():
                ways[w.id] = ([n.ref for n in w.nodes], dict(w.tags))
                keep_nodes.update(n.ref for n in w.nodes)
    written = set()
    with osmium.SimpleWriter(str(osm_file), overwrite=True) as writer:
        for url in GEOFABRIK:
            for n in osmium.FileProcessor(str(folder / url.rsplit("/", 1)[1]), osmium.osm.NODE):
                if n.id in keep_nodes and n.id not in written:
                    writer.add_node(n)
                    written.add(n.id)
        for wid, (refs, tags) in sorted(ways.items()):
            writer.add_way(osmium.osm.mutable.Way(id=wid, version=1, nodes=refs, tags=tags))
    G = ox.graph_from_xml(osm_file, simplify=True)
    G = ox.routing.add_edge_speeds(G)
    G = ox.routing.add_edge_travel_times(G)
    G = ox.truncate.largest_component(G, strongly=True)
    return ox.project_graph(G, to_crs=CRS)


def read_zip_csv(path, name_part):
    with zipfile.ZipFile(path) as z:
        name = next(n for n in z.namelist() if name_part in n and n.endswith(".csv"))
        return pd.read_csv(z.open(name), low_memory=False, encoding="utf-8-sig")


def osm_places(pbf, out):
    """Community buildings and hospitals from an OSM extract: nodes, and ways at the mean of their nodes."""
    if not out.exists():
        rows = []
        for o in osmium.FileProcessor(str(pbf), osmium.osm.NODE | osmium.osm.WAY).with_locations():
            amenity = o.tags.get("amenity")
            kind = ("candidate" if amenity in CANDIDATE_AMENITY else
                    "hospital" if amenity == "hospital" and o.tags.get("healthcare") != "clinic" else None)
            if kind is None:
                continue
            if o.is_node():
                lon, lat = o.location.lon, o.location.lat
            else:
                xy = np.array([(n.lon, n.lat) for n in o.nodes])
                lon, lat = xy.mean(axis=0)
            rows.append({"osm": f"{'n' if o.is_node() else 'w'}{o.id}", "kind": kind, "amenity": amenity,
                         "name": o.tags.get("name", ""), "lon": lon, "lat": lat})
        df = pd.DataFrame(rows)
        gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.lon, df.lat), crs=4326).to_crs(CRS).to_file(out)
    return gpd.read_file(out)


def thin(gdf, min_dist):
    """Greedy thinning: keep a point only if it is at least min_dist from every point already kept."""
    gdf = gdf.assign(named=gdf.name.fillna("").ne("")).sort_values(["named", "osm"], ascending=[False, True])
    kept, xy = [], gdf.get_coordinates().to_numpy()
    for i in range(len(gdf)):
        if not kept or np.min(np.hypot(*(xy[kept] - xy[i]).T)) >= min_dist:
            kept.append(i)
    return gdf.iloc[kept].drop(columns="named").reset_index(drop=True)


section("versions")
print(f"python {sys.version.split()[0]} | spopt {version('spopt')} | pulp {version('pulp')} | "
      f"osmnx {version('osmnx')} | geopandas {version('geopandas')} | numpy {version('numpy')}")
print(f"CBC available through PuLP: {pulp.PULP_CBC_CMD().available()}")

# ---------------------------------------------------------------- Step 1: demand (Cornwall LSOAs)
section("demand")
lads = arcgis(LAD_URL, cached("lads.geojson"), **envelope(130000, 0, 300000, 140000))
cornwall = lads[lads.LAD26CD == LAD].geometry.iloc[0]
devon_lads = lads[lads.intersects(cornwall.buffer(BUFFER_M)) & (lads.LAD26CD != LAD)]
oas = arcgis(OA_URL, cached("cornwall_oas.geojson"), where="LSOA21NM LIKE 'Cornwall %'")
ts001 = download(TS001, cached("census2021-ts001.zip"))
col = {"geography code": "code", "Residence type: Total; measures: Value": "pop"}
lsoa = arcgis(LSOA_PWC_URL, DATA / "cornwall_lsoa_pwc.geojson", **envelope(*cornwall.bounds))
lsoa = lsoa[lsoa.LSOA21CD.isin(oas.LSOA21CD)]
lsoa = lsoa.merge(read_zip_csv(ts001, "-lsoa.csv").rename(columns=col)[["code", "pop"]],
                  left_on="LSOA21CD", right_on="code").drop(columns="code").reset_index(drop=True)
print(f"{len(lsoa)} Cornwall LSOAs, {lsoa['pop'].sum():,} residents (Census 2021), "
      f"median {lsoa['pop'].median():,.0f} per LSOA")
oa = arcgis(OA_PWC_URL, cached("oa_pwc.geojson"), **envelope(*cornwall.buffer(BUFFER_M).bounds))
oa = oa[oa.OA21CD.isin(oas.OA21CD)].merge(
    read_zip_csv(ts001, "-oa.csv").rename(columns=col)[["code", "pop"]], left_on="OA21CD", right_on="code")
oa = oa.reset_index(drop=True)
print(f"for the problem-size benchmark: {len(oa):,} Cornwall OAs, {oa['pop'].sum():,} residents")

# ---------------------------------------------------------------- Step 2: candidate sites
section("candidate sites (OpenStreetMap)")
pbf = cached("cornwall-latest.osm.pbf")
if not pbf.exists():
    download(GEOFABRIK[0], pbf)
places = osm_places(pbf, DATA / "osm_places.gpkg")
places = places[places.within(cornwall)]
print("OSM features in Cornwall:", places.groupby(["kind", "amenity"]).size().to_dict())
# amenity=hospital also catches a surgery, military medical centres and a WW2 site; keep names ending
# "Hospital" and drop the one private hospital, leaving Cornwall's NHS acute and community hospitals
hosp = places[(places.kind == "hospital") & places.name.str.endswith("Hospital")
              & (places.name != "The Duchy Hospital")]
hosp = thin(hosp, 500).sort_values("name").reset_index(drop=True)
print(f"NHS hospital sites kept: {', '.join(hosp.name)}")
cand = thin(places[places.kind == "candidate"], THIN_M)
print(f"{len(cand)} candidate community buildings after {THIN_M / 1000:.0f} km thinning "
      f"(from {(places.kind == 'candidate').sum()}); {len(hosp)} existing hospital sites")

# ---------------------------------------------------------------- Step 3: travel-time matrices
section("travel times")
t0 = time.time()
graph_file = cached("cornwall_drive.graphml")
if not graph_file.exists():
    print("Part 1's graph not found: building it from Geofabrik extracts (a few minutes)")
    ox.save_graphml(build_drive_graph(cornwall.buffer(BUFFER_M), DATA), graph_file)
G = ox.load_graphml(graph_file)
print(f"graph from {graph_file.as_posix()}: {G.number_of_nodes():,} nodes, {G.number_of_edges():,} edges; "
      f"{time.time() - t0:.0f} s")
nodes = list(G.nodes)
pos = {n: i for i, n in enumerate(nodes)}
edges = pd.DataFrame([(pos[a], pos[b], d["travel_time"]) for a, b, d in G.edges(data=True)], columns=["u", "v", "t"])
edges = edges.groupby(["u", "v"], as_index=False).t.min()
WT = csr_matrix((edges.t, (edges.u, edges.v)), shape=(len(nodes), len(nodes))).T.tocsr()


def snap(points):
    ids, dist = ox.distance.nearest_nodes(G, points.geometry.x, points.geometry.y, return_dist=True)
    return np.array([pos[i] for i in ids]), np.asarray(dist) / 1000 / SNAP_KMH * 60


def minutes(demand, sites):
    """Drive time in minutes from every demand point (rows) to every site (columns)."""
    d_node, d_hop = snap(demand)
    s_node, s_hop = snap(sites)
    to_site = dijkstra(WT, directed=True, indices=s_node) / 60
    return to_site[:, d_node].T + d_hop[:, None] + s_hop[None, :]


sites = pd.concat([cand.assign(kind="candidate"), hosp.assign(kind="hospital")], ignore_index=True)
t0 = time.time()
T_all = minutes(lsoa, sites)
T_oa_all = minutes(oa, sites)
print(f"LSOA x site matrix {T_all.shape} and OA x site matrix {T_oa_all.shape} in {time.time() - t0:.1f} s; "
      f"unreachable pairs: {np.isinf(T_all).sum() + np.isinf(T_oa_all).sum()}")
is_c = (sites.kind == "candidate").to_numpy()
T, T_oa, T_h = T_all[:, is_c], T_oa_all[:, is_c], T_all[:, ~is_c]
w, w_oa = lsoa["pop"].to_numpy(float), oa["pop"].to_numpy(float)


def evaluate(open_cols, M=T, wt=w):
    """Population-weighted mean, worst and coverage, each person using the nearest open site."""
    t = M[:, open_cols].min(axis=1)
    return {"mean_min": np.average(t, weights=wt), "max_min": t.max(),
            "covered_pct": wt[t <= RADIUS].sum() / wt.sum() * 100}


def solve(kind, p, M=T, wt=w, predefined=None):
    """Build and solve one spopt model with CBC; return the open columns and the solve time."""
    kw = {} if predefined is None else {"predefined_facilities_arr": predefined}
    if kind == "p-median":
        m = PMedian.from_cost_matrix(M, wt, p_facilities=p, **kw)
    elif kind == "p-center":
        m = PCenter.from_cost_matrix(M, p_facilities=p, **kw)
    else:
        m = MCLP.from_cost_matrix(M, wt, service_radius=RADIUS, p_facilities=p, **kw)
    t0 = time.perf_counter()
    m = m.solve(CBC())
    secs = time.perf_counter() - t0
    status = pulp.LpStatus[m.problem.status]
    opened = [j for j, clients in enumerate(m.fac2cli) if len(clients)] if kind != "mclp" else \
        [j for j in range(M.shape[1]) if m.fac_vars[j].value() > 0.5]
    return sorted(opened), secs, status, m.problem.numVariables(), m.problem.numConstraints()


def pcenter_by_covering(M, p):
    """Exact p-center by bisection on the radius: the smallest r whose set-covering problem (LSCP)
    needs no more than p sites. Each LSCP is a small, fast model; the direct p-center model is not."""
    r = np.unique(M[M >= M.min(axis=1).max()])          # every row must reach some site, so r >= this
    lo, hi, best, n_solves = 0, len(r) - 1, None, 0
    t0 = time.perf_counter()
    while lo <= hi:
        mid = (lo + hi) // 2
        m = LSCP.from_cost_matrix(M, service_radius=r[mid]).solve(CBC())
        n_solves += 1
        opened = [j for j in range(M.shape[1]) if m.fac_vars[j].value() > 0.5]
        if len(opened) <= p:
            best, hi = opened, mid - 1
        else:
            lo = mid + 1
    return sorted(best), time.perf_counter() - t0, n_solves


existing = evaluate(np.arange(T_h.shape[1]), M=T_h)
print(f"existing {len(hosp)} hospital sites: mean {existing['mean_min']:.1f} min | worst {existing['max_min']:.1f} min | "
      f"{existing['covered_pct']:.1f}% within {RADIUS} min")

# ---------------------------------------------------------------- Step 4: the three models at p = P
section(f"p-median, p-center and MCLP with p = {P} (greenfield)")
MODELS = ["p-median", "p-center", "mclp"]
sol = {}
for k in MODELS:
    opened, secs, status, nv, nc = solve(k, P)
    sol[k] = opened
    e = evaluate(opened)
    print(f"{k:8s} {status} in {secs:5.1f} s ({nv:,} variables, {nc:,} constraints) | mean {e['mean_min']:.1f} min | "
          f"worst {e['max_min']:.1f} min | {e['covered_pct']:.1f}% within {RADIUS} min")
    print("         sites: " + "; ".join(cand.name[j] or cand.osm[j] for j in opened))
opened, secs, n = pcenter_by_covering(T, P)
print(f"p-center by LSCP bisection: worst {evaluate(opened)['max_min']:.2f} min "
      f"(direct model {evaluate(sol['p-center'])['max_min']:.2f} min) from {n} set-covering solves in {secs:.1f} s")
shared = set(sol["p-median"]) & set(sol["p-center"]) & set(sol["mclp"])
print(f"sites chosen by all three models: {len(shared)}")

# ---------------------------------------------------------------- Step 5: trade-off as p varies
section(f"trade-off curves, p = {P_RANGE.start}-{P_RANGE.stop - 1}")
rows = []
for p in P_RANGE:
    for k in MODELS:
        if k == "p-center":
            opened, secs, _ = pcenter_by_covering(T, p)
            status = "Optimal"
        else:
            opened, secs, status, *_ = solve(k, p)
        rows.append({"model": k, "p": p, "secs": secs, "status": status, **evaluate(opened)})
curve = pd.DataFrame(rows)
curve.to_csv(OUT / "tradeoff.csv", index=False)
print(curve.pivot(index="p", columns="model", values=["mean_min", "max_min", "covered_pct"]).round(1).to_string())
st = curve.groupby("model").secs.agg(["median", "max"]).round(2)
print("solve time (s) by model (p-center = LSCP bisection):", st.to_dict(orient="index"))
gain = curve[curve.model == "p-median"].set_index("p").mean_min
print("p-median: minutes saved per extra site: " + ", ".join(f"{p}->{p + 1}: {gain[p] - gain[p + 1]:.2f}"
                                                            for p in list(P_RANGE)[:-1]))

# ---------------------------------------------------------------- Step 6: problem size
section("problem size: LSOAs vs output areas")
bench = []
for k in MODELS:
    for label, M, wt in (("LSOA", T, w), ("OA", T_oa, w_oa)):
        if k == "p-center":   # the direct model took minutes on LSOAs; use the covering bisection here
            opened, secs, n = pcenter_by_covering(M, P)
            status, nv, nc, note = "Optimal", np.nan, np.nan, f"{n} LSCPs"
        else:
            opened, secs, status, nv, nc = solve(k, P, M=M, wt=wt)
            note = f"{nv / 1000:,.0f}k vars" if nv >= 1000 else f"{nv:.0f} vars"
        e = evaluate(opened, M=T_oa, wt=w_oa)     # judge every solution on the finer OA demand
        bench.append({"model": k, "demand": label, "rows": M.shape[0], "variables": nv, "constraints": nc, "note": note,
                      "secs": secs, "status": status, **e})
bench = pd.DataFrame(bench)
print(bench.round(2).to_string(index=False))

# ---------------------------------------------------------------- Step 7: adding to the existing network
section(f"adding 3 sites to the {len(hosp)} existing hospital sites")
T_ext = np.hstack([T_h, T])
fixed = np.arange(T_h.shape[1])
ext = {}
for k in ("p-median", "mclp"):
    opened, secs, status, *_ = solve(k, len(hosp) + 3, M=T_ext, predefined=fixed)
    new = [j - len(hosp) for j in opened if j >= len(hosp)]
    ext[k] = new
    e = evaluate(opened, M=T_ext)
    print(f"{k:8s} {status} in {secs:.1f} s | mean {e['mean_min']:.1f} min (was {existing['mean_min']:.1f}) | "
          f"worst {e['max_min']:.1f} min (was {existing['max_min']:.1f}) | {e['covered_pct']:.1f}% within {RADIUS} min "
          f"(was {existing['covered_pct']:.1f}%) | new: " + "; ".join(cand.name[j] or cand.osm[j] for j in new))

# ---------------------------------------------------------------- Figures
section("figures")
XLIM, YLIM = (cornwall.bounds[0] - 3000, cornwall.bounds[2] + 3000), (cornwall.bounds[1] - 3000, cornwall.bounds[3] + 3000)
COL = {"p-median": "#1543b0", "p-center": "#d95f02", "mclp": "#1b9e77"}
LABEL = {"p-median": "p-median", "p-center": "p-center", "mclp": f"MCLP ({RADIUS} min)"}


def frame(ax, title):
    devon_lads.boundary.plot(ax=ax, color="#9aa0a0", linewidth=0.5)
    gpd.GeoSeries([cornwall]).plot(ax=ax, color="#f4f4f0", edgecolor="#333333", linewidth=0.6, zorder=0)
    ax.set_xlim(XLIM); ax.set_ylim(YLIM); ax.set_axis_off(); ax.set_title(title, fontsize=10)
    sx, sy = XLIM[0] + 3000, YLIM[0] + 3000
    ax.plot([sx, sx + 20000], [sy, sy], color="black", lw=2)
    ax.text(sx + 10000, sy + 2500, "20 km", ha="center", fontsize=8)


# Figure 1: demand and candidates
fig, ax = plt.subplots(figsize=(8, 6))
frame(ax, f"Demand ({len(lsoa)} LSOA population-weighted centroids) and candidate sites")
lsoa.plot(ax=ax, markersize=lsoa["pop"] / 60, color="#6baed6", alpha=0.6, edgecolor="none")
cand.plot(ax=ax, marker="^", markersize=10, color="#444444", zorder=3)
hosp.plot(ax=ax, marker="P", markersize=45, color="#c51b7d", edgecolor="white", linewidth=0.5, zorder=4)
ax.legend(handles=[
    Line2D([], [], marker="o", ls="", color="#6baed6", markersize=7, label="LSOA residents (size = population)"),
    Line2D([], [], marker="^", ls="", color="#444444", markersize=5, label=f"Candidate community buildings ({len(cand)})"),
    Line2D([], [], marker="P", ls="", color="#c51b7d", markersize=8, label=f"Existing hospital sites in OSM ({len(hosp)})")],
    loc="lower right", fontsize=8, frameon=True)
fig.savefig(OUT / "01-demand-and-candidates.png"); plt.close(fig)

# Figure 2: the three solutions at p = P with allocation lines
lxy = lsoa.get_coordinates().to_numpy()
cxy = cand.get_coordinates().to_numpy()
fig, axes = plt.subplots(1, 3, figsize=(15, 4.8))
for ax, k in zip(axes, MODELS):
    opened = np.array(sol[k])
    near = opened[T[:, opened].argmin(axis=1)]
    t = T[np.arange(len(lsoa)), near]
    segs = np.stack([lxy, cxy[near]], axis=1)
    ax.add_collection(LineCollection(segs, colors=np.where(t <= RADIUS, "#9ecae1", "#fb6a4a"), linewidths=0.5))
    e = evaluate(opened)
    frame(ax, f"{LABEL[k]}: mean {e['mean_min']:.1f} min, worst {e['max_min']:.1f} min,\n"
              f"{e['covered_pct']:.0f}% of residents within {RADIUS} min")
    ax.scatter(lxy[:, 0], lxy[:, 1], s=lsoa["pop"] / 150, color="#555555", zorder=3)
    ax.scatter(cxy[opened, 0], cxy[opened, 1], s=110, marker="*", color=COL[k], edgecolor="black", lw=0.6, zorder=5)
fig.legend(handles=[Line2D([], [], color="#9ecae1", label=f"LSOA within {RADIUS} min of its nearest site"),
                    Line2D([], [], color="#fb6a4a", label=f"LSOA more than {RADIUS} min away"),
                    Line2D([], [], marker="*", ls="", color="#555555", markersize=10, label="chosen site")],
           loc="lower center", ncol=3, fontsize=8, frameon=False, bbox_to_anchor=(0.5, 0.07))
fig.suptitle(f"Where to put {P} sites among {len(cand)} candidates: three objectives, three answers", y=1.02)
fig.savefig(OUT / "02-solutions-p6.png"); plt.close(fig)

# Figure 3: metrics of each solution, plus the existing hospital network
metrics = pd.DataFrame({LABEL[k].replace(" (", "\n("): evaluate(sol[k]) for k in MODELS} |
                       {f"Existing {len(hosp)}\nhospital sites": existing}).T
fig, axes = plt.subplots(1, 3, figsize=(12, 3.6))
colours = [COL[k] for k in MODELS] + ["#c51b7d"]
for ax, (c, lab) in zip(axes, [("mean_min", "Population-weighted mean time (min)"), ("max_min", "Worst-served LSOA (min)"),
                               ("covered_pct", f"Residents within {RADIUS} min (%)")]):
    bars = ax.bar(metrics.index, metrics[c], color=colours)
    ax.bar_label(bars, fmt="%.1f", fontsize=8)
    ax.set_title(lab, fontsize=9); ax.tick_params(axis="x", labelsize=7); ax.grid(axis="y", alpha=0.3)
    ax.set_ylim(0, metrics[c].max() * 1.15)
fig.suptitle(f"Each model wins on its own objective (p = {P} new sites)", y=1.03)
fig.savefig(OUT / "03-solution-metrics.png"); plt.close(fig)

# Figure 4: trade-off curves as p varies
fig, axes = plt.subplots(1, 3, figsize=(13, 3.8))
for ax, (c, lab) in zip(axes, [("mean_min", "Population-weighted mean time (min)"), ("max_min", "Worst-served LSOA (min)"),
                               ("covered_pct", f"Residents within {RADIUS} min (%)")]):
    for k in MODELS:
        g = curve[curve.model == k]
        ax.plot(g.p, g[c], marker="o", ms=3.5, color=COL[k], label=LABEL[k])
    ax.axhline(existing[c], color="#c51b7d", ls=":", lw=1.2, label=f"existing {len(hosp)} hospital sites")
    ax.set_xlabel("Number of sites, p"); ax.set_title(lab, fontsize=9); ax.grid(alpha=0.3)
    ax.set_xticks(list(P_RANGE))
axes[0].legend(fontsize=8)
fig.suptitle("Trade-off curves: diminishing returns from each extra site", y=1.03)
fig.savefig(OUT / "04-tradeoff-curves.png"); plt.close(fig)

# Figure 5: solve time and problem size
fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), gridspec_kw={"width_ratios": [1.4, 1]})
for k in MODELS:
    g = curve[curve.model == k]
    axes[0].plot(g.p, g.secs, marker="o", ms=3.5, color=COL[k], label=LABEL[k])
axes[0].set_yscale("log"); axes[0].set_xlabel("Number of sites, p"); axes[0].set_ylabel("CBC solve time (s, log scale)")
axes[0].set_title(f"Solve time, {len(lsoa)} LSOAs x {len(cand)} candidates", fontsize=9)
axes[0].set_xticks(list(P_RANGE)); axes[0].grid(alpha=0.3, which="both"); axes[0].legend(fontsize=8)
xb = np.arange(len(MODELS))
for i, (label, hatch) in enumerate((("LSOA", ""), ("OA", "//"))):
    g = bench[bench.demand == label].set_index("model").loc[MODELS]
    bars = axes[1].bar(xb + (i - 0.5) * 0.38, g.secs, 0.38, color=[COL[k] for k in MODELS], hatch=hatch,
                       edgecolor="white", label=f"{label}s ({g.rows.iloc[0]:,} demand points)")
    axes[1].bar_label(bars, labels=g.note.tolist(), fontsize=6.5, rotation=90,
                      padding=2)
axes[1].set_yscale("log"); axes[1].set_xticks(xb, [LABEL[k] for k in MODELS])
axes[1].set_ylabel("CBC solve time (s, log scale)"); axes[1].set_ylim(top=bench.secs.max() * 30)
axes[1].set_title(f"Problem size at p = {P}: LSOAs vs output areas", fontsize=9)
axes[1].legend(fontsize=7, loc="upper left"); axes[1].grid(axis="y", alpha=0.3, which="both")
fig.savefig(OUT / "05-solve-time.png"); plt.close(fig)

# Figure 6: adding 3 sites to the existing network
hxy = hosp.get_coordinates().to_numpy()
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6))
for ax, k in zip(axes, ("p-median", "mclp")):
    opened = list(range(len(hosp))) + [j + len(hosp) for j in ext[k]]
    t = T_ext[:, opened].min(axis=1)
    e = evaluate(opened, M=T_ext)
    frame(ax, f"{LABEL[k]}: 3 new sites added to {len(hosp)} hospitals\nmean {e['mean_min']:.1f} min "
              f"(was {existing['mean_min']:.1f})\n{e['covered_pct']:.0f}% within {RADIUS} min "
              f"(was {existing['covered_pct']:.0f}%)")
    sc = ax.scatter(lxy[:, 0], lxy[:, 1], c=t, s=lsoa["pop"] / 80, cmap="YlOrRd", vmin=0, vmax=40, zorder=3)
    ax.scatter(hxy[:, 0], hxy[:, 1], marker="P", s=45, color="#c51b7d", edgecolor="white", lw=0.5, zorder=4)
    ax.scatter(cxy[ext[k], 0], cxy[ext[k], 1], marker="*", s=180, color=COL[k], edgecolor="black", lw=0.6, zorder=5)
cb = fig.colorbar(sc, ax=axes, shrink=0.7, pad=0.01, extend="max")
cb.set_label("Drive time to nearest site (min)")
axes[0].legend(handles=[Line2D([], [], marker="P", ls="", color="#c51b7d", markersize=8, label="existing hospital site"),
                        Line2D([], [], marker="*", ls="", color="#555555", markersize=11, label="new site")],
               loc="lower right", fontsize=7)
fig.savefig(OUT / "06-existing-plus-new.png"); plt.close(fig)
print("saved " + ", ".join(sorted(p.name for p in OUT.glob("*.png"))))