"""Monitoring reservoir water quality from space: Sentinel-2 over four Peak District reservoirs.

Gets reservoir outlines from OpenStreetMap, searches the Microsoft Planetary Computer STAC API
for every Sentinel-2 L2A scene since 2019, masks cloud and shadow with the scene classification
(SCL) band, keeps only strictly-water pixels well inside each shoreline, and computes relative
water-quality indices (NDCI for chlorophyll, NDTI and red reflectance for turbidity, a green/blue
colour ratio for peat colour). It then builds a multi-year time series and monthly climatology per
reservoir, ranks the most anomalous dates, maps one summer date and compares the results with
public Environment Agency samples.
Run with:  python reservoir_water_quality.py
"""
import io
import platform
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.patheffects as pe
import matplotlib.pyplot as plt
import numpy as np
import osmnx as ox
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import rasterio.features
import requests
from matplotlib.patches import Rectangle
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.windows import from_bounds
from scipy import ndimage
from scipy.stats import spearmanr
from shapely.geometry import box

HERE = Path(__file__).parent
OUT, DATA, CACHE = HERE / "output", HERE / "data", HERE / "cache"
for d in (OUT, DATA, CACHE):
    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()} | pystac-client {pystac_client.__version__} | "
      f"planetary-computer {version('planetary-computer')} | rasterio {rasterio.__version__} | "
      f"osmnx {ox.__version__} | numpy {np.__version__} | pandas {pd.__version__}")


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


# --- Step 1: reservoir outlines from OpenStreetMap ---------------------------------
RESERVOIRS = ["Howden", "Derwent", "Ladybower", "Torside"]   # north to south, then Longdendale
COLOURS = {"Howden": "#1d6fa3", "Derwent": "#2f9e44", "Ladybower": "#e8590c", "Torside": "#862e9c"}
CRS = "EPSG:32630"            # UTM zone 30N, the native grid of Sentinel-2 tile 30UWE
INSET_M = 40                  # shrink each outline by 40 m (two 20 m pixels) to avoid mixed shoreline pixels

outlines = DATA / "reservoirs.gpkg"
if not outlines.exists():
    ox.settings.cache_folder = CACHE / "osmnx"
    osm = None
    for attempt in range(5):   # public Overpass servers are often busy: retry with backoff and a mirror
        for url in ("https://overpass-api.de/api", "https://overpass.kumi.systems/api"):
            ox.settings.overpass_url = url
            try:
                osm = ox.features_from_bbox((-1.95, 53.33, -1.62, 53.55), {"water": "reservoir"})
                break
            except Exception as exc:
                print(f"{url} failed: {str(exc).splitlines()[0][:70]}")
        if osm is not None:
            break
        time.sleep(30 * (attempt + 1))
    osm = osm[osm.geometry.type.isin(["Polygon", "MultiPolygon"]) & osm["name"].notna()]
    osm = osm[osm["name"].isin([f"{r} Reservoir" for r in RESERVOIRS])].reset_index()
    osm[["id", "name", "geometry"]].to_file(outlines, driver="GPKG")
res = gpd.read_file(outlines).to_crs(CRS)
res["reservoir"] = res["name"].str.replace(" Reservoir", "")
res = res.dissolve("reservoir").loc[RESERVOIRS].reset_index()
res["area_ha"] = res.area / 1e4
res["core"] = res.buffer(-INSET_M)
res["core_ha"] = res["core"].area / 1e4
section("reservoir outlines (OpenStreetMap)")
print(res[["reservoir", "id", "area_ha", "core_ha"]].round(1).to_string(index=False))

# --- Step 2: one analysis grid at 20 m ------------------------------------------------
RES = 20
x0, y0, x1, y1 = res.total_bounds
x0, y0 = np.floor((x0 - 600) / RES) * RES, np.floor((y0 - 600) / RES) * RES
x1, y1 = np.ceil((x1 + 600) / RES) * RES, np.ceil((y1 + 600) / RES) * RES
WIDTH, HEIGHT = int((x1 - x0) / RES), int((y1 - y0) / RES)
TRANSFORM = from_origin(x0, y1, RES, RES)
BOUNDS = (x0, y0, x1, y1)
# Label raster: 1..4 = core (inset) pixels of each reservoir, 0 = everything else.
labels = rasterio.features.rasterize(((g, i + 1) for i, g in enumerate(res["core"])), out_shape=(HEIGHT, WIDTH),
                                     transform=TRANSFORM, fill=0, dtype="uint8")
