"""
H3 hexagon grids in Python: index, fill, aggregate, smooth and compare with squares.

Companion script for https://technicalgis.com/tutorials/h3-hexagon-grid-python/

Data:
  - Overture Maps Foundation places theme (latest release, read anonymously from the public
    AWS bucket with DuckDB). Overture places are CDLA-Permissive-2.0.
  - Greater London boundary: ONS Regions (December 2023) Boundaries EN BGC, Open Government
    Licence v3.0. Source: Office for National Statistics licensed under the OGL v3.0.
    Contains OS data (c) Crown copyright and database right 2023.

Run:  python h3_hexagon_grid.py      (under a minute; the first run saves about 53 MB of places)
"""
from __future__ import annotations

import statistics
import sys
import time
from pathlib import Path

import duckdb
import geopandas as gpd
import h3
import matplotlib
import numpy as np
import pandas as pd
import pyproj
import requests
import shapely
from shapely.geometry import Polygon, shape

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.colors import LogNorm  # noqa: E402

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"      # downloads (not committed)
OUT = HERE / "output"     # figures and console log
DATA.mkdir(exist_ok=True)
OUT.mkdir(exist_ok=True)

BBOX = (-0.52, 51.28, 0.34, 51.70)       # lon/lat box around Greater London
TRAFALGAR = (51.50809, -0.12804)         # lat, lng: Nelson's Column, Trafalgar Square
CENTRAL = (523_000, 175_000, 538_000, 186_000)  # BNG box for the central London zoom (15 x 11 km)
BNG = "EPSG:27700"
RES = 8                                  # main analysis resolution: about 0.74 km2 per hexagon
CAFES = ("cafe", "coffee_shop")          # Overture basic_category values used for the smoothing example
STACK_LIMIT = 50                         # drop coordinates shared by this many places or more
ONS = ("https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services/"
       "Regions_December_2023_Boundaries_EN_BGC/FeatureServer/0/query")

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
TO_BNG = pyproj.Transformer.from_crs("EPSG:4326", BNG, always_xy=True)


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


def timed(fn, repeat: int = 3):
    """Run fn() `repeat` times; return the last result and the median wall-clock time in seconds."""
    times = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        result = fn()
        times.append(time.perf_counter() - t0)
    return result, statistics.median(times)


def cell_polygon(cell: str) -> Polygon:
    """H3 returns (lat, lng) pairs; shapely wants (x=lng, y=lat)."""
    return Polygon([(lng, lat) for lat, lng in h3.cell_to_boundary(cell)])


def cells_to_gdf(cells, **columns) -> gpd.GeoDataFrame:
    cells = list(cells)
    gdf = gpd.GeoDataFrame({"h3": cells, **columns}, geometry=[cell_polygon(c) for c in cells], crs="EPSG:4326")
    return gdf.to_crs(BNG)


def zoom(ax, box=CENTRAL) -> None:
    ax.set_xlim(box[0], box[2])
    ax.set_ylim(box[1], box[3])
    ax.set_axis_off()


section("versions")
print(f"python {sys.version.split()[0]} | h3 {h3.__version__} (C library {h3.versions()['c']}) | "
      f"duckdb {duckdb.__version__} | geopandas {gpd.__version__} | shapely {shapely.__version__}")

# ---------------------------------------------------------------------------- 1. data
section("1. Overture places for Greater London")
places_file = DATA / "london_places.parquet"
con = duckdb.connect()
con.sql("INSTALL spatial; LOAD spatial; INSTALL httpfs; LOAD httpfs; INSTALL h3 FROM community; LOAD h3;")
con.sql("SET s3_region = 'us-west-2'; SET geometry_always_xy = true")
if not places_file.exists():
    release = requests.get("https://stac.overturemaps.org/catalog.json", timeout=60).json()["latest"]
    xmin, ymin, xmax, ymax = BBOX
    files = [r[0] for r in con.execute(f"""
        SELECT assets.aws.alternate.s3.href FROM read_parquet('https://stac.overturemaps.org/{release}/collections.parquet')
        WHERE collection = 'place' AND bbox.xmin <= {xmax} AND bbox.xmax >= {xmin}
          AND bbox.ymin <= {ymax} AND bbox.ymax >= {ymin}""").fetchall()]
    t0 = time.perf_counter()
    con.execute(f"""
        COPY (
          SELECT id, names.primary AS name, basic_category AS category, confidence, operating_status,
                 ST_X(geometry) AS lng, ST_Y(geometry) AS lat, '{release}' AS release
          FROM read_parquet({files})
          WHERE bbox.xmin <= {xmax} AND bbox.xmax >= {xmin} AND bbox.ymin <= {ymax} AND bbox.ymax >= {ymin}
        ) TO '{places_file.as_posix()}' (FORMAT parquet)""")
    print(f"downloaded from {len(files)} Overture file(s) in {time.perf_counter() - t0:.1f} s")
