"""
Build, validate and query a static STAC catalogue in Python with pystac and rio-stac, then compare it
with the public Earth Search and Microsoft Planetary Computer STAC APIs through pystac-client.

Companion script for https://technicalgis.com/tutorials/stac-catalogue-pystac-python/
(part 2 of the series "Cloud-native rasters in Python").

Input: the 25 Cloud Optimised GeoTIFF tiles of Eryri (Snowdonia) made by part 1, cog_mosaic.py.
If ../cloud-optimised-geotiff-python/output/tiles/ is missing, this script runs part 1 first.

Data: Copernicus DEM GLO-30 (c) DLR e.V. 2010-2014 and (c) Airbus Defence and Space GmbH 2014-2018,
provided under COPERNICUS by the European Union and ESA; all rights reserved. Sentinel-2 L2A from the
Earth Search catalogue (Copernicus Sentinel data 2026, processed by ESA).

Run:  python stac_catalogue.py
Outputs (in output/): stac/ (the catalogue), six figures and this console log.
"""
from __future__ import annotations

import http.server
import json
import logging
import re
import shutil
import subprocess
import sys
import threading
import time
import warnings
from datetime import datetime, timezone
from functools import partial
from pathlib import Path

import jsonschema
import matplotlib
import numpy as np
import pandas as pd
import geopandas as gpd
import pystac
import rasterio
import requests
from matplotlib.colors import LightSource, ListedColormap
from matplotlib.patches import FancyBboxPatch, Rectangle
from pystac.extensions.projection import ProjectionExtension
from pystac_client import Client
from rasterio.merge import merge
from rasterio.vrt import WarpedVRT
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rio_stac import create_stac_item
from shapely.geometry import Point, box, shape

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

HERE = Path(__file__).resolve().parent
OUT = HERE / "output"
STAC_DIR = OUT / "stac"
PART1 = HERE.parent / "cloud-optimised-geotiff-python"
TILE_DIR = PART1 / "output" / "tiles"
OUT.mkdir(exist_ok=True)

EARTH_SEARCH = "https://earth-search.aws.element84.com/v1"
PLANETARY_COMPUTER = "https://planetarycomputer.microsoft.com/api/stac/v1"
SNOWDON_LONLAT = (-4.0763, 53.0685)           # Yr Wyddfa summit
GDAL_ENV = dict(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".tif")

plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "axes.facecolor": "white"})
pd.set_option("display.width", 140)
pd.set_option("display.max_columns", 20)
LIGHT = LightSource(azdeg=315, altdeg=45)
DEM_CMAP = ListedColormap(plt.cm.terrain(np.linspace(0.25, 1.0, 256)))
SEA = ListedColormap(["#9ecae1"])
ASPECT = 1 / np.cos(np.radians(53.07))


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


def kb(n: float) -> float:
    return n / 1e3


# GDAL logs every HTTP range request it makes; count them as in part 1.
RANGES: list[tuple[int, int]] = []


class RangeLog(logging.Handler):
    def emit(self, record):
        for a, b in re.findall(r"Downloading (\d+)-(\d+)", record.getMessage()):
            RANGES.append((int(a), int(b)))


logging.getLogger("rasterio").addHandler(RangeLog())


# ---------------------------------------------------------------- Step 0: make sure part 1's tiles exist
section("Step 0: find the COG tiles from part 1")
if len(list(TILE_DIR.glob("SH*.tif"))) < 25:
    print("part 1 outputs not found: running cog_mosaic.py (a few minutes)")
    subprocess.run([sys.executable, "cog_mosaic.py"], cwd=PART1, check=True, stdout=subprocess.DEVNULL)
tiles = sorted(TILE_DIR.glob("SH*.tif"))
print(f"{len(tiles)} COG tiles in {TILE_DIR.relative_to(HERE.parent)}, {sum(t.stat().st_size for t in tiles) / 1e6:.2f} MB")


# ---------------------------------------------------------------- Step 1: build the catalogue
section("Step 1: build a Catalog -> Collection -> Items tree with pystac and rio-stac")
if STAC_DIR.exists():
    shutil.rmtree(STAC_DIR)
ACQUIRED = (datetime(2011, 1, 1, tzinfo=timezone.utc), datetime(2015, 12, 31, 23, 59, 59, tzinfo=timezone.utc))
PROCESSED = datetime(2026, 11, 3, tzinfo=timezone.utc)


def source_glo30(bounds) -> list[str]:
    """Earth Search IDs of the 1 x 1 degree GLO-30 tiles that a footprint touches."""
    ids = []
    for lat in range(int(np.floor(bounds[1])), int(np.floor(bounds[3])) + 1):
        for lon in range(int(np.floor(bounds[0])), int(np.floor(bounds[2] - 1e-9)) + 1):
            ids.append(f"Copernicus_DSM_COG_10_N{lat:02d}_00_W{-lon:03d}_00_DEM")
    return ids