section("analysis grid")
print(f"{WIDTH} x {HEIGHT} pixels of {RES} m ({(x1 - x0) / 1000:.1f} x {(y1 - y0) / 1000:.1f} km), {CRS}")
for i, name in enumerate(RESERVOIRS, 1):
    print(f"{name:<10} {(labels == i).sum():>5} core water pixels")

# --- Step 3: search the Sentinel-2 L2A catalogue -----------------------------------------
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
PERIOD = "2019-01-01/2026-10-14"
bbox_ll = gpd.GeoSeries([box(*BOUNDS)], crs=CRS).to_crs(4326).total_bounds
catalog = pystac_client.Client.open(STAC_URL)
search = catalog.search(collections=["sentinel-2-l2a"], bbox=bbox_ll, datetime=PERIOD,
                        query={"eo:cloud_cover": {"lt": 95}, "s2:mgrs_tile": {"eq": "30UWE"}})
items = list(search.items())
# Some dates exist twice (an original and a reprocessed baseline): keep the newest processing.
best = {}
for it in items:
    key = (it.datetime.date(), it.properties["sat:relative_orbit"])
    if key not in best or it.properties["s2:processing_baseline"] > best[key].properties["s2:processing_baseline"]:
        best[key] = it
items = sorted(best.values(), key=lambda i: i.datetime)
section("STAC search")
print(f"{len(items)} Sentinel-2 L2A scenes of tile 30UWE with < 95% cloud, "
      f"{items[0].datetime:%d %b %Y} to {items[-1].datetime:%d %b %Y}")
print(pd.Series([i.properties["s2:processing_baseline"] for i in items]).value_counts().sort_index().to_dict())

# --- Step 4: read, mask and summarise every scene ----------------------------------------------
BANDS = ["B02", "B03", "B04", "B05", "B06", "B08", "B11"]
BAD_SCL = [0, 1, 3, 8, 9, 10, 11]   # no data, saturated, cloud shadow, cloud (medium, high), cirrus, snow
CLOUD_BUFFER_PX = 5                  # grow cloud and shadow by 100 m: their edges are rarely classified well
MIN_CLEAR = 0.5                      # need at least half of a reservoir's core pixels to be clear
NIR_MAX = 0.04                       # water is dark in near-infrared; exposed mud, weed and boats are not
GDAL_ENV = dict(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", GDAL_HTTP_MAX_RETRY=5, GDAL_HTTP_RETRY_DELAY=2,
                GDAL_HTTP_MERGE_CONSECUTIVE_RANGES="YES")


def read_band(item, band: str, res: int = RES) -> np.ndarray:
    """Read one band over the analysis grid, averaged to 20 m (10 m bands) or as-is (20 m bands)."""
    href = planetary_computer.sign(item.assets[band].href)   # sign at read time: tokens expire
    shape = (int(HEIGHT * RES / res), int(WIDTH * RES / res))
    for attempt in range(4):
        try:
            with rasterio.Env(**GDAL_ENV), rasterio.open(href) as src:
                win = from_bounds(*BOUNDS, src.transform)
                return src.read(1, window=win, out_shape=shape, boundless=True, fill_value=0,
                                resampling=Resampling.average if band != "SCL" else Resampling.nearest)
        except rasterio.errors.RasterioIOError:
            time.sleep(5 * (attempt + 1))
    raise RuntimeError(f"could not read {band} of {item.id}")


def to_reflectance(dn: np.ndarray, item) -> np.ndarray:
    """L2A digital numbers to surface reflectance. Baseline 04.00+ (from 25 Jan 2022) adds a -1000 offset."""
    offset = -1000 if float(item.properties["s2:processing_baseline"]) >= 4.0 else 0
    return np.where(dn > 0, (dn.astype("float32") + offset) / 10000, np.nan)