raw = pd.read_parquet(places_file)
print(f"Overture release {raw['release'].iloc[0]}: {len(raw):,} places in the bounding box "
      f"({places_file.stat().st_size / 1e6:.1f} MB on disk)")

boundary_file = DATA / "london_boundary.geojson"
if not boundary_file.exists():
    r = requests.get(ONS, params={"where": "RGN23NM='London'", "outFields": "RGN23CD,RGN23NM",
                                  "outSR": 4326, "f": "geojson"}, timeout=120)
    r.raise_for_status()
    boundary_file.write_bytes(r.content)
london = gpd.read_file(boundary_file)
london_bng = london.to_crs(BNG)
london_geom = london.geometry.iloc[0]
london_area_km2 = london_bng.area.iloc[0] / 1e6
print(f"ONS boundary {london['RGN23CD'].iloc[0]} {london['RGN23NM'].iloc[0]}: {london_area_km2:,.1f} km2")

x, y = TO_BNG.transform(raw["lng"].to_numpy(), raw["lat"].to_numpy())
inside = shapely.contains_xy(london_bng.geometry.iloc[0], x, y)
closed = raw["operating_status"].isin(["permanently_closed", "temporarily_closed"])
keep = inside & ~closed & (raw["confidence"] >= 0.5)
places = raw.loc[keep].reset_index(drop=True)
places["x"], places["y"] = x[keep], y[keep]
print(f"inside London: {inside.sum():,} | closed: {(inside & closed).sum():,} | "
      f"confidence < 0.5: {(inside & ~closed & (raw['confidence'] < 0.5)).sum():,} | left: {len(places):,}")
stack = places.groupby(["lat", "lng"])["id"].transform("size")
biggest = places.loc[stack.idxmax()]
print(f"largest stack: {stack.max():,} places at ({biggest['lat']:.5f}, {biggest['lng']:.5f}), e.g. "
      f"{', '.join(places.loc[(places['lat'] == biggest['lat']) & (places['lng'] == biggest['lng']), 'name'].head(3))}")
farm = stack >= STACK_LIMIT
print(f"dropping {farm.sum():,} places at {len(places.loc[farm, ['lat', 'lng']].drop_duplicates()):,} coordinates "
      f"shared by {STACK_LIMIT}+ places (mostly registered-office addresses) | kept: {(~farm).sum():,}")
places = places.loc[~farm].reset_index(drop=True)
print("top categories:", ", ".join(f"{k} {v:,}" for k, v in places["category"].value_counts().head(6).items()))
print(f"cafes ({' + '.join(CAFES)}): {places['category'].isin(CAFES).sum():,}")

# ---------------------------------------------------------------------------- 2. H3 basics
section("2. one point, one cell: the h3-py v4 API")
lat, lng = TRAFALGAR
cell = h3.latlng_to_cell(lat, lng, 9)
print(f"latlng_to_cell({lat}, {lng}, 9) = {cell}")
print(f"as a 64-bit integer: {h3.str_to_int(cell)} | resolution {h3.get_resolution(cell)} | "
      f"base cell {h3.get_base_cell_number(cell)} | pentagon? {h3.is_pentagon(cell)}")
clat, clng = h3.cell_to_latlng(cell)
offset = h3.great_circle_distance((lat, lng), (clat, clng), unit="m")
print(f"cell centre ({clat:.5f}, {clng:.5f}), {offset:.0f} m from the input point; "
      f"boundary has {len(h3.cell_to_boundary(cell))} vertices")
print(f"cell_to_parent(cell, 7) = {h3.cell_to_parent(cell, 7)} | children at res 10: "
      f"{len(h3.cell_to_children(cell, 10))} | grid_disk k=1: {len(h3.grid_disk(cell, 1))} cells, "
      f"k=2: {len(h3.grid_disk(cell, 2))}")