t0 = time.perf_counter()
items = []
for path in tiles:
    name = path.stem
    item = create_stac_item(
        str(path),
        id=f"eryri-dem-25m-{name}",
        input_datetime=ACQUIRED[0],
        asset_name="dem",
        asset_roles=["data"],
        asset_media_type=pystac.MediaType.COG,
        with_proj=True,
        with_raster=True,
        geom_densify_pts=10,
        geom_precision=6,
        properties={"os:grid_square": name, "processing:software": {"rasterio": rasterio.__version__}},
    )
    # A DEM built from 2011-2015 acquisitions has a date range, not an instant.
    item.datetime = None
    item.common_metadata.start_datetime, item.common_metadata.end_datetime = ACQUIRED
    item.common_metadata.created = PROCESSED
    item.common_metadata.title = f"Eryri DEM 25 m, OS square {name}"
    # rio-stac 0.12 writes the projection extension v1.1 (proj:epsg); upgrade to v2.0 (proj:code), as Earth Search uses.
    item.properties["proj:code"] = f"EPSG:{item.properties.pop('proj:epsg')}"
    item.stac_extensions = [u.replace("projection/v1.1.0", "projection/v2.0.0") for u in item.stac_extensions]
    with rasterio.open(path) as src:
        full = src.read(1)
        small = src.read(1, out_shape=(100, 100), resampling=Resampling.average)
    # rio-stac 0.12 reports valid_percent as a fraction divided by the pixel count; store the true percentage.
    item.assets["dem"].extra_fields["raster:bands"][0]["statistics"]["valid_percent"] = \
        round(100 * float(np.isfinite(full).mean()), 2)
    # Thumbnail asset: a 100 x 100 px hillshaded preview.
    thumb = OUT / "thumbs" / f"{name}.png"
    thumb.parent.mkdir(exist_ok=True)
    rgb = LIGHT.shade(np.where(small > 0, small, 0), cmap=DEM_CMAP, vmin=0, vmax=1100, blend_mode="soft",
                      vert_exag=1.5, dx=100, dy=100)
    rgb[small <= 0, :3] = (0.62, 0.79, 0.88)
    plt.imsave(thumb, rgb)
    item.add_asset("thumbnail", pystac.Asset(href=str(thumb), media_type=pystac.MediaType.PNG,
                                             roles=["thumbnail"], title="Hillshaded preview, 100 x 100 px"))
    for src_id in source_glo30(item.bbox):
        item.add_link(pystac.Link("derived_from", f"{EARTH_SEARCH}/collections/cop-dem-glo-30/items/{src_id}",
                                  media_type=pystac.MediaType.GEOJSON, title=f"Copernicus GLO-30 {src_id[22:36]}"))
    items.append(item)
build_s = time.perf_counter() - t0

collection = pystac.Collection(
    id="eryri-dem-25m",
    title="Eryri DEM, 25 m, British National Grid",
    description=("Copernicus DEM GLO-30 over a 50 x 50 km area of Eryri (Snowdonia), mosaicked, reprojected to "
                 "EPSG:27700 at 25 m and cut into 25 Cloud Optimised GeoTIFFs on the Ordnance Survey 10 km grid."),
    extent=pystac.Extent(
        pystac.SpatialExtent([[min(i.bbox[0] for i in items), min(i.bbox[1] for i in items),
                               max(i.bbox[2] for i in items), max(i.bbox[3] for i in items)]]),
        pystac.TemporalExtent([[ACQUIRED[0], ACQUIRED[1]]]),
    ),
    license="other",
    keywords=["DEM", "Copernicus", "Eryri", "Snowdonia", "COG"],
    providers=[
        pystac.Provider("European Space Agency / Airbus", roles=[pystac.ProviderRole.LICENSOR,
                                                                 pystac.ProviderRole.PRODUCER],
                        url="https://spacedata.copernicus.eu/collections/copernicus-digital-elevation-model"),
        pystac.Provider("TechnicalGIS", roles=[pystac.ProviderRole.PROCESSOR, pystac.ProviderRole.HOST],
                        url="https://technicalgis.com/tutorials/cloud-optimised-geotiff-python/"),
    ],
)
collection.add_link(pystac.Link("license", "https://spacedata.copernicus.eu/documents/20123/121286/CSCDA_ESA_Mission-specific+Annex_31_Oct_22.pdf",
                                media_type="application/pdf", title="Copernicus DEM licence"))
collection.add_items(items)
collection.summaries.add("os:grid_square", sorted(i.properties["os:grid_square"] for i in items))
collection.summaries.add("proj:code", ["EPSG:27700"])
collection.summaries.add("gsd", [25])

catalog = pystac.Catalog(id="technicalgis-eryri", title="TechnicalGIS Eryri example catalogue",
                         description="A static STAC catalogue built for the TechnicalGIS cloud-native rasters series.")
catalog.add_child(collection)
catalog.normalize_hrefs(str(STAC_DIR))

# Copy each COG and thumbnail next to its item JSON so the whole tree can be moved or uploaded as one folder.
for item in collection.get_items():
    item_dir = Path(item.get_self_href()).parent
    item_dir.mkdir(parents=True, exist_ok=True)
    for key, asset in item.assets.items():
        dst = item_dir / Path(asset.href).name
        shutil.copy2(asset.href, dst)
        asset.href = str(dst)