def process(item) -> list[dict]:
    """Cloud-mask one scene and return the median water spectrum of each reservoir."""
    p = item.properties
    base = {"date": item.datetime.strftime("%Y-%m-%d"), "item_id": item.id, "orbit": p["sat:relative_orbit"],
            "baseline": p["s2:processing_baseline"], "sun_zenith": round(p["s2:mean_solar_zenith"], 1)}
    scl = read_band(item, "SCL")
    bad = ndimage.binary_dilation(np.isin(scl, BAD_SCL[1:]), iterations=CLOUD_BUFFER_PX) | (scl == 0)
    clear_frac = {i: float((~bad[labels == i]).mean()) for i in range(1, len(RESERVOIRS) + 1)}
    rows = [dict(base, reservoir=n, clear_frac=round(clear_frac[i], 3)) for i, n in enumerate(RESERVOIRS, 1)]
    if max(clear_frac.values()) < MIN_CLEAR:
        return rows                      # too cloudy: skip the other bands entirely
    r = {b: to_reflectance(read_band(item, b), item) for b in BANDS}
    with np.errstate(divide="ignore", invalid="ignore"):
        mndwi = (r["B03"] - r["B11"]) / (r["B03"] + r["B11"])
    water = ~bad & (r["B08"] < NIR_MAX) & np.isfinite(r["B11"])   # NIR-dark: no exposed bank, weed or boats
    for i, row in enumerate(rows, 1):
        core, w = labels == i, water & (labels == i)
        clear = core & ~bad
        row.update(n_pixels=int(w.sum()), water_frac=round(float(w.sum() / core.sum()), 3),
                   scl_water_frac=round(float((scl[clear] == 6).mean()), 3) if clear.any() else np.nan,
                   mndwi_pos_frac=round(float((mndwi[clear] > 0).mean()), 3) if clear.any() else np.nan)
        if clear_frac[i] >= MIN_CLEAR and w.sum() >= 20:
            row.update({f"rho_{b}": round(float(np.median(r[b][w])), 5) for b in BANDS})
    return rows


cache_file = CACHE / "observations.csv"
cached = pd.read_csv(cache_file) if cache_file.exists() else pd.DataFrame(columns=["item_id"])
todo = [it for it in items if it.id not in set(cached["item_id"])]
section("reading Sentinel-2 over the reservoirs")
print(f"{len(items) - len(todo)} scenes already cached, {len(todo)} to read")
new_rows, t0 = [], time.time()
with ThreadPoolExecutor(max_workers=8) as pool:
    futures = [pool.submit(process, it) for it in todo]
    for k, fut in enumerate(as_completed(futures), 1):
        new_rows.extend(fut.result())
        if k % 100 == 0 or k == len(todo):   # save progress so an interrupted run can resume
            print(f"  {k}/{len(todo)} scenes, {time.time() - t0:.0f} s")
            pd.concat([cached, pd.DataFrame(new_rows)], ignore_index=True).to_csv(cache_file, index=False)
cached = pd.read_csv(cache_file)
obs = cached.sort_values(["date", "reservoir"]).reset_index(drop=True)
obs["date"] = pd.to_datetime(obs["date"])

# --- Step 5: quality control and water-quality indices ------------------------------------------
MAX_ZENITH = 70        # low winter sun: long paths through the atmosphere and weak water signal
MIN_WATER = 0.6        # at least 60% of a reservoir's core pixels must pass the cloud and water tests
MAX_SWIR = 0.03        # water absorbs SWIR almost completely: a bright median B11 means haze, glint or thin cloud
MIN_RHO = 0.003        # B2-B5 darker than this: ratios are dominated by noise (Sen2Cor water error ~0.002-0.005)
BANDS_RHO = [f"rho_{b}" for b in BANDS]


def add_indices(df: pd.DataFrame) -> pd.DataFrame:
    """Relative indices computed from the median water spectrum of each reservoir and date."""
    b = {k: df[f"rho_{k}"] for k in BANDS}
    df["ndci"] = (b["B05"] - b["B04"]) / (b["B05"] + b["B04"])                  # chlorophyll-a proxy
    df["mci"] = b["B05"] - b["B04"] - (b["B06"] - b["B04"]) * (705 - 665) / (740 - 665)   # red-edge peak height
    df["ndti"] = (b["B04"] - b["B03"]) / (b["B04"] + b["B03"])                  # turbidity proxy
    df["red"] = b["B04"]                                                         # turbidity proxy (reflectance)
    df["green_blue"] = b["B03"] / b["B02"]                                       # colour (CDOM, peat) proxy
    df["mndwi"] = (b["B03"] - b["B11"]) / (b["B03"] + b["B11"])                  # kept only as a diagnostic
    return df