print(f"area of this cell {h3.cell_area(cell, unit='m^2'):,.0f} m2 vs global average res-9 hexagon "
      f"{h3.average_hexagon_area(9, unit='m^2'):,.0f} m2")

section("H3 resolutions: average size and the cell at Trafalgar Square")
rows = []
for res in range(16):
    c = h3.latlng_to_cell(lat, lng, res)
    rows.append({"res": res,
                 "avg_area_km2": h3.average_hexagon_area(res, unit="km^2"),
                 "avg_edge_km": h3.average_hexagon_edge_length(res, unit="km"),
                 "cells_on_earth": h3.get_num_cells(res),
                 "london_cell_km2": h3.cell_area(c, unit="km^2")})
res_table = pd.DataFrame(rows)
res_table["london/avg"] = res_table["london_cell_km2"] / res_table["avg_area_km2"]
fmt = {"avg_area_km2": "{:,.6g}".format, "avg_edge_km": "{:,.4g}".format, "cells_on_earth": "{:,}".format,
       "london_cell_km2": "{:,.6g}".format, "london/avg": "{:.3f}".format}
print(res_table.to_string(index=False, formatters=fmt))
print(f"pentagons per resolution: {len(h3.get_pentagons(8))}; the grid is oriented so that they all fall in the sea")

# ---------------------------------------------------------------------------- 3. polygon fill
section("3. fill the London polygon with cells")
shape_ll = h3.geo_to_h3shape(london_geom)
print(f"geo_to_h3shape -> {type(shape_ll).__name__} with {len(shape_ll.polys) if hasattr(shape_ll, 'polys') else 1} part(s)")
fills = {}
for res in (7, 8, 9):
    for mode in ("center", "full", "overlap"):
        cells, t = timed(lambda: h3.h3shape_to_cells_experimental(shape_ll, res, contain=mode), repeat=1)
        unique = set(cells)  # a cell overlapping two parts of a MultiPolygon is returned twice
        fills[(res, mode)] = unique
        area = sum(h3.cell_area(c, unit="km^2") for c in unique)
        print(f"res {res} contain={mode:<7}: {len(unique):>6,} cells ({len(cells) - len(unique):>2} duplicates), "
              f"{area:>7,.1f} km2 ({100 * area / london_area_km2:5.1f}% of the polygon area) in {t:.2f} s")
assert fills[(8, "center")] == set(h3.geo_to_cells(london_geom, 8))
print("geo_to_cells(geometry, 8) gives the same cells as contain='center'")

fill9 = list(fills[(9, "center")])
compact = h3.compact_cells(fill9)
print(f"compact_cells: {len(fill9):,} res-9 cells -> {len(compact):,} mixed-resolution cells "
      f"({', '.join(f'res {r}: {n}' for r, n in sorted(pd.Series([h3.get_resolution(c) for c in compact]).value_counts().items()))})")
assert set(h3.uncompact_cells(compact, 9)) == set(fill9)
fill8 = sorted(fills[(8, "center")])
areas8 = np.array([h3.cell_area(c, unit="km^2") for c in fill8])
print(f"res-8 cell areas inside London: {areas8.min():.4f} to {areas8.max():.4f} km2 "
      f"(range {100 * (areas8.max() / areas8.min() - 1):.2f}%), mean {areas8.mean():.4f} km2")

# ---------------------------------------------------------------------------- 4. aggregate
section(f"4. index {len(places):,} points: h3-py loop vs DuckDB h3 extension")
lats, lngs = places["lat"].to_numpy(), places["lng"].to_numpy()
py_cells, t_py = timed(lambda: [h3.latlng_to_cell(a, b, RES) for a, b in zip(lats, lngs)])
print(f"h3-py list comprehension: {len(py_cells):,} cells in {t_py:.2f} s "
      f"({1e9 * t_py / len(py_cells):.0f} ns per point)")
con.register("places_df", places[["id", "lat", "lng", "category"]])
db_cells, t_db = timed(lambda: con.execute(
    f"SELECT h3_h3_to_string(h3_latlng_to_cell(lat, lng, {RES})) AS h3 FROM places_df").df()["h3"].tolist())
print(f"DuckDB h3_latlng_to_cell:  {len(db_cells):,} cells in {t_db:.2f} s "
      f"({1e9 * t_db / len(db_cells):.0f} ns per point)")