catalog.make_all_asset_hrefs_relative()
t0 = time.perf_counter()
catalog.save(catalog_type=pystac.CatalogType.SELF_CONTAINED)
save_s = time.perf_counter() - t0

json_files = sorted(STAC_DIR.rglob("*.json"))
data_files = [p for p in STAC_DIR.rglob("*") if p.is_file() and p.suffix != ".json"]
print(f"pystac {pystac.__version__}, STAC version {pystac.get_stac_version()}; built 25 items in {build_s:.2f} s, "
      f"saved in {save_s:.2f} s")
print(f"{len(json_files)} JSON files ({sum(p.stat().st_size for p in json_files) / 1e3:.1f} kB) and "
      f"{len(data_files)} asset files ({sum(p.stat().st_size for p in data_files) / 1e6:.2f} MB)")
for p in json_files[:3] + [json_files[-1]]:
    print(f"  {p.relative_to(OUT).as_posix():48s} {p.stat().st_size:6,d} bytes")
sh65 = json.loads((STAC_DIR / "eryri-dem-25m" / "eryri-dem-25m-SH65" / "eryri-dem-25m-SH65.json").read_text())
print("\nSH65 item, abridged:")
print(json.dumps({
    "stac_extensions": sh65["stac_extensions"],
    "bbox": sh65["bbox"],
    "properties": {k: v for k, v in sh65["properties"].items() if k not in ("proj:geometry", "proj:projjson", "proj:wkt2")},
    "assets.dem": {k: v for k, v in sh65["assets"]["dem"].items() if k not in ("raster:bands", "bands")},
    "dem band statistics": (sh65["assets"]["dem"].get("raster:bands") or sh65["assets"]["dem"]["bands"])[0]["statistics"],
    "links (rel)": [l["rel"] for l in sh65["links"]],
}, indent=None, default=str).replace(", \"", ",\n  \"")[:2500])


# ---------------------------------------------------------------- Step 2: validate
section("Step 2: validate every object against the STAC JSON schemas")
t0 = time.perf_counter()
n_valid = catalog.validate_all()
first_s = time.perf_counter() - t0
t0 = time.perf_counter()
catalog.validate_all()
second_s = time.perf_counter() - t0
print(f"validate_all: {n_valid} items valid (plus the catalogue and collection); "
      f"first pass {first_s:.2f} s (fetches the schemas), second pass {second_s:.2f} s (schemas cached)")

no_start = json.loads(json.dumps(sh65))
no_start["properties"].pop("start_datetime")       # datetime is null, so start and end become mandatory
bad_shape = json.loads(json.dumps(sh65))
bad_shape["properties"]["proj:shape"] = "400 x 400"  # must be an array of integers
for label, doc in [("start_datetime removed", no_start), ("proj:shape as a string", bad_shape)]:
    try:
        pystac.validation.validate_dict(doc)
        print(f"  {label}: passed")
    except pystac.STACValidationError as err:
        best = jsonschema.exceptions.best_match(err.source)
        print(f"  {label}: STACValidationError")
        print(f"    schema: {str(err).splitlines()[0].split('schema at ')[-1]}")
        print(f"    reason: {best.message[:110]} (at {'/'.join(map(str, best.absolute_path)) or 'item root'})")


# ---------------------------------------------------------------- Step 3: read back and query locally
section("Step 3: read the catalogue back and query it with GeoPandas")
t0 = time.perf_counter()
cat = pystac.Catalog.from_file(str(STAC_DIR / "catalog.json"))
local_items = list(cat.get_items(recursive=True))
read_s = time.perf_counter() - t0
rows = []
for it in local_items:
    band = (it.assets["dem"].extra_fields.get("raster:bands") or it.assets["dem"].extra_fields["bands"])[0]
    rows.append({"id": it.id, "square": it.properties["os:grid_square"], "min_m": band["statistics"]["minimum"],
                 "max_m": band["statistics"]["maximum"], "mean_m": band["statistics"]["mean"],
                 "href": it.assets["dem"].get_absolute_href(), "thumb": it.assets["thumbnail"].get_absolute_href(),
                 "geometry": shape(it.geometry)})
gdf = gpd.GeoDataFrame(rows, crs="EPSG:4326")
proj = ProjectionExtension.ext(local_items[0])
print(f"read {len(local_items)} items from disk in {read_s:.3f} s; projection extension on {local_items[0].id}: "
      f"code {proj.code}, shape {proj.shape}, bbox {proj.bbox}")