def qc_reason(row) -> str:
    if row["clear_frac"] < MIN_CLEAR or pd.isna(row["rho_B04"]):
        return "cloud or shadow"
    if row["sun_zenith"] > MAX_ZENITH:
        return "low sun"
    if row["water_frac"] < MIN_WATER:
        return "too few water pixels"
    if row["rho_B11"] > MAX_SWIR:
        return "haze or glint"
    if (row[BANDS_RHO[:-1]] <= 0).any():
        return "negative reflectance"
    if (row[["rho_B02", "rho_B03", "rho_B04", "rho_B05"]] < MIN_RHO).any():
        return "signal too dark"
    return "ok"


obs = add_indices(obs)
obs["qc"] = obs.apply(qc_reason, axis=1)
good = obs[obs["qc"] == "ok"].copy()
section("quality control")
print(pd.crosstab(obs["qc"], obs["reservoir"])[RESERVOIRS].to_string())
print(f"\n{len(good)} usable reservoir-dates on {good['date'].nunique()} dates, "
      f"{good['date'].min():%d %b %Y} to {good['date'].max():%d %b %Y}")
section("why MNDWI fails on peaty water")
print(good.groupby("reservoir")[["rho_B03", "rho_B11", "mndwi", "scl_water_frac", "mndwi_pos_frac"]]
      .median().loc[RESERVOIRS].round(4).to_string())

# --- Step 6: anomalies: robust z-score against each reservoir's own month-of-year -----------------
good["month"] = good["date"].dt.month
grp = good.groupby(["reservoir", "month"])["ndci"]
good["clim_med"] = grp.transform("median")
good["clim_mad"] = grp.transform(lambda x: (x - x.median()).abs().median())
good["clim_n"] = grp.transform("size")
# Months with fewer than five observations fall back to the reservoir's whole-record statistics.
whole = good.groupby("reservoir")["ndci"]
few = good["clim_n"] < 5
good.loc[few, "clim_med"] = whole.transform("median")[few]
good.loc[few, "clim_mad"] = whole.transform(lambda x: (x - x.median()).abs().median())[few]
good["ndci_z"] = (good["ndci"] - good["clim_med"]) / (1.4826 * good["clim_mad"])
# If most reservoirs spike on the same day, suspect the atmosphere (haze, aerosol model) before the water.
high = good[good["ndci_z"] > 2]
good["n_high_same_day"] = good["date"].map(high.groupby("date").size()).fillna(0).astype(int)
good["flag"] = np.where(good["n_high_same_day"] >= 3, "check atmosphere (≥3 reservoirs high)",
                        np.where(good["ndci_z"] > 2, "investigate", ""))
section("dates most worth investigating (NDCI robust z > 2)")
cols = ["date", "reservoir", "ndci", "ndci_z", "ndti", "green_blue", "sun_zenith", "flag"]
top = good.sort_values("ndci_z", ascending=False).head(12)[cols]
print(top.assign(date=top["date"].dt.strftime("%Y-%m-%d")).round(3).to_string(index=False))

# --- Step 7: one tidy CSV per reservoir ---------------------------------------------------------
keep = ["date", "reservoir", "item_id", "orbit", "baseline", "sun_zenith", "clear_frac", "water_frac", "n_pixels",
        *BANDS_RHO, "ndci", "mci", "ndti", "red", "green_blue", "mndwi", "qc"]
tidy = obs[keep].merge(good[["item_id", "reservoir", "ndci_z", "flag"]], on=["item_id", "reservoir"], how="left")
section("tidy CSV exports")
for name in RESERVOIRS:
    f = OUT / f"{name.lower()}_sentinel2_water_quality.csv"
    t = tidy[tidy["reservoir"] == name]
    t.round(5).to_csv(f, index=False, date_format="%Y-%m-%d")
    print(f"{f.name:<42} {len(t):>4} rows, {int((t['qc'] == 'ok').sum()):>3} usable")

# --- Step 8: seasonal summary --------------------------------------------------------------------
SEASON = {3: "spring", 4: "spring", 5: "spring", 6: "summer", 7: "summer", 8: "summer",
          9: "autumn", 10: "autumn", 11: "autumn", 12: "winter", 1: "winter", 2: "winter"}
good["season"] = good["month"].map(SEASON)
section("median NDCI by season (n in brackets)")
med = good.pivot_table(index="reservoir", columns="season", values="ndci", aggfunc="median")
cnt = good.pivot_table(index="reservoir", columns="season", values="ndci", aggfunc="size")
seasons = [s for s in ["spring", "summer", "autumn"] if s in med]
print(pd.DataFrame({s: [f"{med.loc[r, s]:.3f} ({cnt.loc[r, s]:.0f})" for r in RESERVOIRS] for s in seasons},
                   index=RESERVOIRS).to_string())
