"""ERA5 climate reanalysis in Python with xarray: the July 2022 UK heatwave, from a GIS perspective.

Streams hourly ERA5 maximum 2 m temperature for July 2022 from Google's public ARCO-ERA5 Zarr store and
the 1991-2020 July climatology from the Microsoft Planetary Computer's ERA5 Zarr copy (no account or key
for either), turns hourly values into daily maxima, clips the grid to the UK polygon, computes
cos(latitude) area-weighted means, maps anomalies against the 1991-2020 baseline, exports them as Cloud
Optimised GeoTIFFs and compares ERA5 with the 37 Met Office historic stations.
Run with:  python era5_uk_heatwave.py
"""
import os
import platform
import re
import sys
import time
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import requests
import rioxarray  # noqa: F401  (registers the .rio accessor)
import xarray as xr
import zarr
from matplotlib.colors import ListedColormap, TwoSlopeNorm

HERE = Path(__file__).parent
OUT, DATA = HERE / "output", HERE / "data"
for d in (OUT, DATA):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)

print("=== versions ===")
print(f"python {platform.python_version()} | xarray {xr.__version__} | zarr {zarr.__version__} | "
      f"rioxarray {rioxarray.__version__} | rasterio {rasterio.__version__} | geopandas {gpd.__version__} | "
      f"numpy {np.__version__} | pandas {pd.__version__}")


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


# --- Step 1: open ARCO-ERA5 lazily ---------------------------------------------------
ARCO = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
LON, LAT = (-8.75, 2.0), (61.0, 49.75)  # UK box; latitude runs north to south in ERA5

section("Step 1: open ARCO-ERA5")
t0 = time.perf_counter()
arco = xr.open_zarr(ARCO, chunks=None, storage_options={"token": "anon"})
print(f"opened in {time.perf_counter() - t0:.1f} s: {len(arco.data_vars)} variables, dims {dict(arco.sizes)}")
print(f"valid ERA5 data: {arco.attrs['valid_time_start']} to {arco.attrs['valid_time_stop']} "
      f"(ERA5T, preliminary, to {arco.attrs['valid_time_stop_era5t']})")
mx2t = arco["maximum_2m_temperature_since_previous_post_processing"]
print(f"mx2t: {mx2t.attrs['long_name']} [{mx2t.attrs['units']}], chunks {mx2t.encoding['chunks']}, "
      f"{mx2t.size * 4 / 1e12:.1f} TB if loaded")


def to_uk(da: xr.DataArray) -> xr.DataArray:
    """Convert 0..360 longitudes to -180..180 and cut out the UK box (still lazy)."""
    da = da.rename({k: v for k, v in {"lat": "latitude", "lon": "longitude"}.items() if k in da.dims})
    lon = da["longitude"]
    da = da.assign_coords(longitude=xr.where(lon > 180, lon - 360, lon)).sortby("longitude")
    return da.sel(longitude=slice(*LON), latitude=slice(*LAT))


def daily_max(hourly: xr.DataArray) -> xr.DataArray:
    """mx2t at hh:00 is the maximum over the previous hour, so shift by 1 h before taking the daily max."""
    shifted = hourly.assign_coords(time=hourly["time"] - np.timedelta64(1, "h"))
    return shifted.resample(time="1D").max()


# --- Step 2: hourly July 2022 -> daily maximum temperature -----------------------------
section("Step 2: July 2022 daily maximum temperature from ARCO-ERA5")
f2022 = DATA / "era5_uk_mx2t_2022-07_daily.nc"
if not f2022.exists():
    t0 = time.perf_counter()
    hourly = to_uk(mx2t).sel(time=slice("2022-07-01T01", "2022-08-01T00")).load() - 273.15
    print(f"downloaded {hourly.sizes['time']} hourly UK grids {hourly.shape} in {time.perf_counter() - t0:.0f} s")
    daily_max(hourly).rename("tmax").to_netcdf(f2022)
    static = arco[["land_sea_mask", "geopotential_at_surface"]].sel(time="2022-07-01T00").load()
    static.map(to_uk).drop_vars("time").to_netcdf(DATA / "era5_uk_static.nc")
tmax22 = xr.open_dataarray(f2022).load()
static = xr.open_dataset(DATA / "era5_uk_static.nc").load()
print(f"daily Tmax grid: {dict(tmax22.sizes)}; {tmax22.longitude.values[0]}..{tmax22.longitude.values[-1]} E, "
      f"{tmax22.latitude.values[-1]}..{tmax22.latitude.values[0]} N")