print(f"identical cell IDs from both: {py_cells == db_cells}")
places["h3_8"] = py_cells
places["h3_9"] = [h3.latlng_to_cell(a, b, 9) for a, b in zip(lats, lngs)]
places["h3_7"] = [h3.latlng_to_cell(a, b, 7) for a, b in zip(lats, lngs)]

counts = (places.groupby("h3_8").size().rename("places")
          .reindex(fill8, fill_value=0).rename_axis("h3").reset_index())
outside = ~places["h3_8"].isin(set(fill8))
counts["area_km2"] = areas8
counts["per_km2"] = counts["places"] / counts["area_km2"]
print(f"res-8 cells: {len(counts):,} | empty: {(counts['places'] == 0).sum():,} "
      f"({100 * (counts['places'] == 0).mean():.1f}%) | points whose cell centre is outside the polygon: "
      f"{outside.sum():,}")
q = counts["places"].quantile([0.5, 0.9, 0.99]).to_dict()
print(f"places per cell: median {q[0.5]:.0f}, 90th pct {q[0.9]:.0f}, 99th pct {q[0.99]:.0f}, max {counts['places'].max():,}")
top = counts.nlargest(5, "places").copy()
top["centre"] = [f"{a:.4f}, {b:.4f}" for a, b in (h3.cell_to_latlng(c) for c in top["h3"])]
top["top_category"] = [places.loc[places["h3_8"] == c, "category"].value_counts().index[0] for c in top["h3"]]
print(top[["h3", "places", "per_km2", "centre", "top_category"]].to_string(index=False, float_format="{:,.0f}".format))
in_top = places[places["h3_8"] == top["h3"].iloc[0]]
stacks = in_top.groupby(["lat", "lng"]).size().sort_values(ascending=False)
print(f"busiest cell: {len(in_top):,} places at {len(stacks):,} distinct coordinates; the 5 largest stacks hold "
      f"{', '.join(f'{n:,}' for n in stacks.head(5))} places "
      f"({100 * stacks.head(5).sum() / len(in_top):.0f}% of the cell)")

sql_counts = con.execute(f"""
    SELECT h3_h3_to_string(h3_latlng_to_cell(lat, lng, {RES})) AS h3, count(*) AS places
    FROM places_df GROUP BY 1""").df().set_index("h3")["places"]
print(f"DuckDB GROUP BY gives the same counts: {sql_counts.reindex(fill8, fill_value=0).tolist() == counts['places'].tolist()}")

section("hierarchy: is the parent of a point's res-9 cell its res-8 cell?")
parent_of_9 = np.array([h3.cell_to_parent(c, 8) for c in places["h3_9"]])
mismatch = parent_of_9 != places["h3_8"].to_numpy()
print(f"cell_to_parent(res-9 cell, 8) != latlng_to_cell(point, 8) for {mismatch.sum():,} of {len(places):,} "
      f"points ({100 * mismatch.mean():.1f}%)")
rollup = pd.Series(parent_of_9).value_counts().reindex(fill8, fill_value=0)
diff = (rollup.to_numpy() - counts["places"].to_numpy())
print(f"rolling res-9 counts up to res 8 changes {np.count_nonzero(diff):,} of {len(fill8):,} cell counts "
      f"(largest change {np.abs(diff).max()} places)")

# ---------------------------------------------------------------------------- 5. k-ring smoothing
section("5. k-ring smoothing of cafe counts at res 9")
fill9 = sorted(fill9)
cafes = places[places["category"].isin(CAFES)]
raw9 = cafes.groupby("h3_9").size().reindex(fill9, fill_value=0)
lookup = cafes.groupby("h3_9").size().to_dict()  # include cafes in edge cells whose centre is outside London


def disk_mean(cell: str, k: int, decay: float = 1.0) -> float:
    """Weighted mean over grid_disk(cell, k); ring d gets weight decay**d, missing cells count as 0."""
    total = weight = 0.0
    for d in range(k + 1):
        w = decay ** d
        ring = h3.grid_ring(cell, d)
        total += w * sum(lookup.get(c, 0) for c in ring)
        weight += w * len(ring)
    return total / weight