snowdon = Point(*SNOWDON_LONLAT)
q_point = gdf[gdf.intersects(snowdon)]
q_bbox_geom = box(-3.95, 53.10, -3.75, 53.25)          # Conwy valley
q_bbox = gdf[gdf.intersects(q_bbox_geom)]
q_high = gdf[gdf.max_m > 800].sort_values("max_m", ascending=False)
print(f"point query (Yr Wyddfa {SNOWDON_LONLAT}): {q_point.square.tolist()}")
print(f"bbox query (-3.95, 53.10, -3.75, 53.25): {len(q_bbox)} items {q_bbox.square.tolist()}")
print(f"property query (maximum elevation > 800 m): {len(q_high)} items")
print(q_high[["square", "min_m", "mean_m", "max_m"]].round(1).to_string(index=False))


# ---------------------------------------------------------------- Step 4: serve it as static files
section("Step 4: serve the catalogue as static files and crawl it over HTTP")
HITS: list[tuple[str, int]] = []


class CountingHandler(http.server.SimpleHTTPRequestHandler):
    """Static file server that honours single HTTP Range requests and records every GET and its bytes."""

    def log_message(self, *args):
        pass

    def send_head(self):
        path = self.translate_path(self.path.split("?")[0])
        if not Path(path).is_file():
            self.send_error(404)
            return None
        size = Path(path).stat().st_size
        rng = re.match(r"bytes=(\d+)-(\d*)", self.headers.get("Range", ""))
        f = open(path, "rb")
        start, end = 0, size - 1
        if rng:
            start = int(rng.group(1))
            end = min(int(rng.group(2)) if rng.group(2) else size - 1, size - 1)
            self.send_response(206)
            self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
            f.seek(start)
        else:
            self.send_response(200)
        self.send_header("Content-Type", self.guess_type(path))
        self.send_header("Accept-Ranges", "bytes")
        self.send_header("Content-Length", str(end - start + 1))
        self.end_headers()
        self._remaining = end - start + 1
        if self.command == "GET":
            HITS.append((self.path, self._remaining))
        return f

    def copyfile(self, source, outputfile):
        outputfile.write(source.read(self._remaining))


server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), partial(CountingHandler, directory=str(STAC_DIR)))
threading.Thread(target=server.serve_forever, daemon=True).start()
base = f"http://127.0.0.1:{server.server_address[1]}"

t0 = time.perf_counter()
remote_cat = pystac.Catalog.from_file(f"{base}/catalog.json")
crawled = list(remote_cat.get_items(recursive=True))
crawl_s = time.perf_counter() - t0
crawl_requests, crawl_bytes = len(HITS), sum(s for _, s in HITS)
hits = [h for h, _ in HITS]
print(f"crawl all items: {len(crawled)} items, {crawl_requests} GET requests, {crawl_bytes / 1e3:.1f} kB of JSON, "
      f"{crawl_s:.2f} s")
print(f"  first requests: {hits[:3]}")

HITS.clear()
t0 = time.perf_counter()
one = pystac.Item.from_file(f"{base}/eryri-dem-25m/eryri-dem-25m-SH65/eryri-dem-25m-SH65.json")
one_s = time.perf_counter() - t0
print(f"fetch one item you already know: {len(HITS)} request, {sum(s for _, s in HITS) / 1e3:.1f} kB, {one_s * 1000:.0f} ms")

try:
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")      # NoConformsTo: a static catalogue advertises no API conformance
        static_client = Client.open(f"{base}/catalog.json")
    static_client.search(bbox=[-4.1, 53.0, -4.0, 53.1]).item_collection()
    print("pystac-client search on the static catalogue: worked")
except Exception as err:  # noqa: BLE001
    print(f"pystac-client search on the static catalogue: {type(err).__name__}: {str(err).splitlines()[0][:100]}")

HITS.clear()
with rasterio.Env(CPL_DEBUG=True, **GDAL_ENV):
    logging.getLogger("rasterio").setLevel(logging.DEBUG)
    RANGES.clear()
    t0 = time.perf_counter()
    with rasterio.open(one.assets["dem"].get_absolute_href()) as src:
        win = src.read(1, window=rasterio.windows.from_bounds(260_000, 353_000, 262_000, 355_000, src.transform))
    cog_s = time.perf_counter() - t0
    logging.getLogger("rasterio").setLevel(logging.WARNING)
print(f"read a 2 x 2 km window of SH65 through its asset href: {len(RANGES)} range requests, "
      f"{sum(b - a + 1 for a, b in RANGES) / 1e3:.1f} kB of a "
      f"{(STAC_DIR / 'eryri-dem-25m' / 'eryri-dem-25m-SH65' / 'SH65.tif').stat().st_size / 1e3:.1f} kB file, "
      f"max {win.max():.1f} m, {cog_s * 1000:.0f} ms")
server.shutdown()


# ---------------------------------------------------------------- Step 5: the same questions to public STAC APIs
section("Step 5: search public STAC APIs with pystac-client")
API_CALLS: list[tuple[str, float, int]] = []


def count_requests(client: Client):
    session = client._stac_io.session

    def hook(resp, *args, **kwargs):
        API_CALLS.append((resp.request.method, resp.elapsed.total_seconds(), len(resp.content)))
    session.hooks["response"].append(hook)
    return client


aoi_bbox = [round(float(v), 4) for v in gdf.total_bounds]
searches = []