hot = tmax22.sel(time="2022-07-19").squeeze(drop=True)
imax = np.unravel_index(int(hot.argmax()), hot.shape)
print(f"hottest ERA5 cell on 19 Jul 2022: {float(hot.max()):.1f} °C at {float(hot.latitude[imax[0]]):.2f} N, "
      f"{float(hot.longitude[imax[1]]):.2f} E")

# --- Step 3: 1991-2020 July climatology from the Planetary Computer ---------------------
section("Step 3: 1991-2020 July climatology from the Planetary Computer")
fclim = DATA / "era5_uk_mx2t_july_1991-2020_daily.nc"
if not fclim.exists():
    catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1",
                                        modifier=planetary_computer.sign_inplace)
    years, t0 = [], time.perf_counter()
    for year in range(1991, 2021):
        item = next(catalog.search(collections=["era5-pds"], datetime=f"{year}-07",
                                   query={"era5:kind": {"eq": "fc"}}).items())
        asset = item.assets["air_temperature_at_2_metres_1hour_Maximum"]
        name = "air_temperature_at_2_metres_1hour_Maximum"
        ds = xr.open_dataset(asset.href, chunks={}, **asset.extra_fields["xarray:open_kwargs"])
        if year == 1991:
            print(f"{item.id}: chunks {ds[name].encoding['preferred_chunks']}")
        hourly = to_uk(ds[name]).load() - 273.15
        years.append(daily_max(hourly).sel(time=f"{year}-07"))
    clim_daily = xr.concat(years, "time").rename("tmax")
    print(f"downloaded 30 Julys in {time.perf_counter() - t0:.0f} s")
    clim_daily.to_netcdf(fclim)
clim_daily = xr.open_dataarray(fclim).load()
print(f"climatology: {clim_daily.sizes['time']} July days, {clim_daily.time.dt.year.min().item()}-"
      f"{clim_daily.time.dt.year.max().item()}")

# Both stores hold the same ERA5 grid: check they agree where they overlap
pc20 = clim_daily.sel(time="2020-07-15").squeeze(drop=True)
arco20 = daily_max(to_uk(mx2t).sel(time=slice("2020-07-15T01", "2020-07-16T00")).load() - 273.15)
print(f"ARCO vs Planetary Computer, 15 Jul 2020 daily Tmax: max |diff| "
      f"{float(abs(arco20.squeeze(drop=True) - pc20).max()):.3f} °C")

clim_mean = clim_daily.mean("time")                       # mean July daily maximum, 1991-2020
clim_record = clim_daily.max("time")                      # highest July daily maximum, 1991-2020
anom_month = tmax22.mean("time") - clim_mean              # July 2022 mean Tmax anomaly
anom_19 = hot - clim_mean                                 # 19 July 2022 anomaly
above_record = hot - clim_record                          # > 0 where 19 July beat every July day of 1991-2020

# --- Step 4: clip to the UK polygon ---------------------------------------------------
section("Step 4: clip to the UK polygon")
ne_zip = DATA / "ne_10m_admin_0_countries.zip"
if not ne_zip.exists():
    ne_zip.write_bytes(requests.get("https://naciscdn.org/naturalearth/10m/cultural/ne_10m_admin_0_countries.zip",
                                    timeout=120).content)
countries = gpd.read_file(ne_zip)
uk = countries[countries["ADMIN"] == "United Kingdom"].to_crs(4326)
neighbours = countries.cx[LON[0] - 3:LON[1] + 3, LAT[1] - 3:LAT[0] + 3]


def georef(da: xr.DataArray) -> xr.DataArray:
    return da.rio.set_spatial_dims(x_dim="longitude", y_dim="latitude").rio.write_crs("EPSG:4326")


print(f"grid transform: {tuple(round(v, 3) for v in georef(clim_mean).rio.transform())[:6]}")
clip_centre = georef(clim_mean).rio.clip(uk.geometry, all_touched=False, drop=False)
clip_touch = georef(clim_mean).rio.clip(uk.geometry, all_touched=True, drop=False)
land = static["land_sea_mask"] > 0.5
inside = clip_centre.notnull()
n_c, n_t = int(inside.sum()), int(clip_touch.notnull().sum())
print(f"cells with centre inside the UK: {n_c}; touching the UK: {n_t}; land-sea mask > 0.5 in box: {int(land.sum())}")
print(f"of the centre-inside cells, {int((inside & land).sum())} are also land (lsm > 0.5)")

