"""Geocoding in Python: batch geocode every open school in Sheffield with geopy + Nominatim and postcodes.io,
reverse geocode a sample, and measure accuracy against the coordinates the Department for Education publishes.

Data: Get Information about Schools (GIAS) all-establishments extract, Department for Education, Open Government
Licence v3.0. Geocoders: OpenStreetMap Nominatim (results are ODbL) and postcodes.io (ONS Postcode Directory, OGL).
Nominatim's usage policy is respected: at most 1 request per second, an identifying user agent, and every answer
is cached on disk so a rerun sends no requests at all.
Run: python geocode_schools.py   (first run takes about 12 minutes because of the 1 request per second limit)
"""
import json
import time
from datetime import date, timedelta
from pathlib import Path

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import requests
from geopy.extra.rate_limiter import RateLimiter
from geopy.geocoders import Nominatim

HERE = Path(__file__).parent
DATA, OUT, CACHE = HERE / "data", HERE / "output", HERE / "cache"
for p in (DATA, OUT, CACHE):
    p.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

USER_AGENT = "TechnicalGIS-tutorial"
GIAS = "https://ea-edubase-api-prod.azurewebsites.net/edubase/downloads/public/edubasealldata{}.csv"
POSTCODES_IO = "https://api.postcodes.io"
METHODS = ["Nominatim free text", "Nominatim structured", "Nominatim by name", "postcodes.io postcode"]
COLOURS = dict(zip(METHODS, ["#d53e4f", "#f46d43", "#3288bd", "#66a61e"]))
LAD_URL = ("https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services/"
           "Local_Authority_Districts_May_2026_Boundaries_UK_BGC/FeatureServer/0/query")


# ---------------------------------------------------------------- 1. the addresses and the reference coordinates
def gias_extract():
    """GIAS publishes a fresh extract every day; take the newest one that exists."""
    path = DATA / "edubasealldata.csv"
    if not path.exists():
        for back in range(10):
            stamp = (date.today() - timedelta(days=back)).strftime("%Y%m%d")
            r = requests.get(GIAS.format(stamp), timeout=300)
            if r.ok and len(r.content) > 1_000_000:
                path.write_bytes(r.content)
                print(f"downloaded GIAS extract {stamp}")
                break
        else:
            raise RuntimeError("no GIAS extract found in the last 10 days")
    return path


cols = ["URN", "EstablishmentName", "TypeOfEstablishment (name)", "Street", "Locality", "Address3", "Town",
        "Postcode", "Easting", "Northing", "LA (name)", "EstablishmentStatus (name)"]
gias = pd.read_csv(gias_extract(), encoding="cp1252", usecols=cols, dtype=str)
schools = gias[(gias["LA (name)"] == "Sheffield") & (gias["EstablishmentStatus (name)"] == "Open")
               & (gias["TypeOfEstablishment (name)"] != "Secure units")].copy()
schools = schools.rename(columns={"EstablishmentName": "name", "TypeOfEstablishment (name)": "type"})
schools[["Easting", "Northing"]] = schools[["Easting", "Northing"]].astype(float)
schools = schools.dropna(subset=["Easting", "Postcode"]).fillna("").reset_index(drop=True)
print(f"open establishments in Sheffield in GIAS: {len(schools)} (secure units left out)")
print(schools[["URN", "name", "Street", "Locality", "Town", "Postcode", "Easting", "Northing"]].head(4).to_string(index=False))
ref = gpd.GeoDataFrame(schools, geometry=gpd.points_from_xy(schools.Easting, schools.Northing), crs=27700)


def full_address(r):
    return ", ".join(p for p in (r.Street, r.Locality, r.Address3, r.Town, r.Postcode) if p)


schools["address"] = schools.apply(full_address, axis=1)

# ---------------------------------------------------------------- 2. a polite, cached Nominatim client
cache_file = CACHE / "nominatim.json"
cache = json.loads(cache_file.read_text(encoding="utf-8")) if cache_file.exists() else {}
geolocator = Nominatim(user_agent=USER_AGENT, timeout=20)
geocode_live = RateLimiter(geolocator.geocode, min_delay_seconds=1, max_retries=2, error_wait_seconds=10)
reverse_live = RateLimiter(geolocator.reverse, min_delay_seconds=1, max_retries=2, error_wait_seconds=10)
calls = []