API_CALLS.clear()
t0 = time.perf_counter()
es = count_requests(Client.open(EARTH_SEARCH))
open_s = time.perf_counter() - t0
print(f"Earth Search: opened landing page in {open_s:.2f} s; {len(list(es.get_collections()))} collections")

API_CALLS.clear()
t0 = time.perf_counter()
dem_search = es.search(collections=["cop-dem-glo-30"], bbox=aoi_bbox)
es_dem = list(dem_search.items())
s = time.perf_counter() - t0
searches.append(("Earth Search\ncop-dem-glo-30", len(es_dem), len(API_CALLS), s, sum(c[2] for c in API_CALLS)))
print(f"cop-dem-glo-30 in bbox {aoi_bbox}: {len(es_dem)} items, {len(API_CALLS)} POST request(s), {s:.2f} s")
for it in es_dem:
    print(f"  {it.id}  datetime {it.datetime:%Y-%m-%d}  {it.assets['data'].href}")
our_sources = sorted({l.href.rsplit('/', 1)[-1] for i in local_items for l in i.get_links('derived_from')})
print(f"our derived_from links point at the same {len(our_sources)} items: "
      f"{sorted(our_sources) == sorted(i.id for i in es_dem)}")

API_CALLS.clear()
t0 = time.perf_counter()
s2_search = es.search(collections=["sentinel-2-l2a"], intersects=box(*aoi_bbox).__geo_interface__,
                      datetime="2026-05-01/2026-09-30", query={"eo:cloud_cover": {"lt": 10}},
                      sortby=[{"field": "properties.eo:cloud_cover", "direction": "asc"}], max_items=100)
s2 = list(s2_search.items())
s = time.perf_counter() - t0
searches.append(("Earth Search\nsentinel-2-l2a", len(s2), len(API_CALLS), s, sum(c[2] for c in API_CALLS)))
s2_df = pd.DataFrame([{"id": i.id, "date": i.datetime.date(), "mgrs": i.properties.get("grid:code", "")[5:],
                       "cloud_%": i.properties["eo:cloud_cover"], "geometry": shape(i.geometry)} for i in s2])
print(f"sentinel-2-l2a, May-Sep 2026, cloud < 10%: {len(s2)} items, {len(API_CALLS)} request(s), {s:.2f} s")
print(s2_df.drop(columns="geometry").sort_values(["date", "mgrs"]).to_string(index=False))

pc_items = []
try:
    API_CALLS.clear()
    t0 = time.perf_counter()
    pc = count_requests(Client.open(PLANETARY_COMPUTER))
    pc_items = list(pc.search(collections=["cop-dem-glo-30"], bbox=aoi_bbox).items())
    s = time.perf_counter() - t0
    searches.append(("Planetary Computer\ncop-dem-glo-30", len(pc_items), len(API_CALLS), s,
                     sum(c[2] for c in API_CALLS)))
    print(f"Planetary Computer cop-dem-glo-30: {len(pc_items)} items, {len(API_CALLS)} requests "
          f"(including the landing page), {s:.2f} s")
    print(f"  same item IDs as Earth Search: {sorted(i.id for i in pc_items) == sorted(i.id for i in es_dem)}")
    print(f"  example href: {pc_items[0].assets['data'].href.split('?')[0]}")
except Exception as err:  # noqa: BLE001
    print(f"Planetary Computer search failed: {type(err).__name__}: {err}")


# ---------------------------------------------------------------- Step 6: read data through the search results
section("Step 6: read pixels through asset hrefs from both catalogues")
t0 = time.perf_counter()
mosaic, mosaic_tf = merge(list(q_high.href), nodata=np.nan)   # NaN where no item covers the bounding box
mosaic = mosaic[0]
print(f"merged the {len(q_high)} items with maximum > 800 m from local hrefs: {mosaic.shape[1]} x {mosaic.shape[0]} px "
      f"in {time.perf_counter() - t0:.2f} s; max {np.nanmax(mosaic):.1f} m; {np.isnan(mosaic).mean() * 100:.0f}% of the box has no item")

best_date = s2_df.groupby("date").agg(n=("mgrs", "nunique"), cloud=("cloud_%", "mean")).query("n >= 2") \
    .sort_values("cloud").index[0]
# The same date can hold more than one processing of a tile (IDs ending _0 and _1); keep the least cloudy per tile.
best_df = s2_df[s2_df.date == best_date].sort_values("cloud_%").drop_duplicates("mgrs")
best = [i for i in s2 if i.id in set(best_df.id)]
print(f"clearest date covering both MGRS tiles: {best_date}; "
      f"{(s2_df.date == best_date).sum()} items, using {', '.join(i.id for i in best)}")