section("whole-record medians")
summary = good.groupby("reservoir")[["ndci", "mci", "ndti", "red", "green_blue"]].median().loc[RESERVOIRS]
summary["n"] = good.groupby("reservoir").size()
print(summary.round(4).to_string())

# --- Step 9: maps for one clear summer date -------------------------------------------------------
summer = good[good["month"].between(6, 9)]
full = summer.groupby("date").agg(n=("reservoir", "size"), water=("water_frac", "min"), swir=("rho_B11", "mean"))
MAP_DATE = full[full["n"] == len(RESERVOIRS)].sort_values(["water", "swir"], ascending=[False, True]).index[0]
map_item = next(it for it in items if it.datetime.strftime("%Y-%m-%d") == MAP_DATE.strftime("%Y-%m-%d"))
section("map date")
print(f"{MAP_DATE:%d %B %Y}  {map_item.id}  sun zenith {map_item.properties['s2:mean_solar_zenith']:.1f}°")

scl = read_band(map_item, "SCL")
bad = ndimage.binary_dilation(np.isin(scl, BAD_SCL[1:]), iterations=CLOUD_BUFFER_PX) | (scl == 0)
r = {b: to_reflectance(read_band(map_item, b), map_item) for b in BANDS}
water = (labels > 0) & ~bad & (r["B08"] < NIR_MAX)


def smooth(a: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """3 x 3 mean over water pixels only (60 m), to tame pixel noise in very dark water."""
    v = ndimage.uniform_filter(np.where(mask, a, 0), 3)
    w = ndimage.uniform_filter(mask.astype("float32"), 3)
    with np.errstate(divide="ignore", invalid="ignore"):
        return np.where(mask, v / w, np.nan)


s = {b: smooth(r[b], water) for b in ("B03", "B04", "B05")}
with np.errstate(divide="ignore", invalid="ignore"):
    ndci_map = (s["B05"] - s["B04"]) / (s["B05"] + s["B04"])
    ndti_map = (s["B04"] - s["B03"]) / (s["B04"] + s["B03"])
rgb10 = np.dstack([to_reflectance(read_band(map_item, b, res=10), map_item) for b in ("B04", "B03", "B02")])
rgb20 = np.dstack([r[b] for b in ("B04", "B03", "B02")])
for name in RESERVOIRS:
    m = labels == RESERVOIRS.index(name) + 1
    print(f"{name:<10} NDCI p10/median/p90 {np.nanpercentile(ndci_map[m], 10):6.3f} {np.nanmedian(ndci_map[m]):6.3f} "
          f"{np.nanpercentile(ndci_map[m], 90):6.3f} | NDTI median {np.nanmedian(ndti_map[m]):6.3f}")


def stretch(rgb: np.ndarray, lo=1, hi=99, gamma=0.8) -> np.ndarray:
    a, b = np.nanpercentile(rgb, [lo, hi])
    return np.clip((rgb - a) / (b - a), 0, 1) ** gamma


EXTENT = (x0, x1, y0, y1)
HALO = {"white": [pe.withStroke(linewidth=2, foreground="black")], "black": [pe.withStroke(linewidth=2, foreground="white")]}
GROUPS = {"Upper Derwent Valley": ["Howden", "Derwent", "Ladybower"], "Longdendale": ["Torside"]}


def group_window(names, pad=500):
    gx0, gy0, gx1, gy1 = res[res["reservoir"].isin(names)].total_bounds
    return (max(gx0 - pad, x0), min(gx1 + pad, x1)), (max(gy0 - pad, y0), min(gy1 + pad, y1))


def map_panels(title: str, layer=None, cmap=None, vmin=None, vmax=None, label=None, fname=None):
    wins = [group_window(n) for n in GROUPS.values()]
    ratios = [(xl[1] - xl[0]) / (yl[1] - yl[0]) for xl, yl in wins]
    fig, axes = plt.subplots(1, 2, figsize=(10, 5.4), gridspec_kw={"width_ratios": ratios}, layout="constrained")
    for ax, (gname, names) in zip(axes, GROUPS.items()):
        if layer is None:
            ax.imshow(stretch(rgb10), extent=EXTENT)
        else:
            grey = np.nanmean(stretch(rgb20), axis=2)
            ax.imshow(grey, cmap="gray", extent=EXTENT, vmin=0, vmax=1.4)
            im = ax.imshow(layer, cmap=cmap, vmin=vmin, vmax=vmax, extent=EXTENT, interpolation="nearest")
        sub = res[res["reservoir"].isin(names)]
        sub.boundary.plot(ax=ax, color="white" if layer is None else "#333333", lw=0.8)
        if layer is None:
            gpd.GeoSeries(sub["core"], crs=CRS).boundary.plot(ax=ax, color="#ffd43b", lw=0.6, ls="--")
        for _, row in sub.iterrows():
            c = row.geometry.representative_point()
            ax.annotate(row["reservoir"], (c.x, c.y), xytext=(14, 0), textcoords="offset points", fontsize=8,
                        color="black", bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.85))
        xl, yl = group_window(names)
        ax.set_xlim(xl), ax.set_ylim(yl)
        ax.set_title(gname)
        ax.set_xticks([]), ax.set_yticks([]), ax.set_xlabel(""), ax.set_ylabel("")
        ink = "white" if layer is None else "black"
        h = (yl[1] - yl[0]) / 100
        ax.add_patch(Rectangle((xl[0] + 200, yl[0] + 3 * h), 1000, h, color=ink))
        ax.text(xl[0] + 700, yl[0] + 5 * h, "1 km", ha="center", va="bottom", fontsize=7, color=ink, path_effects=HALO[ink])
    if layer is not None:
        fig.colorbar(im, ax=axes, shrink=0.75, label=label, location="bottom", aspect=40)
    fig.suptitle(f"{title}, {MAP_DATE:%d %B %Y}")
    fig.savefig(OUT / fname)
    plt.close(fig)