def cached(kind, query, **kw):
    key = f"{kind}|{json.dumps(query, sort_keys=True)}|{json.dumps(kw, sort_keys=True)}"
    if key not in cache:
        t0 = time.perf_counter()
        fn = geocode_live if kind == "geocode" else reverse_live
        loc = fn(query, **kw)
        calls.append(time.perf_counter() - t0)
        cache[key] = loc.raw if loc else None
        if len(calls) % 25 == 0:
            cache_file.write_text(json.dumps(cache), encoding="utf-8")
            print(f"  {len(calls)} live requests so far")
    return cache[key]


def nominatim(query):
    raw = cached("geocode", query, country_codes="gb", addressdetails=True)
    if raw is None:
        return {}
    return {"lat": float(raw["lat"]), "lon": float(raw["lon"]), "addresstype": raw.get("addresstype", ""),
            "osm_class": f"{raw.get('class')}={raw.get('type')}", "place_rank": raw.get("place_rank"),
            "display_name": raw.get("display_name", "")}


t_start = time.perf_counter()
rows = []
for r in schools.itertuples():
    variants = {
        "Nominatim free text": f"{r.address}, United Kingdom",
        "Nominatim structured": {"street": r.Street, "city": r.Town or "Sheffield", "postalcode": r.Postcode,
                                 "country": "United Kingdom"},
        "Nominatim by name": f"{r.name}, Sheffield",
    }
    for method, q in variants.items():
        rows.append({"URN": r.URN, "method": method, **nominatim(q)})
cache_file.write_text(json.dumps(cache), encoding="utf-8")
print(f"\nNominatim: {len(rows)} forward queries, {len(calls)} sent live, {len(rows) - len(calls)} from the cache, "
      f"{time.perf_counter() - t_start:.0f} s")
if calls:
    print(f"live request wall time: median {np.median(calls):.2f} s (the rate limiter keeps them at least 1 s apart)")

# ---------------------------------------------------------------- 3. postcodes.io: bulk postcode lookup
pc_cache_file = CACHE / "postcodes_io.json"
pc_cache = json.loads(pc_cache_file.read_text(encoding="utf-8")) if pc_cache_file.exists() else {}
todo = sorted({p for p in schools.Postcode if p not in pc_cache})
t0 = time.perf_counter()
for i in range(0, len(todo), 100):
    resp = requests.post(f"{POSTCODES_IO}/postcodes", json={"postcodes": todo[i:i + 100]},
                         headers={"User-Agent": USER_AGENT}, timeout=60)
    resp.raise_for_status()
    for item in resp.json()["result"]:
        pc_cache[item["query"]] = item["result"]
pc_cache_file.write_text(json.dumps(pc_cache), encoding="utf-8")
print(f"postcodes.io: {schools.Postcode.nunique()} unique postcodes, {len(todo)} looked up live in "
      f"{-(-len(todo) // 100)} bulk request(s), {time.perf_counter() - t0:.1f} s")
for r in schools.itertuples():
    res = pc_cache.get(r.Postcode)
    rows.append({"URN": r.URN, "method": "postcodes.io postcode",
                 **({"lat": res["latitude"], "lon": res["longitude"], "addresstype": "postcode",
                     "osm_class": "postcode centroid", "display_name": res["postcode"]} if res else {})})

# ---------------------------------------------------------------- 4. accuracy against the GIAS coordinates
res = pd.DataFrame(rows).merge(schools[["URN", "name", "Postcode", "Easting", "Northing", "address"]], on="URN")
found = res.lat.notna()
pts = gpd.GeoSeries(gpd.points_from_xy(res.lon.fillna(0), res.lat.fillna(0)), crs=4326).to_crs(27700)
res["x"], res["y"] = pts.x.where(found), pts.y.where(found)
res["error_m"] = np.hypot(res.x - res.Easting, res.y - res.Northing)
res.drop(columns=["x", "y"]).to_csv(OUT / "geocoded_schools.csv", index=False)