# --- Step 5: area-weighted means ------------------------------------------------------
section("Step 5: area-weighted UK means")
weights = np.cos(np.deg2rad(clim_mean.latitude))
print(f"cos(lat) weight: {float(weights.min()):.3f} at {float(clim_mean.latitude.max()):.2f} N, "
      f"{float(weights.max()):.3f} at {float(clim_mean.latitude.min()):.2f} N")
R = 6_371_008.8
cell_km2 = (R ** 2 * np.deg2rad(0.25) * (np.sin(np.deg2rad(clim_mean.latitude + 0.125))
                                         - np.sin(np.deg2rad(clim_mean.latitude - 0.125)))) / 1e6
print(f"cell area: {float(cell_km2.max()):.0f} km² in the south, {float(cell_km2.min()):.0f} km² in the north; "
      f"UK cells total {float((cell_km2 * inside).sum()):,.0f} km²")


def uk_mean(da: xr.DataArray, weighted: bool = True) -> xr.DataArray:
    da = da.where(inside)
    return da.weighted(weights).mean(("latitude", "longitude")) if weighted else da.mean(("latitude", "longitude"))


m22, mclim = uk_mean(tmax22.mean("time")), uk_mean(clim_mean)
print(f"UK mean July daily Tmax 1991-2020: {float(mclim):.2f} °C (unweighted {float(uk_mean(clim_mean, False)):.2f})")
print(f"UK mean July daily Tmax 2022:      {float(m22):.2f} °C, anomaly {float(m22 - mclim):+.2f} °C")
d22 = uk_mean(tmax22)
dclim = uk_mean(clim_daily)
by_day = dclim.groupby("time.day")
clim_q = pd.DataFrame({"mean": by_day.mean().values, "p05": by_day.quantile(0.05).values,
                       "p95": by_day.quantile(0.95).values, "max": by_day.max().values}, index=range(1, 32))
ser = pd.Series(d22.values, index=d22.time.dt.day.values)
for day, v in ser.sort_values(ascending=False).head(3).items():
    print(f"  {day:2d} Jul 2022: UK-mean Tmax {v:.1f} °C, 1991-2020 mean for the day {clim_q.loc[day, 'mean']:.1f}, "
          f"highest of the 30 years {clim_q.loc[day, 'max']:.1f}")
print(f"warmest UK-mean July day in 1991-2020: {float(dclim.max()):.1f} °C on "
      f"{pd.Timestamp(dclim.idxmax().values):%d %b %Y}")
print(f"19 Jul 2022 anomaly: mean {float(uk_mean(anom_19)):+.1f} °C, max {float(anom_19.where(inside).max()):+.1f} °C; "
      f"{float((above_record > 0).where(inside).weighted(weights).mean()) * 100:.0f}% of UK area above its "
      f"1991-2020 July record")

# --- Step 6: export anomalies as Cloud Optimised GeoTIFFs -------------------------------
section("Step 6: export to Cloud Optimised GeoTIFF")
stack = xr.concat([anom_month, anom_19, above_record], dim="band").assign_coords(band=[1, 2, 3])
stack = georef(stack.astype("float32").where(clip_touch.notnull())).rio.write_nodata(np.nan)
stack.attrs = {"long_name": ("July 2022 mean Tmax anomaly", "19 July 2022 Tmax anomaly",
                             "19 July 2022 Tmax minus 1991-2020 July record"), "units": "degC"}
cog = OUT / "era5_uk_tmax_anomaly_2022-07.tif"
stack.rio.to_raster(cog, driver="COG", compress="DEFLATE", predictor=3)
with rasterio.open(cog) as src:
    print(f"{cog.name}: {src.driver}, {src.count} bands {src.dtypes[0]}, {src.width} x {src.height}, {src.crs}, "
          f"res {src.res}, layout {src.tags(ns='IMAGE_STRUCTURE').get('LAYOUT')}, "
          f"{cog.stat().st_size / 1024:.0f} KB")
    print(f"bounds {tuple(round(b, 3) for b in src.bounds)}; band 2 at Heathrow: "
          f"{next(src.sample([(-0.449, 51.479)], indexes=2))[0]:.1f} °C")
    print("band descriptions:", src.descriptions)