AOI_BNG = (250_000, 330_000, 300_000, 380_000)
S2_RES = 100
s2_shape = (int((AOI_BNG[3] - AOI_BNG[1]) / S2_RES), int((AOI_BNG[2] - AOI_BNG[0]) / S2_RES))
s2_tf = from_origin(AOI_BNG[0], AOI_BNG[3], S2_RES, S2_RES)
rgb = np.zeros((3, *s2_shape), dtype=np.uint8)
RANGES.clear()
logging.getLogger("rasterio").setLevel(logging.DEBUG)
t0 = time.perf_counter()
with rasterio.Env(CPL_DEBUG=True, **GDAL_ENV):
    for it in best:
        href = it.assets["visual"].href
        with rasterio.open(href) as src:
            ovr = src.overviews(1)
            size_mb = int(requests.head(href, timeout=60).headers["Content-Length"]) / 1e6
            print(f"  {it.id}: visual asset {src.width} x {src.height} px at 10 m, overviews {ovr}, {size_mb:.0f} MB")
        # Overview level 2 is the 1/8 overview (80 m), the closest one finer than our 100 m target.
        with rasterio.open(href, overview_level=2) as src, \
                WarpedVRT(src, crs="EPSG:27700", transform=s2_tf, width=s2_shape[1], height=s2_shape[0],
                          resampling=Resampling.average) as vrt:
            part = vrt.read()
            fill = (rgb.sum(axis=0) == 0) & (part.sum(axis=0) > 0)
            rgb[:, fill] = part[:, fill]
s2_s = time.perf_counter() - t0
logging.getLogger("rasterio").setLevel(logging.WARNING)
s2_bytes = sum(b - a + 1 for a, b in RANGES)
print(f"Sentinel-2 true colour over the 50 x 50 km area at {S2_RES} m: {len(RANGES)} range requests, "
      f"{s2_bytes / 1e6:.2f} MB, {s2_s:.1f} s")


# ---------------------------------------------------------------- Figures
section("figures")
gdf["cx"] = gdf.geometry.representative_point().x
gdf["cy"] = gdf.geometry.representative_point().y

# 01: the catalogue tree as it sits on disk
fig, ax = plt.subplots(figsize=(11, 6.2))
ax.set_xlim(0, 11)
ax.set_ylim(0, 6.2)
ax.axis("off")


def node(x, y, w, h, text, fc, fs=8.5):
    ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle="round,pad=0.04,rounding_size=0.08", fc=fc, ec="#333",
                                lw=1))
    ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=fs)


sizes = {p.parent.name.rsplit("-", 1)[-1]: p.stat().st_size for p in STAC_DIR.rglob("eryri-dem-25m-SH*.json")}
cat_size = (STAC_DIR / "catalog.json").stat().st_size
col_size = (STAC_DIR / "eryri-dem-25m" / "collection.json").stat().st_size
node(0.2, 4.9, 2.6, 0.95, f"Catalog\ntechnicalgis-eryri\ncatalog.json ({cat_size:,} B)", "#fee08b")
node(0.2, 3.25, 2.6, 1.15, f"Collection\neryri-dem-25m\ncollection.json ({col_size:,} B)\nextent, licence, providers",
     "#abdda4")
ax.annotate("", (1.5, 4.4), (1.5, 4.9), arrowprops=dict(arrowstyle="->", lw=1.2))
item_kb = np.mean(list(sizes.values())) / 1e3
node(0.2, 0.5, 2.6, 2.2, f"Each item folder holds\nitem JSON (~{item_kb:.1f} kB)\n+ dem COG asset\n+ thumbnail PNG asset\n"
     "\nLinks: root, parent,\ncollection, derived_from\n(the source GLO-30 items)", "#e6f5d0", fs=8)
ax.annotate("", (1.5, 2.7), (1.5, 3.25), arrowprops=dict(arrowstyle="->", lw=1.2))
ax.annotate("", (3.3, 3.82), (2.8, 3.82), arrowprops=dict(arrowstyle="->", lw=1.2))
for _, r in gdf.iterrows():
    e, n = int(r.square[2]), int(r.square[3])
    x, y = 3.4 + (e - 5) * 1.5, 0.4 + (n - 3) * 1.15
    col = DEM_CMAP(min(r.max_m, 1100) / 1100)
    node(x, y, 1.35, 0.95, f"Item {r.square}\nmax {r.max_m:.0f} m\n{sizes[r.square] / 1e3:.1f} kB JSON", col, fs=7.5)
ax.set_title(f"The static catalogue: 1 catalogue, 1 collection and 25 items in {len(json_files)} JSON files "
             f"(items laid out on the OS 10 km grid, coloured by maximum elevation)", fontsize=9.5)
fig.savefig(OUT / "01-catalogue-tree.png")
plt.close(fig)

# 02: item footprints by max elevation, and the thumbnails placed at their bounding boxes
fig, axs = plt.subplots(1, 2, figsize=(12.5, 5.6))
gdf.plot(ax=axs[0], column="max_m", cmap=DEM_CMAP, vmin=0, vmax=1100, edgecolor="#333", linewidth=0.8, legend=True,
         legend_kwds={"label": "Maximum elevation in item (m), from raster statistics", "shrink": 0.8})
for _, r in gdf.iterrows():
    axs[0].text(r.cx, r.cy, f"{r.square}\n{r.max_m:.0f} m", ha="center", va="center", fontsize=7)