def summarise(g):
    e = g.error_m.dropna()
    return pd.Series({"found": f"{len(e)}/{len(g)}", "median_m": e.median(), "p90_m": e.quantile(0.9),
                      "within_50m": 100 * (e <= 50).mean(), "within_100m": 100 * (e <= 100).mean(),
                      "within_250m": 100 * (e <= 250).mean(), "over_1km": int((e > 1000).sum())})


table = res.groupby("method", sort=False).apply(summarise, include_groups=False)
pd.set_option("display.width", 160)
print("\ndistance from the GIAS coordinates (m; percentages of the schools that were found):")
print(table.round(1).to_string())

# cross-check two independent geocoders: accept the name match only where it agrees with the postcode centroid
ft = res[res.method == "Nominatim free text"].set_index("URN")
pc = res[res.method == "postcodes.io postcode"].set_index("URN")
nm = res[res.method == "Nominatim by name"].set_index("URN").reindex(pc.index)
to_bng = lambda d: gpd.GeoSeries(gpd.points_from_xy(d.lon.fillna(0), d.lat.fillna(0)), crs=4326, index=d.index).to_crs(27700)
nm["agree_m"] = to_bng(nm).distance(to_bng(pc)).where(nm.lat.notna())
AGREE = 250
accept = nm.agree_m <= AGREE
combo = pc.copy()
combo.loc[accept, ["lat", "lon", "error_m", "addresstype", "osm_class"]] = nm.loc[accept, ["lat", "lon", "error_m",
                                                                                         "addresstype", "osm_class"]]
combo["method"] = "Name if it agrees with postcode, else postcode"
print(f"\ncross-check: Nominatim by name vs postcodes.io centroid for the same school: median {nm.agree_m.median():.0f} m apart")
for lo, hi in [(0, AGREE), (AGREE, 1000), (1000, 1e9)]:
    s = nm[(nm.agree_m > lo) & (nm.agree_m <= hi)]
    print(f"  {lo:>5.0f}-{hi if hi < 1e9 else float('inf'):>5.0f} m apart: {len(s):3d} schools, error of the name match "
          f"median {s.error_m.median():.0f} m, max {s.error_m.max():.0f} m")
print(f"combined rule (name match if within {AGREE} m of the postcode centroid, else postcode centroid): "
      f"{accept.sum()} from Nominatim by name, {(~accept).sum()} from postcodes.io")
print(summarise(combo.reset_index()).to_string())

print("\nNominatim free-text matches by addresstype (what the result actually is):")
print(ft.groupby(ft.addresstype.replace("", "no result")).error_m
      .agg(["size", "median", lambda s: s.quantile(0.9)]).rename(columns={"<lambda_0>": "p90"})
      .sort_values("size", ascending=False).round(0).to_string())

worst = res[res.method != "postcodes.io postcode"].sort_values("error_m", ascending=False).head(8)
print("\nlargest errors:")
for w in worst.itertuples():
    print(f"  {w.error_m / 1000:6.1f} km  {w.method:<22} {w.name} | {w.address}\n"
          f"            -> {w.display_name[:110]}")
far_pc = res[(res.method == "postcodes.io postcode") & (res.error_m > 500)]
on_centroid = (res.method == "postcodes.io postcode") & (res.error_m < 5)
print(f"\nGIAS coordinates within 5 m of their own postcode centroid: {on_centroid.sum()} of {len(schools)}")
print("postcode centroids more than 500 m from the GIAS coordinate:")
for w in far_pc.itertuples():
    others = res[(res.URN == w.URN) & (res.method != w.method)].error_m.round(0).tolist()
    print(f"  {w.error_m:6.0f} m  {w.name} | {w.address}  (Nominatim methods: {others} m)")
missed = res[res.lat.isna()].groupby("method").size()
print("\nno result at all:", missed.to_dict())

# ---------------------------------------------------------------- 5. reverse geocoding
sample = ref.sample(40, random_state=1).to_crs(4326)
rev = []
for r in sample.itertuples():
    raw = cached("reverse", f"{r.geometry.y:.6f}, {r.geometry.x:.6f}", zoom=18, addressdetails=True)
    a = (raw or {}).get("address", {})
    rev.append({"URN": r.URN, "name": r.name, "Postcode": r.Postcode, "lat": r.geometry.y, "lon": r.geometry.x,
                "nom_postcode": a.get("postcode", ""), "nom_class": f"{raw.get('class')}={raw.get('type')}" if raw else "",
                "nom_name": (raw or {}).get("name", ""), "nom_road": a.get("road", ""),
                "nom_display": (raw or {}).get("display_name", "")})