# --- Step 7: compare with Met Office station records ----------------------------------
section("Step 7: compare with Met Office historic stations")
STATIONS = ("aberporth armagh ballypatrick bradford braemar camborne cambridge cardiff chivenor cwmystwyth "
            "dunstaffnage durham eastbourne eskdalemuir heathrow hurn lerwick leuchars lowestoft manston nairn "
            "newtonrigg oxford paisley ringway rossonwye shawbury sheffield southampton stornoway suttonbonington "
            "tiree valley waddington whitby wickairport yeovilton").split()
MO = "https://www.metoffice.gov.uk/pub/data/weather/uk/climate/stationdata/{}data.txt"
fst = DATA / "metoffice_july_tmax.csv"
if not fst.exists():
    rows = []
    for s in STATIONS:
        txt = requests.get(MO.format(s), timeout=60).text
        lines = txt.splitlines()
        lat, lon = map(float, re.search(r"Lat\s*([-\d.]+)\s*,?\s*Lon\s*([-\d.]+)", txt).groups())
        alt = float(re.search(r"(-?[\d.]+)\s*m(?:etres)?\s*amsl", txt).group(1))
        for line in lines:
            p = line.split()
            if len(p) >= 3 and re.fullmatch(r"\d{4}", p[0]) and p[1] == "7":
                tx = p[2].rstrip("*#")
                rows.append({"station": re.split(r"\s*[/(]", lines[0].strip())[0], "lat": lat, "lon": lon, "alt": alt, "year": int(p[0]),
                             "tmax": float(tx) if re.fullmatch(r"-?[\d.]+", tx) else np.nan})
    pd.DataFrame(rows).to_csv(fst, index=False)
st = pd.read_csv(fst)
meta = st.groupby("station")[["lat", "lon", "alt"]].first()
print(f"{len(meta)} stations, {st['tmax'].notna().sum()} July means; {st.query('year == 2022')['tmax'].notna().sum()} "
      f"stations reported July 2022")

z_era5 = static["geopotential_at_surface"] / 9.80665       # surface geopotential (m2 s-2) / g -> height (m)
pts = dict(latitude=xr.DataArray(meta["lat"].values, dims="station"),
           longitude=xr.DataArray(meta["lon"].values, dims="station"))
# Nearest land cell (lsm > 0.5): coastal stations otherwise pick up cells that are mostly sea
cells = land.stack(cell=("latitude", "longitude"))
cells = cells[cells.values]
dy = meta["lat"].values[:, None] - cells["latitude"].values[None, :]
dx = (meta["lon"].values[:, None] - cells["longitude"].values[None, :]) * np.cos(np.deg2rad(meta["lat"].values))[:, None]
j = np.argmin(dx ** 2 + dy ** 2, axis=1)
land_pts = dict(latitude=xr.DataArray(cells["latitude"].values[j], dims="station"),
                longitude=xr.DataArray(cells["longitude"].values[j], dims="station"))
moved = int((np.hypot(dx, dy)[np.arange(len(j)), j] > 0.18).sum())
print(f"{moved} of {len(meta)} stations are nearer a sea or coastal cell than the centre of a land cell")


def at(da, p):
    return da.sel(**p, method="nearest").values.astype(float)


lapse = lambda p: 0.0065 * (at(z_era5, p) - meta["alt"].values)  # noqa: E731  ERA5 cell higher -> ERA5 too cool
cmp = meta.assign(era_z=at(z_era5, pts).round(), lsm=at(static["land_sea_mask"], pts).round(2))
cmp["obs_2022"] = st.query("year == 2022").set_index("station")["tmax"].reindex(meta.index)
cmp["era5_2022"] = at(tmax22.mean("time"), pts)
cmp["era5_2022_lapse"] = cmp["era5_2022"] + lapse(pts)
cmp["era5_2022_land"] = at(tmax22.mean("time"), land_pts) + lapse(land_pts)
base = st.query("1991 <= year <= 2020").pivot(index="year", columns="station", values="tmax")[meta.index]
era_yr = clim_daily.groupby("time.year").mean()
era_base = pd.DataFrame(at(era_yr, pts), index=era_yr.year.values, columns=meta.index)
era_base_land = pd.DataFrame(at(era_yr, land_pts), index=era_yr.year.values, columns=meta.index)
ok = base.notna().sum() >= 25                                                # >= 25 of the 30 Julys
cmp["obs_clim"] = base.mean().where(ok)
cmp["era5_clim_land"] = era_base_land.where(base.notna()).mean().where(ok) + lapse(land_pts)