# Figure 1: locator map with a Great Britain inset.
fig, ax = plt.subplots(figsize=(8, 7))
ax.imshow(stretch(rgb20), extent=EXTENT)
res.boundary.plot(ax=ax, color="white", lw=0.9)
for _, row in res.iterrows():
    c = row.geometry.centroid
    ax.annotate(f"{row['reservoir']}\n{row['area_ha']:.0f} ha", (c.x, c.y), xytext=(18, -4), textcoords="offset points",
                fontsize=8, color=COLOURS[row["reservoir"]], fontweight="bold",
                bbox=dict(boxstyle="round,pad=0.25", fc="white", ec=COLOURS[row["reservoir"]], lw=0.8))
ax.set_xticks([]), ax.set_yticks([]), ax.set_xlabel(""), ax.set_ylabel("")
ax.set_title(f"Peak District study reservoirs on Sentinel-2 true colour, {MAP_DATE:%d %B %Y}")
ax.add_patch(Rectangle((x0 + 500, y0 + 500), 2000, 120, color="white"))
ax.text(x0 + 1500, y0 + 700, "2 km", ha="center", va="bottom", color="white", fontsize=8, path_effects=HALO["white"])
try:
    ne = DATA / "ne_50m_admin_0_countries.gpkg"
    if not ne.exists():
        gpd.read_file("https://naciscdn.org/naturalearth/50m/cultural/ne_50m_admin_0_countries.zip").to_file(ne)
    uk = gpd.read_file(ne)
    uk = uk[uk["ADM0_A3"].isin(["GBR", "IRL", "IMN"])].to_crs(27700)
    inset = ax.inset_axes([0.76, 0.62, 0.24, 0.38])
    uk.plot(ax=inset, color="#e9ecef", edgecolor="#868e96", lw=0.4)
    p = gpd.GeoSeries([box(*BOUNDS)], crs=CRS).to_crs(27700).centroid.iloc[0]
    inset.plot(p.x, p.y, "o", ms=6, mfc="#e8590c", mec="white")
    inset.set_xlim(-50_000, 700_000), inset.set_ylim(0, 1_000_000)
    inset.set_xticks([]), inset.set_yticks([]), inset.set_xlabel(""), inset.set_ylabel("")
    inset.set_facecolor("#d0ebff")
except Exception as exc:   # the inset is decoration: never fail the tutorial over it
    print(f"(GB inset skipped: {exc})")
fig.savefig(OUT / "01-locator-map.png")
plt.close(fig)

map_panels("Sentinel-2 true colour (B4, B3, B2 at 10 m); dashed = 40 m inset core", fname="02-true-colour.png")
lo, hi = np.nanpercentile(ndci_map[labels > 0], [2, 98])
map_panels("NDCI (chlorophyll-a proxy), 60 m smoothed", ndci_map, "Spectral_r", lo, hi,
           "NDCI = (B5 − B4) / (B5 + B4), relative", "03-ndci-map.png")