cache_file.write_text(json.dumps(cache), encoding="utf-8")
rev = pd.DataFrame(rev)
geoloc = [{"longitude": r.lon, "latitude": r.lat, "limit": 1, "radius": 2000} for r in rev.itertuples()]
resp = requests.post(f"{POSTCODES_IO}/postcodes", json={"geolocations": geoloc},
                     headers={"User-Agent": USER_AGENT}, timeout=60)
resp.raise_for_status()
rev["pio_postcode"] = [(it["result"] or [{}])[0].get("postcode", "") for it in resp.json()["result"]]
rev["pio_distance_m"] = [(it["result"] or [{}])[0].get("distance", np.nan) for it in resp.json()["result"]]
norm = lambda s: s.str.replace(" ", "").str.upper()
rev["nom_pc_match"] = norm(rev.nom_postcode) == norm(rev.Postcode)
rev["nom_district_match"] = norm(rev.nom_postcode).str[:-3] == norm(rev.Postcode).str[:-3]
rev["pio_pc_match"] = norm(rev.pio_postcode) == norm(rev.Postcode)
SCHOOL_CLASSES = ["amenity=school", "amenity=college", "amenity=university", "amenity=kindergarten", "building=school"]
rev["nom_is_school"] = rev.nom_class.isin(SCHOOL_CLASSES)
print(f"\nreverse geocoding {len(rev)} GIAS school coordinates:")
print(f"  Nominatim (zoom=18): returned a school feature for {rev.nom_is_school.sum()}, the same full postcode as "
      f"GIAS for {rev.nom_pc_match.sum()}, the same postcode district for {rev.nom_district_match.sum()}, "
      f"no postcode at all for {(rev.nom_postcode == '').sum()}")
print(f"  postcodes.io (nearest postcode centroid): same full postcode as GIAS for {rev.pio_pc_match.sum()}; "
      f"nearest centroid median {rev.pio_distance_m.median():.0f} m away")
print(rev.nom_class.value_counts().head(8).to_string())
print(rev[["name", "Postcode", "nom_class", "nom_name", "nom_postcode", "pio_postcode"]].head(10).to_string(index=False))
rev.to_csv(OUT / "reverse_geocoded_sample.csv", index=False)

# ---------------------------------------------------------------- 6. figures
lad_file = DATA / "sheffield_lad.geojson"
if not lad_file.exists():
    r = requests.get(LAD_URL, params={"where": "LAD26NM='Sheffield'", "outFields": "LAD26CD,LAD26NM", "outSR": 27700,
                                      "f": "geojson"}, timeout=120)
    r.raise_for_status()
    lad_file.write_bytes(r.content)
boundary = gpd.read_file(lad_file).set_crs(27700, allow_override=True)
centre = gpd.GeoSeries(gpd.points_from_xy([-1.4701], [53.3811]), crs=4326).to_crs(27700).iloc[0]  # city centre


def frame(ax, extent=None):
    boundary.plot(ax=ax, facecolor="#f4f2ec", edgecolor="#9a9488", lw=0.8, zorder=0)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_xlabel("")
    ax.set_ylabel("")
    ax.set_aspect("equal")
    if extent:
        ax.set_xlim(extent[0], extent[2])
        ax.set_ylim(extent[1], extent[3])


def save(fig, name):
    for ax in fig.axes:
        if "[metre]" in ax.get_xlabel():
            ax.set_xlabel("")
            ax.set_ylabel("")
    fig.savefig(OUT / name)
    plt.close(fig)


def km_bar(ax, length=5000):
    x0, x1 = ax.get_xlim()
    y0, y1 = ax.get_ylim()
    x, y = x0 + 0.05 * (x1 - x0), y0 + 0.05 * (y1 - y0)
    ax.plot([x, x + length], [y, y], color="#222", lw=2.5, solid_capstyle="butt")
    label = f"{length / 1000:g} km" if length >= 1000 else f"{length:g} m"
    ax.text(x + length / 2, y + 0.015 * (y1 - y0), label, ha="center", va="bottom", fontsize=8)