def stats(o, e):
    m = o.notna() & e.notna()
    d = e[m] - o[m]
    return f"n={m.sum()}, bias {d.mean():+.2f} °C, MAE {d.abs().mean():.2f} °C, r={np.corrcoef(o[m], e[m])[0, 1]:.2f}"


print("July 2022 mean Tmax, ERA5 vs station")
print("  nearest cell:                       ", stats(cmp["obs_2022"], cmp["era5_2022"]))
print("  + 6.5 °C/km lapse-rate correction:  ", stats(cmp["obs_2022"], cmp["era5_2022_lapse"]))
print("  nearest land cell + lapse rate:     ", stats(cmp["obs_2022"], cmp["era5_2022_land"]))
print("1991-2020 July mean Tmax, nearest land cell + lapse rate:", stats(cmp["obs_clim"], cmp["era5_clim_land"]))
cmp["diff_2022"] = cmp["era5_2022_land"] - cmp["obs_2022"]
print(cmp.dropna(subset=["diff_2022"]).sort_values("diff_2022")[
    ["alt", "era_z", "lsm", "obs_2022", "era5_2022", "era5_2022_land", "diff_2022"]].round(1)
    .iloc[[0, 1, 2, -3, -2, -1]].to_string())
hr = "Heathrow"
yr_ok = base[hr].notna()
print(f"Heathrow July means 1991-2020: r={np.corrcoef(base[hr][yr_ok], era_base[hr][yr_ok])[0, 1]:.2f} "
      f"between station and ERA5 year to year")
print(f"Heathrow 19 Jul 2022: ERA5 cell {float(hot.sel(latitude=51.479, longitude=-0.449, method='nearest')):.1f} °C "
      f"(the station recorded 40.2 °C that day)")

# --- Figures -------------------------------------------------------------------------
section("Figures")
EXT = dict(xlim=(LON[0], LON[1]), ylim=(LAT[1], LAT[0]))


def basemap(ax, title):
    neighbours.boundary.plot(ax=ax, color="#555", lw=0.4)
    uk.boundary.plot(ax=ax, color="black", lw=0.6)
    ax.set(**EXT, title=title, xlabel="", ylabel="")
    ax.set_aspect(1 / np.cos(np.deg2rad(55)))
    ax.tick_params(labelsize=7)


# 1. Daily maximum on 19 July 2022 (cells drawn as the areas they are)
fig, ax = plt.subplots(figsize=(5.2, 6))
im = hot.plot.pcolormesh(ax=ax, cmap="Spectral_r", vmin=14, vmax=40, add_colorbar=False)
basemap(ax, "ERA5 daily maximum 2 m temperature, 19 July 2022")
ax.plot(hot.longitude[imax[1]], hot.latitude[imax[0]], marker="^", color="black", ms=6)
ax.annotate(f"{float(hot.max()):.1f} °C", (hot.longitude[imax[1]], hot.latitude[imax[0]]), xytext=(6, -10),
            textcoords="offset points", fontsize=8, fontweight="bold")
fig.colorbar(im, ax=ax, shrink=0.7, label="°C")
fig.savefig(OUT / "01-era5-tmax-19-july-2022.png")
plt.close(fig)

# 2. Clipping and weights
fig, axes = plt.subplots(1, 3, figsize=(11, 4.6))
both = xr.where(inside, 2, xr.where(clip_touch.notnull(), 1, 0))
axes[0].pcolormesh(both.longitude, both.latitude, both, cmap=ListedColormap(["#f2f4f3", "#fdae61", "#2b83ba"]),
                   shading="nearest")