cb = collection.extent.spatial.bboxes[0]
axs[0].add_patch(Rectangle((cb[0], cb[1]), cb[2] - cb[0], cb[3] - cb[1], fill=False, ec="#d53e4f", lw=1.6, ls="--"))
axs[0].set_title("Item geometries (EPSG:4326) coloured by a property;\ndashed: the collection's spatial extent",
                 fontsize=9)
for _, r in gdf.iterrows():
    b = r.geometry.bounds
    axs[1].imshow(plt.imread(r.thumb), extent=[b[0], b[2], b[1], b[3]])
    axs[1].plot(*r.geometry.exterior.xy, color="white", lw=0.6)
axs[1].set_title("The 25 thumbnail assets, each drawn at its item's bbox", fontsize=9)
for a in axs:
    a.set_xlim(cb[0] - 0.02, cb[2] + 0.02)
    a.set_ylim(cb[1] - 0.02, cb[3] + 0.02)
    a.set_aspect(ASPECT)
    a.set_xlabel("Longitude (°)")
axs[0].set_ylabel("Latitude (°)")
fig.savefig(OUT / "02-item-footprints-thumbnails.png")
plt.close(fig)

# 03: three local queries
fig, axs = plt.subplots(1, 3, figsize=(13.5, 4.9), gridspec_kw={"wspace": 0.12})
qs = [(q_point, "Point: which item contains Yr Wyddfa?", lambda a: a.plot(*SNOWDON_LONLAT, "k^", ms=8)),
      (q_bbox, "Bbox: items touching the Conwy valley box",
       lambda a: a.plot(*q_bbox_geom.exterior.xy, color="black", lw=1.6, ls="--")),
      (q_high, "Property: maximum elevation > 800 m", lambda a: None)]
for a, (sel, title, extra) in zip(axs, qs):
    gdf.plot(ax=a, color="#f0f0f0", edgecolor="#999", lw=0.6)
    gdf[gdf.id.isin(sel.id)].plot(ax=a, color="#3288bd", edgecolor="#08306b", lw=1, alpha=0.75)
    for _, r in gdf.iterrows():
        a.text(r.cx, r.cy, r.square, ha="center", va="center", fontsize=6.5,
               color="white" if r.id in set(sel.id) else "#555")
    extra(a)
    a.set_title(f"{title}\n{len(sel)} of 25 items: {', '.join(sel.square) if len(sel) <= 6 else ''}".rstrip(": "),
                fontsize=9)
    a.set_aspect(ASPECT)
    a.set_xticks([])
    a.set_yticks([])
    a.set_xlabel("")
    a.set_ylabel("")
fig.savefig(OUT / "03-local-queries.png")
plt.close(fig)

# 04: requests, bytes and time: static crawl against API searches
labels = ["Static catalogue\ncrawl (localhost)", "Static catalogue\none known item"] + [s[0] for s in searches]
reqs = [crawl_requests, 1] + [s[2] for s in searches]
secs = [crawl_s, one_s] + [s[3] for s in searches]
nitems = [len(crawled), 1] + [s[1] for s in searches]
fig, axs = plt.subplots(1, 2, figsize=(12, 4.6), sharey=True)
colors = ["#fdae61", "#fee08b", "#3288bd", "#66c2a5", "#5e4fa2"][:len(labels)]
y = np.arange(len(labels))[::-1]
axs[0].barh(y, reqs, color=colors)
for yi, v, n in zip(y, reqs, nitems):
    axs[0].text(v + 0.4, yi, f"{v} request{'s' if v != 1 else ''} for {n} item{'s' if n != 1 else ''}", va="center",
                fontsize=8)
axs[0].set_yticks(y, labels, fontsize=8)
axs[0].set_xlim(0, max(reqs) * 1.6)
axs[0].set_xlabel("HTTP requests")
axs[0].set_title("Requests needed to find the items", fontsize=9)
axs[1].barh(y, secs, color=colors)
for yi, v in zip(y, secs):
    axs[1].text(v + max(secs) * 0.02, yi, f"{v:.2f} s" if v >= 0.1 else f"{v * 1000:.0f} ms", va="center", fontsize=8)
axs[1].set_xlim(0, max(secs) * 1.3)
axs[1].set_xlabel("Wall-clock seconds")
axs[1].set_title("Time taken (APIs over the internet from the UK)", fontsize=9)
fig.savefig(OUT / "04-static-vs-api-requests.png")
plt.close(fig)

# 05: our items against the Earth Search results
fig, axs = plt.subplots(1, 2, figsize=(13, 5.6), gridspec_kw={"width_ratios": [1.15, 1]})
es_gdf = gpd.GeoDataFrame([{"id": i.id, "geometry": shape(i.geometry)} for i in es_dem], crs="EPSG:4326")
s2_gdf = gpd.GeoDataFrame(s2_df, crs="EPSG:4326")
es_gdf.plot(ax=axs[0], facecolor="none", edgecolor="#9e0142", lw=2)
for _, r in es_gdf.iterrows():
    c = r.geometry.centroid
    axs[0].text(c.x, c.y + 0.32, r.id.replace("Copernicus_DSM_COG_10_", "").replace("_DEM", "").replace("_00", ""),
                ha="center", fontsize=8, color="#9e0142", fontweight="bold")