# figure 1: the input
fig, ax = plt.subplots(figsize=(7.5, 6.5))
frame(ax)
for t, grp in ref.groupby(ref["type"].str.contains("Academy")):
    grp.plot(ax=ax, markersize=10, color="#3288bd" if t else "#d53e4f", edgecolor="white", linewidth=0.3,
             label=f"{'academies' if t else 'other school types'} ({len(grp)})", zorder=2)
ax.legend(loc="upper right", fontsize=8)
km_bar(ax)
ax.set_title(f"{len(ref)} open establishments in Sheffield with the DfE's own coordinates (GIAS)")
save(fig, "01-gias-schools.png")

# figure 2: where each method put each school, as error vectors
city = (centre.x - 6000, centre.y - 6000, centre.x + 6000, centre.y + 6000)
fig, axs = plt.subplots(2, 2, figsize=(10, 10))
for ax, method in zip(axs.flat, METHODS):
    m = res[(res.method == method) & res.lat.notna()]
    p = gpd.GeoSeries(gpd.points_from_xy(m.lon, m.lat), crs=4326).to_crs(27700)
    frame(ax, city)
    for (ex, ny), q, err in zip(zip(m.Easting, m.Northing), p, m.error_m):
        ax.plot([ex, q.x], [ny, q.y], color=COLOURS[method] if err > 100 else "#999", lw=0.8, zorder=1)
    ax.scatter(m.Easting, m.Northing, s=6, color="#222", zorder=2)
    ax.scatter(p.x, p.y, s=10, color=COLOURS[method], edgecolor="white", linewidth=0.3, zorder=3)
    e = m.error_m
    ax.set_title(f"{method}: {len(m)} of {len(schools)} found\nmedian {e.median():.0f} m, {(e > 100).sum()} more than "
                 f"100 m out (coloured lines)", fontsize=8.5)
    km_bar(ax, 2000)
fig.suptitle("Central Sheffield: black dot = GIAS coordinate, coloured dot = geocoded position", y=0.995)
fig.tight_layout()
save(fig, "02-error-vectors.png")

# figure 3: ECDF of the distance error
fig, ax = plt.subplots(figsize=(8, 4.8))
for method in METHODS + [combo.method.iloc[0]]:
    e = (combo if method.startswith("Name if") else res[res.method == method]).error_m
    e = np.sort(e.dropna().clip(lower=1))
    y = np.arange(1, len(e) + 1) / len(schools) * 100
    ax.step(e, y, where="post", lw=2 if method.startswith("Name if") else 1.5,
            color=COLOURS.get(method, "#222"), ls="--" if method.startswith("Name if") else "-", label=method)
for d in (50, 100, 250, 1000):
    ax.axvline(d, color="#bbb", lw=0.7, zorder=0)
    ax.text(d, 2, f" {d} m", fontsize=7, color="#777")
ax.set_xscale("log")
ax.set_xlim(1, 30000)
ax.set_ylim(0, 100)
ax.set_xlabel("distance from the GIAS coordinate (m, log scale)")
ax.set_ylabel("schools geocoded within that distance (%)")
ax.legend(loc="upper left", fontsize=7.5)
ax.grid(alpha=0.25)
ax.set_title("Cumulative distance error by method (schools with no result never reach 100%)")
save(fig, "03-error-ecdf.png")

# figure 4: error by what Nominatim matched
nom = res[res.method.str.startswith("Nominatim") & res.lat.notna()]
order = nom.addresstype.value_counts().index.tolist()
fig, ax = plt.subplots(figsize=(8, 4.6))
rng = np.random.default_rng(0)
for i, at in enumerate(order):
    for j, method in enumerate(METHODS[:3]):
        e = nom[(nom.addresstype == at) & (nom.method == method)].error_m.clip(lower=1)
        ax.scatter(e, i + (j - 1) * 0.22 + rng.uniform(-0.07, 0.07, len(e)), s=9, color=COLOURS[method],
                   alpha=0.75, label=method if i == 0 else None)