lo, hi = np.nanpercentile(ndti_map[labels > 0], [2, 98])
map_panels("NDTI (turbidity proxy), 60 m smoothed", ndti_map, "BrBG_r", lo, hi,
           "NDTI = (B4 − B3) / (B4 + B3), relative", "04-turbidity-map.png")

# Figure 5: multi-year NDCI time series.
fig, axes = plt.subplots(len(RESERVOIRS), 1, figsize=(10, 7.5), sharex=True, sharey=True)
for ax, name in zip(axes, RESERVOIRS):
    g = good[good["reservoir"] == name].sort_values("date")
    for yr in range(2019, 2027):
        ax.axvspan(pd.Timestamp(f"{yr}-06-01"), pd.Timestamp(f"{yr}-09-01"), color="#fff3bf", lw=0)
    ax.plot(g["date"], g["ndci"], "-", color=COLOURS[name], lw=0.6, alpha=0.5)
    ax.plot(g["date"], g["ndci"], "o", color=COLOURS[name], ms=3)
    hi_ = g[g["ndci_z"] > 2]
    ax.plot(hi_["date"], hi_["ndci"], "o", ms=8, mfc="none", mec="black", mew=0.9)
    ax.axhline(g["ndci"].median(), color=COLOURS[name], ls="--", lw=0.8)
    ax.text(0.005, 0.9, f"{name} (n = {len(g)})", transform=ax.transAxes, fontweight="bold", color=COLOURS[name],
            va="top")
    ax.grid(alpha=0.3)