basemap(axes[0], f"Centre inside (blue, {n_c}) and\nonly touching (orange, {n_t - n_c}) cells")
lsm = static["land_sea_mask"]
im = axes[1].pcolormesh(lsm.longitude, lsm.latitude, lsm, cmap="YlGn", shading="nearest", vmin=0, vmax=1)
basemap(axes[1], "ERA5 land-sea mask\n(land fraction of each cell)")
fig.colorbar(im, ax=axes[1], shrink=0.7)
wmap = (weights * xr.ones_like(clim_mean)).where(inside)
im = axes[2].pcolormesh(wmap.longitude, wmap.latitude, wmap, cmap="Spectral", shading="nearest")
basemap(axes[2], "cos(latitude) weight\nof the UK cells")
fig.colorbar(im, ax=axes[2], shrink=0.7)
fig.savefig(OUT / "02-clip-mask-and-weights.png")
plt.close(fig)

# 3. Anomaly maps
fig, axes = plt.subplots(1, 3, figsize=(11, 4.8))
for ax, da, title, lim in [(axes[0], anom_month, "July 2022 mean Tmax\nminus 1991–2020 July mean", 4),
                           (axes[1], anom_19, "19 July 2022 Tmax\nminus 1991–2020 July mean", 18),
                           (axes[2], above_record, "19 July 2022 Tmax minus\n1991–2020 July record", 8)]:
    im = da.where(clip_touch.notnull()).plot.pcolormesh(ax=ax, cmap="RdBu_r", norm=TwoSlopeNorm(0, -lim, lim),
                                                        add_colorbar=False)
    basemap(ax, title)
    fig.colorbar(im, ax=ax, shrink=0.7, orientation="horizontal", pad=0.08, label="°C")
fig.savefig(OUT / "03-anomaly-maps.png")
plt.close(fig)

# 4. UK area-weighted daily series
fig, ax = plt.subplots(figsize=(8, 3.8))
ax.fill_between(clim_q.index, clim_q["p05"], clim_q["p95"], color="#abd9e9", label="1991–2020, 5–95th percentile")
ax.plot(clim_q.index, clim_q["mean"], color="#2c7bb6", lw=1.5, label="1991–2020 mean")
ax.plot(clim_q.index, clim_q["max"], color="#2c7bb6", lw=0.8, ls="--", label="1991–2020 highest")
ax.plot(ser.index, ser.values, color="#d7191c", lw=2, marker="o", ms=3, label="2022")
ax.set(xlabel="Day of July", ylabel="UK area-weighted mean Tmax (°C)", xlim=(1, 31),
       title="UK daily maximum temperature in July: 2022 against the 1991–2020 ERA5 climatology")
ax.legend(loc="upper left", fontsize=8)
ax.grid(alpha=0.3)
fig.savefig(OUT / "04-uk-daily-tmax-july-2022.png")
plt.close(fig)

# 5. Station comparison
fig, axes = plt.subplots(1, 2, figsize=(10, 4.4), gridspec_kw={"width_ratios": [1, 1.4]})
ax = axes[0]
ax.scatter(cmp["obs_2022"], cmp["era5_2022"], s=18, facecolor="none", edgecolor="#f46d43", label="ERA5 nearest cell")
ax.scatter(cmp["obs_2022"], cmp["era5_2022_land"], s=18, color="#2b83ba", label="nearest land cell + lapse rate")
lims = [13, 29]
ax.plot(lims, lims, color="#555", lw=0.8)
ax.set(xlim=lims, ylim=lims, xlabel="Station July 2022 mean Tmax (°C)", ylabel="ERA5 July 2022 mean Tmax (°C)",
       title=f"{int(cmp['obs_2022'].notna().sum())} Met Office stations, July 2022")
ax.set_aspect("equal")
ax.legend(fontsize=7, loc="upper left")
ax = axes[1]
yrs = base.index
ax.plot(yrs, base[hr], color="#333", marker="o", ms=3, lw=1, label="Heathrow station")
ax.plot(yrs, era_base[hr], color="#2b83ba", marker="s", ms=3, lw=1, label="ERA5 nearest cell")
ax.plot([2022], [cmp.loc[hr, "obs_2022"]], color="#333", marker="o", ms=5)
ax.plot([2022], [cmp.loc[hr, "era5_2022"]], color="#2b83ba", marker="s", ms=5)
ax.set(xlabel="Year", ylabel="July mean Tmax (°C)", title="Heathrow: July mean daily maximum, 1991–2020 and 2022")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.savefig(OUT / "05-station-comparison.png")
plt.close(fig)
print("figures written:", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))
sys.stdout.flush()
# gcsfs can hang at interpreter shutdown on Windows while closing its event loop; everything is written by now.
os._exit(0)