smooth = pd.DataFrame({"raw": raw9})
smooth["k1"] = [disk_mean(c, 1) for c in fill9]
smooth["k2_decay"] = [disk_mean(c, 2, decay=0.5) for c in fill9]
for col in ("raw", "k1", "k2_decay"):
    s = smooth[col]
    peak = s.idxmax()
    plat, plng = h3.cell_to_latlng(peak)
    print(f"{col:<9}: cells > 0 {100 * (s > 0).mean():5.1f}% | max {s.max():5.2f} at ({plat:.4f}, {plng:.4f}) | "
          f"std {s.std():.3f} | lag-1 autocorrelation "
          f"{np.corrcoef(s, [np.mean([s.get(n, 0) for n in h3.grid_ring(c, 1)]) for c in fill9])[0, 1]:.2f}")

sql_k1 = con.execute("""
    WITH c AS (SELECT h3_latlng_to_cell(lat, lng, 9) AS cell, count(*) AS n
               FROM places_df WHERE category IN ('cafe', 'coffee_shop') GROUP BY 1),
         spread AS (SELECT unnest(h3_grid_disk(cell, 1)) AS nb, n FROM c)
    SELECT h3_h3_to_string(nb) AS h3, sum(n) / 7.0 AS k1 FROM spread GROUP BY nb""").df().set_index("h3")["k1"]
print(f"DuckDB h3_grid_disk scatter-sum: max |difference| from the Python k=1 mean = "
      f"{np.abs(sql_k1.reindex(fill9, fill_value=0) - smooth['k1']).max():.2e}")

# ---------------------------------------------------------------------------- 6. squares vs hexagons
section("6. hexagons vs squares of the same area")
side = np.sqrt(areas8.mean() * 1e6)
print(f"square side for equal area with the mean London res-8 hexagon: {side:.1f} m")
ring = h3.grid_ring(h3.latlng_to_cell(lat, lng, RES), 1)
centre8 = h3.cell_to_latlng(h3.latlng_to_cell(lat, lng, RES))
nbr = [h3.great_circle_distance(centre8, h3.cell_to_latlng(c), unit="m") for c in ring]
print(f"hexagon: 6 neighbours at {min(nbr):.0f}-{max(nbr):.0f} m from the centre")
print(f"square : 4 edge neighbours at {side:.0f} m and 4 corner neighbours at {side * np.sqrt(2):.0f} m")

hex_bng = cells_to_gdf([h3.latlng_to_cell(lat, lng, RES)])
hg = hex_bng.geometry.iloc[0]
sq = shapely.box(0, 0, side, side)
print(f"compactness 4*pi*A/P^2: hexagon {4 * np.pi * hg.area / hg.length ** 2:.3f}, square {4 * np.pi * sq.area / sq.length ** 2:.3f}")
print(f"farthest point from the centre: hexagon {max(shapely.distance(hg.centroid, shapely.points(hg.exterior.coords))):.0f} m, "
      f"square {side / np.sqrt(2):.0f} m")

hcx, hcy = TO_BNG.transform(*np.array([h3.cell_to_latlng(c)[::-1] for c in places["h3_8"]]).T)
d_hex = np.hypot(places["x"] - hcx, places["y"] - hcy)


def square_counts(offset: float):
    ix = np.floor((places["x"] - offset) / side).astype(int)
    iy = np.floor((places["y"] - offset) / side).astype(int)
    key = ix.astype(str) + "_" + iy.astype(str)
    cx, cy = (ix + 0.5) * side + offset, (iy + 0.5) * side + offset
    return key, np.hypot(places["x"] - cx, places["y"] - cy)


sq_key, d_sq = square_counts(0.0)
sq_key_shift, _ = square_counts(side / 2)
print(f"mean distance from a place to its cell centre: hexagon {d_hex.mean():.0f} m, square {d_sq.mean():.0f} m")

per_point = {}
for name, key in (("square", sq_key), ("square shifted by half a cell", sq_key_shift), ("H3 res 8", places["h3_8"])):
    n = key.map(key.value_counts())
    per_point[name] = n
    print(f"{name:<30}: {key.nunique():>5,} occupied cells, max {key.value_counts().max():>5,} places in a cell")
change = np.abs(per_point["square shifted by half a cell"] / per_point["square"] - 1)
print(f"moving the square grid origin by {side / 2:.0f} m changes the count of the cell a place falls in by a median "
      f"{100 * change.median():.0f}% (90th pct {100 * change.quantile(0.9):.0f}%)")