axes[0].set_ylim(good["ndci"].min() - 0.02, good["ndci"].max() + 0.02)
axes[len(RESERVOIRS) // 2].set_ylabel("NDCI (relative)")
axes[-1].xaxis.set_major_locator(mdates.YearLocator())
axes[-1].xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
axes[0].set_title("Sentinel-2 NDCI per reservoir (median of core water pixels); yellow = June–August, "
                  "circled = robust z > 2")
fig.savefig(OUT / "05-ndci-timeseries.png")
plt.close(fig)

# Figure 6: monthly climatology of three indices.
fig, axes = plt.subplots(1, 3, figsize=(11, 3.8))
specs = [("ndci", "NDCI (chlorophyll proxy)"), ("ndti", "NDTI (turbidity proxy)"),
         ("green_blue", "B3 / B2 (colour proxy)")]
months = range(3, 10)
for ax, (col, lab) in zip(axes, specs):
    for name in RESERVOIRS:
        g = good[(good["reservoir"] == name) & good["month"].isin(months)].groupby("month")[col]
        q = g.quantile([0.25, 0.5, 0.75]).unstack()
        q = q[g.size() >= 3]
        ax.plot(q.index, q[0.5], "-o", color=COLOURS[name], ms=3, label=name)
        ax.fill_between(q.index, q[0.25], q[0.75], color=COLOURS[name], alpha=0.12, lw=0)
    ax.set_xticks(list(months), ["Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep"])
    ax.set_title(lab)
    ax.grid(alpha=0.3)
axes[0].legend(fontsize=8, frameon=False)
fig.suptitle("Monthly median and interquartile range, 2019–2026 (months with at least 3 usable dates)", y=1.02)
fig.tight_layout()
fig.savefig(OUT / "06-monthly-climatology.png")
plt.close(fig)

# --- Step 10: sanity check against public Environment Agency samples ------------------------------
EA = "https://environment.data.gov.uk/water-quality"
POINTS = {"MD-49714900": "Howden", "MD-49714220": "Derwent", "MD-LSN0395": "Ladybower", "NW-88021743": "Torside"}
DETERMINANDS = {"7887": "chlorophyll_ugl", "0072": "colour_hazen"}
ea_file = DATA / "ea_samples.csv"
if not ea_file.exists():
    frames = []
    for point, name in POINTS.items():
        for code in DETERMINANDS:
            skip = 0
            while True:
                resp = requests.get(f"{EA}/sampling-point/{point}/observation", headers={"Accept": "text/csv"},
                                    params={"dateFrom": "2019-01-01", "dateTo": "2026-10-14", "determinand": code,
                                            "limit": 250, "skip": skip}, timeout=60)
                resp.raise_for_status()
                part = pd.read_csv(io.StringIO(resp.text), dtype=str) if resp.text.strip() else pd.DataFrame()
                if part.empty:
                    break
                frames.append(part.assign(point=point, reservoir=name))
                if len(part) < 250:
                    break
                skip += 250
    pd.concat(frames, ignore_index=True).to_csv(ea_file, index=False)
ea = pd.read_csv(ea_file, dtype=str)
ea["date"] = pd.to_datetime(ea["phenomenonTime"].str[:10])
ea["censored"] = ea["result"].str.startswith("<")
ea["value"] = pd.to_numeric(ea["result"].str.lstrip("<>"), errors="coerce")
ea["var"] = ea["determinand.notation"].str.zfill(4).map(DETERMINANDS)
ea = ea.dropna(subset=["var", "value"])
section("Environment Agency samples (2019 onwards)")
print(ea.groupby(["reservoir", "var"])["date"].agg(["size", "min", "max"]).rename(columns={"size": "n"})
      .reset_index().assign(min=lambda d: d["min"].dt.strftime("%Y-%m-%d"),
                            max=lambda d: d["max"].dt.strftime("%Y-%m-%d")).to_string(index=False))

MATCH_DAYS = 3
pairs = []
for _, s_ in ea.iterrows():
    g = good[good["reservoir"] == s_["reservoir"]]
    dt = (g["date"] - s_["date"]).abs()
    if len(g) and dt.min() <= pd.Timedelta(days=MATCH_DAYS):
        m = g.loc[dt.idxmin()]
        pairs.append({"reservoir": s_["reservoir"], "var": s_["var"], "sample_date": s_["date"], "s2_date": m["date"],
                      "value": s_["value"], "censored": s_["censored"], "ndci": m["ndci"], "mci": m["mci"],
                      "green_blue": m["green_blue"], "ndti": m["ndti"]})
pairs = pd.DataFrame(pairs)
section(f"satellite vs sample matches (within ±{MATCH_DAYS} days)")
chl, col = pairs[pairs["var"] == "chlorophyll_ugl"], pairs[pairs["var"] == "colour_hazen"]
show = chl.assign(sample_date=chl["sample_date"].dt.strftime("%Y-%m-%d"), s2_date=chl["s2_date"].dt.strftime("%Y-%m-%d"))
print(show[["reservoir", "sample_date", "s2_date", "value", "censored", "ndci", "mci"]].round(4).to_string(index=False))
for lab, d, x in [("chlorophyll vs NDCI", chl, "ndci"), ("chlorophyll vs MCI", chl, "mci"),
                  ("colour vs B3/B2", col, "green_blue")]:
    if len(d) >= 4:
        rho, pval = spearmanr(d["value"], d[x])
        print(f"{lab:<22} n = {len(d):>2}  Spearman rho = {rho:5.2f}  p = {pval:.2f}")
    else:
        print(f"{lab:<22} n = {len(d):>2}  too few matches for a correlation")

# Figure 7: the in-situ sanity check.
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
for ax, (d, x, xl, yl) in zip(axes, [(chl, "ndci", "Sentinel-2 NDCI", "Chlorophyll-a, acetone extract (µg/L)"),
                                      (col, "green_blue", "Sentinel-2 B3 / B2", "Colour, filtered (Hazen)")]):
    for name in RESERVOIRS:
        g = d[d["reservoir"] == name]
        if len(g):
            ax.scatter(g[x], g["value"], color=COLOURS[name], s=28, label=f"{name} ({len(g)})",
                       edgecolor="white", lw=0.5)
            cz = g[g["censored"]]
            ax.scatter(cz[x], cz["value"], s=60, facecolor="none", edgecolor="black", lw=0.7)
    if len(d) >= 4:
        rho, _ = spearmanr(d["value"], d[x])
        ax.set_title(f"Spearman ρ = {rho:.2f}, n = {len(d)}")
    ax.set_xlabel(xl), ax.set_ylabel(yl)
    ax.grid(alpha=0.3)
    ax.legend(fontsize=7, frameon=False, loc="upper left", bbox_to_anchor=(1.0, 1.0))
fig.suptitle(f"Sentinel-2 vs Environment Agency samples within ±{MATCH_DAYS} days (open circle = below detection limit)",
             y=1.02)
fig.tight_layout()
fig.savefig(OUT / "07-in-situ-check.png")
plt.close(fig)

section("figures written")
for f in sorted(OUT.glob("*.png")):
    print(f.name)