for mg, colr in zip(sorted(s2_gdf.mgrs.unique()), ["#3288bd", "#f46d43", "#66c2a5", "#5e4fa2"]):
    g = s2_gdf[s2_gdf.mgrs == mg]
    g.plot(ax=axs[0], facecolor=colr, edgecolor=colr, alpha=0.06, lw=0.8)
    axs[0].plot([], [], color=colr, lw=4, alpha=0.6, label=f"Sentinel-2 L2A {mg}: {len(g)} scenes")
gdf.plot(ax=axs[0], facecolor="#ffffbf", edgecolor="#333", lw=0.5, alpha=0.9)
axs[0].plot([], [], color="#9e0142", lw=2, label="Copernicus GLO-30 items (Earth Search)")
axs[0].fill_between([], [], color="#ffffbf", label="our 25 items")
axs[0].legend(fontsize=7.5, loc="lower left", framealpha=0.95)
axs[0].set_xlim(-5.4, -2.6)
axs[0].set_ylim(51.85, 54.15)
axs[0].set_aspect(ASPECT)
axs[0].set_xlabel("Longitude (°)")
axs[0].set_ylabel("Latitude (°)")
axs[0].set_title("Search results from Earth Search around our catalogue", fontsize=9)
for mg, mk, colr in zip(sorted(s2_df.mgrs.unique()), ["o", "s"], ["#3288bd", "#f46d43"]):
    g = s2_df[s2_df.mgrs == mg]
    axs[1].scatter(pd.to_datetime(g.date), g["cloud_%"], marker=mk, color=colr, s=40, label=f"tile {mg}",
                   edgecolor="black", lw=0.5)
axs[1].axvline(pd.to_datetime(best_date), color="#999", ls=":", lw=1)
axs[1].text(pd.to_datetime(best_date), 9.6, f" {best_date}", fontsize=8, color="#555", va="top")
axs[1].set_ylim(-0.4, 10)
axs[1].set_ylabel("eo:cloud_cover (%)")
axs[1].set_title("Sentinel-2 L2A, May to September 2026, eo:cloud_cover < 10", fontsize=9)
axs[1].legend(fontsize=8, loc="center right")
axs[1].grid(alpha=0.3)
fig.autofmt_xdate()
fig.savefig(OUT / "05-earth-search-results.png")
plt.close(fig)

# 06: pixels read through the hrefs of the query results
fig, axs = plt.subplots(1, 2, figsize=(12.5, 5.6))
mb_ = rasterio.transform.array_bounds(*mosaic.shape, mosaic_tf)
ext = [mb_[0], mb_[2], mb_[1], mb_[3]]
axs[0].imshow(np.ma.masked_where(np.isnan(mosaic), LIGHT.hillshade(np.nan_to_num(mosaic), vert_exag=1.5, dx=25,
                                                                   dy=25)), cmap="gray", extent=ext)
im = axs[0].imshow(np.ma.masked_invalid(np.where(mosaic > 0, mosaic, np.nan)), cmap=DEM_CMAP, vmin=0, vmax=1100, alpha=0.5, extent=ext)
for _, r in gdf[gdf.id.isin(q_high.id)].to_crs("EPSG:27700").iterrows():
    axs[0].plot(*r.geometry.exterior.xy, color="black", lw=0.6)
fig.colorbar(im, ax=axs[0], shrink=0.75, label="Elevation (m)")
axs[0].set_title(f"Local catalogue: the {len(q_high)} items with maximum > 800 m,\nmerged from their asset hrefs",
                 fontsize=9)
axs[1].imshow(np.transpose(rgb, (1, 2, 0)), extent=[AOI_BNG[0], AOI_BNG[2], AOI_BNG[1], AOI_BNG[3]])
for _, r in gdf.to_crs("EPSG:27700").iterrows():
    axs[1].plot(*r.geometry.exterior.xy, color="white", lw=0.5, alpha=0.8)
axs[1].set_title(f"Earth Search: Sentinel-2 true colour, {best_date}, read at {S2_RES} m\n"
                 f"from {len(best)} 'visual' COG assets ({s2_bytes / 1e6:.1f} MB in {len(RANGES)} range requests)",
                 fontsize=9)
for a in axs:
    a.set_xlim(AOI_BNG[0], AOI_BNG[2])
    a.set_ylim(AOI_BNG[1], AOI_BNG[3])
    a.xaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v / 1000:.0f}"))
    a.yaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v / 1000:.0f}"))
    a.set_xlabel("Easting (km, EPSG:27700)")
axs[0].set_ylabel("Northing (km)")
fig.savefig(OUT / "06-read-through-hrefs.png")
plt.close(fig)

print("wrote", ", ".join(p.name for p in sorted(OUT.glob("*.png"))))