ax.set_yticks(range(len(order)))
ax.set_yticklabels([f"{a} ({(nom.addresstype == a).sum()})" for a in order])
ax.invert_yaxis()
ax.set_xscale("log")
ax.set_xlabel("distance from the GIAS coordinate (m, log scale)")
ax.set_ylabel("Nominatim addresstype of the match")
ax.axvline(100, color="#bbb", lw=0.7)
ax.legend(fontsize=7.5, loc="lower right")
ax.grid(alpha=0.25, axis="x")
ax.set_title("What the match was decides how far out it is")
save(fig, "04-error-by-match-type.png")

# figure 5: reverse geocoding
fig, axs = plt.subplots(1, 2, figsize=(10.5, 4.4), gridspec_kw={"width_ratios": [1, 1.15]})
cats = ["full postcode matches GIAS", "district only", "different district", "no postcode returned"]
def pc_cat(got, want):
    if not got:
        return cats[3]
    g, w = got.replace(" ", "").upper(), want.replace(" ", "").upper()
    return cats[0] if g == w else cats[1] if g[:-3] == w[:-3] else cats[2]
counts = pd.DataFrame({"Nominatim reverse (zoom 18)": rev.apply(lambda r: pc_cat(r.nom_postcode, r.Postcode), axis=1)
                       .value_counts(), "postcodes.io nearest postcode": rev.apply(
                           lambda r: pc_cat(r.pio_postcode, r.Postcode), axis=1).value_counts()}).reindex(cats).fillna(0)
left = np.zeros(2)
for cat, colour in zip(cats, ["#66a61e", "#fee08b", "#d53e4f", "#bbbbbb"]):
    axs[0].barh(counts.columns, counts.loc[cat], left=left, color=colour, label=cat)
    left += counts.loc[cat].values
axs[0].set_xlabel(f"schools (sample of {len(rev)})")
axs[0].legend(fontsize=7, loc="upper center", bbox_to_anchor=(0.5, -0.18), ncol=2)
axs[0].set_title("(a) postcode returned for a school's own coordinate")
kinds = rev.nom_class.replace("", "no result").value_counts()
axs[1].barh(kinds.index[::-1], kinds.values[::-1], color=["#3288bd" if k in SCHOOL_CLASSES
                                                          else "#f46d43" for k in kinds.index[::-1]])
axs[1].set_xlabel("schools")
axs[1].set_title("(b) what Nominatim says is at the school's coordinate")
fig.tight_layout()
save(fig, "05-reverse-geocoding.png")

# figure 6: cross-checking two geocoders flags the bad name matches without any reference data
fig, ax = plt.subplots(figsize=(8, 5))
ok = nm.agree_m.notna()
ax.scatter(nm.agree_m[ok].clip(lower=1), nm.error_m[ok].clip(lower=1), s=14,
           c=np.where(accept[ok], "#3288bd", "#d53e4f"), edgecolor="white", linewidth=0.3)
ax.axvline(AGREE, color="#555", lw=0.9, ls="--")
ax.text(AGREE * 1.08, 1.3, f"accept if within {AGREE} m", fontsize=8, color="#555")
for k, r in enumerate(nm[ok & (nm.error_m > 1000)].sort_values("error_m", ascending=False).itertuples()):
    ax.annotate(r.name, (max(r.agree_m, 1), r.error_m), xytext=(-8, 14 - 12 * k), textcoords="offset points",
                ha="right", va="center", fontsize=7, arrowprops=dict(arrowstyle="-", color="#888", lw=0.5))
ax.plot([1, 1e5], [1, 1e5], color="#bbb", lw=0.7, zorder=0)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlim(1, 1e5)
ax.set_ylim(1, 1e5)
ax.set_xlabel("distance between the Nominatim name match and the postcodes.io centroid (m)")
ax.set_ylabel("true error of the name match vs GIAS (m)")
ax.grid(alpha=0.25, which="major")
ax.set_title(f"Cross-check: {accept.sum()} name matches accepted (blue), {(ok & ~accept).sum()} rejected (red)")
save(fig, "06-cross-check.png")
print(f"\nfigures written to {OUT}; total run time {time.perf_counter() - t_start:.0f} s")