# ---------------------------------------------------------------------------- figures
section("figures")
outline = london_bng.boundary

# Figure 1: nesting and k-rings around Trafalgar Square
fig, axes = plt.subplots(1, 2, figsize=(10, 5.4), layout="constrained")
c7 = h3.latlng_to_cell(lat, lng, 7)
ax = axes[0]
cells_to_gdf(h3.cell_to_children(c7, 9)).boundary.plot(ax=ax, color="#9ecae1", linewidth=0.6)
cells_to_gdf(h3.cell_to_children(c7, 8)).boundary.plot(ax=ax, color="#3182bd", linewidth=1.2)
cells_to_gdf([c7]).boundary.plot(ax=ax, color="#08306b", linewidth=2.4)
tx, ty = TO_BNG.transform(lng, lat)
ax.plot(tx, ty, "o", color="#d7301f", ms=5)
ax.annotate("Trafalgar Square", (tx, ty), xytext=(6, 6), textcoords="offset points", color="#d7301f")
ax.set_title("One res-7 cell (dark), its 7 res-8 children\n"
             "and 49 res-9 grandchildren. Children do not\ntile the parent exactly.", loc="left")
ax.set_aspect("equal")
ax.set_axis_off()
ax = axes[1]
c8 = h3.latlng_to_cell(lat, lng, 8)
colours = ["#d7301f", "#fc8d59", "#fdcc8a", "#fef0d9"]
for k in range(3, -1, -1):
    g = cells_to_gdf(h3.grid_ring(c8, k))
    g.plot(ax=ax, color=colours[k], edgecolor="#555555", linewidth=0.6)
for k in range(4):
    cx_, cy_ = TO_BNG.transform(*h3.cell_to_latlng(h3.grid_ring(c8, k)[0])[::-1])
    ax.annotate(f"k={k}", (cx_, cy_), ha="center", va="center", fontsize=8)
ax.set_title("grid_ring(cell, k) at res 8: 1, 6, 12\nand 18 cells. grid_disk(cell, 3) is all\n37 together.", loc="left")
ax.set_aspect("equal")
ax.set_axis_off()
fig.savefig(OUT / "01-h3-hierarchy-kring.png")
plt.close(fig)

# Figure 2: containment modes
fig, axes = plt.subplots(1, 3, figsize=(12, 4.2), layout="constrained")
for ax, mode in zip(axes, ("center", "full", "overlap")):
    cells_to_gdf(fills[(7, mode)]).plot(ax=ax, color="#6baed6", edgecolor="white", linewidth=0.4)
    outline.plot(ax=ax, color="black", linewidth=0.8)
    ax.set_title(f"contain='{mode}': {len(fills[(7, mode)])} res-7 cells", loc="left")
    ax.set_axis_off()
fig.suptitle("Filling the Greater London boundary with h3shape_to_cells_experimental", x=0.01, ha="left", fontsize=11)
fig.savefig(OUT / "02-polygon-fill-modes.png")
plt.close(fig)

# Figure 3: place density at res 8
g8 = cells_to_gdf(counts["h3"], places=counts["places"].to_numpy(), per_km2=counts["per_km2"].to_numpy())
fig, ax = plt.subplots(figsize=(9, 6.6), layout="constrained")
g8[g8["places"] == 0].plot(ax=ax, color="#eeeeee", edgecolor="white", linewidth=0.2)
g8[g8["places"] > 0].plot(ax=ax, column="per_km2", cmap="magma_r", norm=LogNorm(vmin=5, vmax=g8["per_km2"].max()),
                          edgecolor="white", linewidth=0.15, legend=True,
                          legend_kwds={"label": "Overture places per km² (log scale)", "shrink": 0.6})
outline.plot(ax=ax, color="black", linewidth=0.6)
ax.set_title(f"{len(places):,} Overture places in {len(g8):,} H3 resolution-8 hexagons (≈{areas8.mean():.2f} km² each)",
             loc="left")
ax.set_axis_off()
fig.savefig(OUT / "03-london-places-res8.png")
plt.close(fig)

# Figure 4: k-ring smoothing
g9 = cells_to_gdf(fill9, **{c: smooth[c].to_numpy() for c in smooth})
g9 = g9[g9.intersects(shapely.box(*CENTRAL))]
fig, axes = plt.subplots(1, 3, figsize=(13, 3.9), layout="constrained")
titles = {"raw": "Raw cafe count per res-9 cell", "k1": "Mean over grid_disk(k=1)",
          "k2_decay": "Weighted mean over grid_disk(k=2),\nweights 1, 0.5, 0.25"}
for ax, col in zip(axes, ("raw", "k1", "k2_decay")):
    vmax = g9[col].quantile(0.995)
    g9.plot(ax=ax, column=col, cmap="YlOrBr", vmin=0, vmax=vmax, edgecolor="none", legend=True,
            legend_kwds={"shrink": 0.7, "label": "cafes per cell"})
    outline.plot(ax=ax, color="black", linewidth=0.5)
    ax.set_title(titles[col], loc="left")
    zoom(ax)
fig.suptitle("Central London cafes (Overture basic_category cafe + coffee_shop)", x=0.01, ha="left", fontsize=11)
fig.savefig(OUT / "04-kring-smoothing.png")
plt.close(fig)

# Figure 5: hexagons vs squares, same area, same data
fig, axes = plt.subplots(1, 3, figsize=(13, 4.0), layout="constrained")
norm = LogNorm(vmin=1, vmax=10_000)
vc8 = places["h3_8"].value_counts()
g8c = cells_to_gdf(vc8.index, places=vc8.to_numpy())  # every occupied cell, including those on the Thames
g8c = g8c[g8c.intersects(shapely.box(*CENTRAL))]
for ax, (title, key, off) in zip(axes, (("H3 res 8 hexagons", None, None), ("Squares, origin (0, 0)", sq_key, 0.0),
                                        (f"Squares, origin shifted {side / 2:.0f} m", sq_key_shift, side / 2))):
    if key is None:
        g8c.plot(ax=ax, column="places", cmap="viridis", norm=norm, edgecolor="white", linewidth=0.2)
    else:
        vc = key.value_counts()
        ij = vc.index.str.split("_", expand=True)
        ix = np.array([int(a) for a, _ in ij]); iy = np.array([int(b) for _, b in ij])
        geoms = [shapely.box(a * side + off, b * side + off, (a + 1) * side + off, (b + 1) * side + off) for a, b in zip(ix, iy)]
        gsq = gpd.GeoDataFrame({"places": vc.to_numpy()}, geometry=geoms, crs=BNG)
        gsq = gsq[gsq.intersects(shapely.box(*CENTRAL))]
        gsq.plot(ax=ax, column="places", cmap="viridis", norm=norm, edgecolor="white", linewidth=0.2)
    outline.plot(ax=ax, color="black", linewidth=0.5)
    ax.set_title(title, loc="left")
    zoom(ax)
sm = plt.cm.ScalarMappable(norm=norm, cmap="viridis")
fig.colorbar(sm, ax=axes, shrink=0.7, label="places per cell (log scale)")
fig.suptitle(f"Same {len(places):,} places, cells of equal area ({areas8.mean():.2f} km²)", x=0.01, ha="left", fontsize=11)
fig.savefig(OUT / "05-hex-vs-square.png")
plt.close(fig)

# Figure 6: choice of resolution
fig, axes = plt.subplots(1, 3, figsize=(13, 3.9), layout="constrained")
norm = LogNorm(vmin=10, vmax=20000)
for ax, (res, col) in zip(axes, ((7, "h3_7"), (8, "h3_8"), (9, "h3_9"))):
    vc = places[col].value_counts()
    cells = [c for c in vc.index if c in fills[(res, "center")]]
    g = cells_to_gdf(cells, places=vc[cells].to_numpy())
    g["per_km2"] = g["places"] / np.array([h3.cell_area(c, unit="km^2") for c in cells])
    g.plot(ax=ax, column="per_km2", cmap="magma_r", norm=norm, edgecolor="none")
    outline.plot(ax=ax, color="black", linewidth=0.5)
    fill_area = np.mean([h3.cell_area(c, unit="km^2") for c in fills[(res, "center")]])
    ax.set_title(f"res {res}: {len(fills[(res, 'center')]):,} cells of {fill_area:.3g} km²", loc="left")
    ax.set_axis_off()
sm = plt.cm.ScalarMappable(norm=norm, cmap="magma_r")
fig.colorbar(sm, ax=axes, shrink=0.7, label="places per km² (log scale)")
fig.savefig(OUT / "06-resolution-choice.png")
plt.close(fig)
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))