"""
Cloud Optimised GeoTIFFs in Python: reproject and mosaic Copernicus DEM tiles into a validated COG.

Companion script for https://technicalgis.com/tutorials/cloud-optimised-geotiff-python/
(part 1 of the series "Cloud-native rasters in Python"; part 2 builds a STAC catalogue over the outputs).

Data: Copernicus DEM GLO-30, four 1 x 1 degree COG tiles over Eryri (Snowdonia), read anonymously
from the AWS Open Data bucket copernicus-dem-30m. (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.

Run:  python cog_mosaic.py      (a few minutes; it downloads roughly 200 MB in total, most of it
                                 for the deliberate whole-file comparisons in Step 2)
Outputs (in output/): eryri_dem_25m_cog.tif, the block-size and compression variants, a set of
25 COG tiles on the Ordnance Survey 10 km grid in output/tiles/, six figures and this console log.
"""
from __future__ import annotations

import http.server
import logging
import os
import re
import statistics
import threading
import time
from functools import partial
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd
import rasterio
import rasterio.shutil
import requests
from matplotlib.colors import LightSource, ListedColormap
from matplotlib.patches import Rectangle
from pyproj import Transformer
from rasterio.enums import Resampling
from rasterio.merge import merge
from rasterio.transform import array_bounds, from_origin
from rasterio.warp import reproject, transform_bounds
from rasterio.windows import Window, from_bounds
from rio_cogeo.cogeo import cog_info, cog_translate, cog_validate
from rio_cogeo.profiles import cog_profiles

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

HERE = Path(__file__).resolve().parent
OUT = HERE / "output"
TILE_DIR = OUT / "tiles"
OUT.mkdir(exist_ok=True)
TILE_DIR.mkdir(exist_ok=True)

TILES = ["N53_00_W005_00", "N53_00_W004_00", "N52_00_W005_00", "N52_00_W004_00"]
BUCKET = "https://copernicus-dem-30m.s3.amazonaws.com"
BNG = "EPSG:27700"
AOI = (250_000, 330_000, 300_000, 380_000)   # 50 x 50 km of Eryri in British National Grid (OS square SH)
RES = 25                                      # output pixel size in metres: 10 km = 400 pixels exactly
SNOWDON = (260_968, 354_370)                  # Yr Wyddfa summit trig point, easting/northing
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)
LIGHT = LightSource(azdeg=315, altdeg=45)
DEM_CMAP = ListedColormap(plt.cm.terrain(np.linspace(0.25, 1.0, 256)))   # land colours only; sea drawn separately
SEA = ListedColormap(["#9ecae1"])


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


def tile_url(tile: str) -> str:
    name = f"Copernicus_DSM_COG_10_{tile}_DEM"
    return f"{BUCKET}/{name}/{name}.tif"


def mb(n: float) -> float:
    return n / 1e6


def shade(z: np.ndarray, dx: float, dy: float) -> np.ndarray:
    return LIGHT.hillshade(np.nan_to_num(z, nan=0.0), vert_exag=1.5, dx=dx, dy=dy)


def draw_dem(ax, z, dx, dy, extent, alpha=0.5):
    """Hillshade, then land coloured by height and the sea (0 m or below in GLO-30) in flat blue."""
    ax.imshow(shade(z, dx, dy), cmap="gray", extent=extent, vmin=0, vmax=1)
    im = ax.imshow(np.ma.masked_less_equal(z, 0), cmap=DEM_CMAP, vmin=0, vmax=1100, alpha=alpha, extent=extent)
    ax.imshow(np.ma.masked_greater(z, 0), cmap=SEA, extent=extent)
    return im


# GDAL logs every HTTP range request it makes ("Downloading 0-16383 ..."). We count them without printing URLs.
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())

RUN = iter(range(10_000))


def fresh(url: str) -> str:
    """Add a unique query string so GDAL's in-process HTTP cache cannot serve a repeated read."""
    return f"{url}?run={next(RUN)}"


def measured(fn, repeat: int = 3):
    """Run fn() `repeat` times with GDAL debug logging on; return result, median seconds, requests and bytes."""
    times, counts = [], []
    for _ in range(repeat):
        RANGES.clear()
        logging.getLogger("rasterio").setLevel(logging.DEBUG)
        t0 = time.perf_counter()
        with rasterio.Env(CPL_DEBUG=True, **GDAL_ENV):
            result = fn()
        times.append(time.perf_counter() - t0)
        logging.getLogger("rasterio").setLevel(logging.WARNING)
        counts.append((len(RANGES), sum(b - a + 1 for a, b in RANGES)))
    return result, statistics.median(times), counts[-1][0], counts[-1][1]


# ---------------------------------------------------------------- Step 1: inspect the source COGs
section("Step 1: inspect the four Copernicus GLO-30 tiles on AWS")
rows, thumbs = [], {}
for t in TILES:
    url = tile_url(t)
    size = int(requests.head(url, timeout=60).headers["Content-Length"])

    def read_thumb(url=url):
        with rasterio.open(fresh(url)) as src:
            meta = dict(width=src.width, height=src.height, block=src.block_shapes[0], overviews=src.overviews(1),
                        compress=src.compression.value if src.compression else None,
                        predictor=src.tags(ns="IMAGE_STRUCTURE").get("PREDICTOR"), nodata=src.nodata,
                        res=src.res, bounds=src.bounds, crs=src.crs.to_string())
            arr = src.read(1, out_shape=(src.height // 8, src.width // 8))   # served from the 1/8 overview
            return meta, arr

    (meta, arr), secs, nreq, nbytes = measured(read_thumb, repeat=1)
    thumbs[t] = (arr, meta["bounds"])
    rows.append({"tile": t, "size_px": f"{meta['width']} x {meta['height']}", "blocks": "x".join(map(str, meta["block"])),
                 "overviews": meta["overviews"], "compress": meta["compress"], "predictor": meta["predictor"],
                 "file_MB": round(mb(size), 1), "thumb_requests": nreq, "thumb_MB": round(mb(nbytes), 2)})
print(pd.DataFrame(rows).to_string(index=False))
print("crs:", meta["crs"], "| pixel size (deg):", tuple(round(r, 7) for r in meta["res"]), "| nodata:", meta["nodata"])

aoi_ll = transform_bounds(BNG, "EPSG:4326", *AOI, densify_pts=51)
print("AOI in BNG:", AOI, "-> lon/lat bounds:", tuple(round(v, 4) for v in aoi_ll))

# ---------------------------------------------------------------- Step 2: a windowed read against the whole file
section("Step 2: what a windowed read costs, against downloading the whole file")
big = tile_url("N52_00_W004_00")               # the largest of the four tiles: Arenig, Bala and the Berwyn
big_size = int(requests.head(big, timeout=60).headers["Content-Length"])
ARENIG = (-3.745, 52.917)                      # Arenig Fawr summit, lon/lat
win_ll = (ARENIG[0] - 0.04, ARENIG[1] - 0.025, ARENIG[0] + 0.04, ARENIG[1] + 0.025)   # about 5.4 x 5.6 km


def open_only():
    with rasterio.open(fresh(big)) as src:
        return src.profile


def window_read():
    with rasterio.open(fresh(big)) as src:
        return src.read(1, window=from_bounds(*win_ll, transform=src.transform))


def full_read():
    with rasterio.open(fresh(big)) as src:
        return src.read(1)


remote = []
for label, fn in [("open: header only", open_only), ("5 km window, full resolution", window_read),
                  ("whole tile via rasterio", full_read)]:
    res, secs, nreq, nbytes = measured(fn)
    shape = f"{res.shape[1]} x {res.shape[0]}" if isinstance(res, np.ndarray) else "-"
    remote.append({"job": label, "pixels": shape, "requests": nreq, "MB": mb(nbytes),
                   "pct_of_file": 100 * nbytes / big_size, "seconds": secs})
times = []
for _ in range(3):
    t0 = time.perf_counter()
    nbytes = len(requests.get(big, timeout=600).content)
    times.append(time.perf_counter() - t0)
remote.append({"job": "download whole file (requests.get)", "pixels": "-", "requests": 1, "MB": mb(nbytes),
               "pct_of_file": 100 * nbytes / big_size, "seconds": statistics.median(times)})
remote = pd.DataFrame(remote)
print(f"tile N52_00_W004_00: {mb(big_size):.1f} MB on S3 (eu-central-1); median of 3 runs")
print(remote.round(3).to_string(index=False))

# ---------------------------------------------------------------- Step 3: mosaic and reproject
section("Step 3: mosaic the four tiles and reproject to British National Grid")
pad = 0.01
read_bounds = (aoi_ll[0] - pad, aoi_ll[1] - pad, aoi_ll[2] + pad, aoi_ll[3] + pad)
pieces = {}
for t in TILES:
    with rasterio.Env(**GDAL_ENV), rasterio.open(tile_url(t)) as src:
        w = from_bounds(*read_bounds, transform=src.transform).intersection(Window(0, 0, src.width, src.height))
        w = w.round_offsets().round_lengths()
        pieces[t] = (src.read(1, window=w), src.window_bounds(w))


def do_merge():
    return merge([fresh(tile_url(t)) for t in TILES], bounds=read_bounds, resampling=Resampling.nearest)


(mosaic_ll, ll_transform), secs, nreq, nbytes = measured(do_merge, repeat=1)
mosaic_ll = mosaic_ll[0]
print(f"rasterio.merge: {mosaic_ll.shape[1]} x {mosaic_ll.shape[0]} px in EPSG:4326 from 4 remote tiles, "
      f"{nreq} range requests, {mb(nbytes):.2f} MB, {secs:.1f} s")
for t, (arr, b) in pieces.items():
    print(f"  {t}: window {arr.shape[1]} x {arr.shape[0]} px")

width = height = int((AOI[2] - AOI[0]) / RES)
dst_transform = from_origin(AOI[0], AOI[3], RES, RES)
dem = np.full((height, width), np.nan, dtype="float32")
t0 = time.perf_counter()
reproject(mosaic_ll, dem, src_transform=ll_transform, src_crs="EPSG:4326", dst_transform=dst_transform,
          dst_crs=BNG, resampling=Resampling.bilinear, dst_nodata=np.nan)
print(f"reprojected to {BNG} at {RES} m: {width} x {height} px in {time.perf_counter() - t0:.2f} s; "
      f"NaN pixels: {int(np.isnan(dem).sum())}")
print(f"elevation min/mean/max (m): {np.nanmin(dem):.1f} / {np.nanmean(dem):.1f} / {np.nanmax(dem):.1f}; "
      f"sea and lakes at 0 m or below: {100 * np.mean(dem <= 0):.1f}%")

# ---------------------------------------------------------------- Step 4: write a plain GeoTIFF and COG variants
section("Step 4: write a plain GeoTIFF, then COGs with different compression and block sizes")
plain = OUT / "eryri_dem_25m_plain.tif"
profile = dict(driver="GTiff", dtype="float32", count=1, width=width, height=height, crs=BNG,
               transform=dst_transform, nodata=None)
with rasterio.open(plain, "w", **profile) as dst:
    dst.write(dem, 1)
    dst.set_band_description(1, "elevation (m, EGM2008)")
plain_ok, _, _ = cog_validate(str(plain), quiet=True)
print(f"plain GeoTIFF (GDAL defaults): {mb(plain.stat().st_size):.2f} MB, valid COG: {plain_ok}")

variants = [
    ("DEFLATE, no predictor", 512, dict(COMPRESS="DEFLATE")),
    ("DEFLATE + predictor", 512, dict(COMPRESS="DEFLATE", PREDICTOR="YES")),
    ("ZSTD + predictor", 512, dict(COMPRESS="ZSTD", PREDICTOR="YES")),
    ("LERC_ZSTD, max error 0.1 m", 512, dict(COMPRESS="LERC_ZSTD", MAX_Z_ERROR="0.1")),
    ("DEFLATE + predictor", 256, dict(COMPRESS="DEFLATE", PREDICTOR="YES")),
    ("DEFLATE + predictor", 1024, dict(COMPRESS="DEFLATE", PREDICTOR="YES")),
]
var_rows = []
for label, block, opts in variants:
    tag = re.sub(r"[^a-z0-9]+", "_", f"{label} {block}".lower()).strip("_")
    path = OUT / f"eryri_dem_25m_{tag}.tif"
    t0 = time.perf_counter()
    rasterio.shutil.copy(plain, path, driver="COG", BLOCKSIZE=block, OVERVIEW_RESAMPLING="AVERAGE", **opts)
    secs = time.perf_counter() - t0
    ok, errors, warnings = cog_validate(str(path), quiet=True)
    with rasterio.open(path) as src:
        err = float(np.nanmax(np.abs(src.read(1) - dem)))
        ovr = src.overviews(1)
    var_rows.append({"variant": label, "block": block, "MB": mb(path.stat().st_size), "write_s": secs,
                     "valid": ok, "overviews": ovr, "max_abs_error_m": err, "path": path})
variants_df = pd.DataFrame(var_rows)
print(variants_df.drop(columns="path").round(3).to_string(index=False))

# The main output: rio-cogeo with DEFLATE, a floating-point predictor, 512 px blocks and average overviews.
cog = OUT / "eryri_dem_25m_cog.tif"
cog_profile = cog_profiles.get("deflate")
cog_profile.update(predictor=3, blockxsize=512, blockysize=512)
t0 = time.perf_counter()
cog_translate(plain, cog, cog_profile, overview_resampling="average", quiet=True,
              additional_cog_metadata={"SOURCE": "Copernicus DEM GLO-30 (AWS copernicus-dem-30m), 4 tiles"})
ok, errors, warnings = cog_validate(str(cog), quiet=True)
print(f"rio-cogeo COG: {mb(cog.stat().st_size):.2f} MB in {time.perf_counter() - t0:.2f} s; valid COG: {ok}; "
      f"errors: {errors}; warnings: {warnings}")

# ---------------------------------------------------------------- Step 5: look inside the COG
section("Step 5: internal tiling and overviews")
info = cog_info(str(cog))
print("compression:", info.Compression, "| COG layout:", info.Tags["Image Structure"].get("LAYOUT"),
      "| predictor:", info.Tags["Image Structure"].get("PREDICTOR"))
for ifd in info.IFD:
    print(f"  IFD level {ifd.Level}: {ifd.Width} x {ifd.Height} px, blocks {ifd.Blocksize[1]} x {ifd.Blocksize[0]}, "
          f"decimation {ifd.Decimation}")
with rasterio.open(cog) as src:
    levels = [src.read(1)] + [src.read(1, out_shape=(src.height // f, src.width // f)) for f in src.overviews(1)]
    factors = [1] + src.overviews(1)

# ---------------------------------------------------------------- Step 6: cut a tile set on the OS 10 km grid
section("Step 6: cut 25 COG tiles on the Ordnance Survey 10 km grid (input for part 2)")
for old in TILE_DIR.glob("*.tif"):
    old.unlink()
tile_rows = []
with rasterio.open(cog) as src:
    for e in range(AOI[0], AOI[2], 10_000):
        for n in range(AOI[1], AOI[3], 10_000):
            name = f"SH{(e % 100_000) // 10_000}{(n % 100_000) // 10_000}"     # e.g. SH65 contains Yr Wyddfa
            w = from_bounds(e, n, e + 10_000, n + 10_000, transform=src.transform).round_offsets().round_lengths()
            arr = src.read(1, window=w)
            tmp = TILE_DIR / f"_{name}.tif"
            with rasterio.open(tmp, "w", **{**profile, "width": w.width, "height": w.height,
                                            "transform": src.window_transform(w)}) as dst:
                dst.write(arr, 1)
                dst.set_band_description(1, "elevation (m, EGM2008)")
            path = TILE_DIR / f"{name}.tif"
            rasterio.shutil.copy(tmp, path, driver="COG", BLOCKSIZE=256, COMPRESS="DEFLATE", PREDICTOR="YES",
                                 OVERVIEW_RESAMPLING="AVERAGE")
            tmp.unlink()
            tile_rows.append({"tile": name, "easting": e, "northing": n,
                              "valid": cog_validate(str(path), quiet=True)[0],
                              "kB": path.stat().st_size / 1e3, "min_m": float(arr.min()), "max_m": float(arr.max())})
tiles_df = pd.DataFrame(tile_rows)
print(f"{len(tiles_df)} tiles of 400 x 400 px, all valid COGs: {tiles_df.valid.all()}, "
      f"total {tiles_df.kB.sum() / 1e3:.2f} MB")
print(tiles_df.nlargest(5, "max_m").drop(columns=["easting", "northing"]).round(1).to_string(index=False))

# ---------------------------------------------------------------- Step 7: range reads from our own files
section("Step 7: range reads from our own files over HTTP")
served: list[int] = []


class RangeHandler(http.server.SimpleHTTPRequestHandler):
    """Static file server that honours single HTTP Range requests and records the bytes it sends."""

    def log_message(self, *args):
        pass

    def send_head(self):
        path = self.translate_path(self.path)
        size = os.path.getsize(path)
        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", "image/tiff")
        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":
            served.append(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(RangeHandler, directory=str(OUT)))
threading.Thread(target=server.serve_forever, daemon=True).start()
base = f"http://127.0.0.1:{server.server_address[1]}"
snow_win = (SNOWDON[0] - 1000, SNOWDON[1] - 1000, SNOWDON[0] + 1000, SNOWDON[1] + 1000)

files = [("plain GeoTIFF", plain)] + [(f"COG {r.block} px blocks", r.path) for r in variants_df.itertuples()
                                       if r.variant == "DEFLATE + predictor"]
jobs = {"2 x 2 km window at Yr Wyddfa": lambda s: s.read(1, window=from_bounds(*snow_win, transform=s.transform)),
        "whole area at 1/16 resolution": lambda s: s.read(1, out_shape=(s.height // 16, s.width // 16)),
        "whole file, full resolution": lambda s: s.read(1)}
local = []
for label, path in files:
    for job, fn in jobs.items():
        times = []
        for _ in range(3):
            served.clear()
            t0 = time.perf_counter()
            with rasterio.Env(**GDAL_ENV), rasterio.open(fresh(f"/vsicurl/{base}/{path.name}")) as s:
                arr = fn(s)
            times.append(time.perf_counter() - t0)
        local.append({"file": label, "job": job, "pixels": f"{arr.shape[1]} x {arr.shape[0]}",
                      "requests": len(served), "MB": mb(sum(served)),
                      "pct_of_file": 100 * sum(served) / path.stat().st_size, "ms": 1000 * statistics.median(times)})
server.shutdown()
local = pd.DataFrame(local)
print(local.round(3).to_string(index=False))

with rasterio.open(cog) as src:
    peak = src.read(1, window=from_bounds(*snow_win, transform=src.transform))
print(f"highest pixel within 1 km of Yr Wyddfa: {peak.max():.1f} m (published summit height 1,085 m)")
with rasterio.Env(**GDAL_ENV), rasterio.open(tile_url("N53_00_W005_00")) as src:
    native = src.read(1, window=from_bounds(-4.0913, 53.0595, -4.0613, 53.0775, transform=src.transform))
print(f"highest native GLO-30 pixel in the same area: {native.max():.1f} m")

# ---------------------------------------------------------------- Figures
section("figures")
TILE_COLOURS = {"N53_00_W005_00": "#3288bd", "N53_00_W004_00": "#66c2a5",
                "N52_00_W005_00": "#f46d43", "N52_00_W004_00": "#9e0142"}
aoi_x = np.r_[np.linspace(AOI[0], AOI[2], 50), np.full(50, AOI[2]), np.linspace(AOI[2], AOI[0], 50), np.full(50, AOI[0])]
aoi_y = np.r_[np.full(50, AOI[1]), np.linspace(AOI[1], AOI[3], 50), np.full(50, AOI[3]), np.linspace(AOI[3], AOI[1], 50)]
aoi_lon, aoi_lat = Transformer.from_crs(BNG, "EPSG:4326", always_xy=True).transform(aoi_x, aoi_y)
ext_bng = [AOI[0], AOI[2], AOI[1], AOI[3]]


def km_axes(ax):
    ax.xaxis.set_major_formatter(lambda v, p: f"{v / 1000:.0f}")
    ax.yaxis.set_major_formatter(lambda v, p: f"{v / 1000:.0f}")
    ax.set_xlabel("Easting (km, EPSG:27700)")
    ax.set_ylabel("Northing (km)")


# 01: the four source tiles, read from their 1/8 overviews, with the area of interest
fig, ax = plt.subplots(figsize=(7.5, 6.2))
for t, (arr, b) in thumbs.items():
    ext = [b.left, b.right, b.bottom, b.top]
    dx = 111_320 * np.cos(np.radians(53)) * (b.right - b.left) / arr.shape[1]
    dy = 111_320 * (b.top - b.bottom) / arr.shape[0]
    draw_dem(ax, arr, dx, dy, ext)
    ax.add_patch(Rectangle((b.left, b.bottom), b.right - b.left, b.top - b.bottom, fill=False,
                           ec=TILE_COLOURS[t], lw=2))
    ax.text(b.left + 0.03, b.top - 0.04, t.replace("_00", ""), color=TILE_COLOURS[t], fontsize=8, fontweight="bold",
            va="top", bbox=dict(fc="white", ec="none", alpha=0.85, pad=1.5))
ax.plot(aoi_lon, aoi_lat, color="black", lw=1.6, ls="--", label="50 x 50 km area of interest (BNG)")
ax.set_xlim(-5.05, -2.95)
ax.set_ylim(51.95, 54.05)
ax.set_aspect(1 / np.cos(np.radians(53)))
ax.set_xlabel("Longitude (°)")
ax.set_ylabel("Latitude (°)")
ax.legend(loc="lower left", fontsize=8, framealpha=0.9)
ax.set_title("Four Copernicus GLO-30 tiles, each read from its 1/8 overview")
fig.savefig(OUT / "01-source-tiles-overviews.png")
plt.close(fig)

# 02: mosaic steps
fig, axs = plt.subplots(1, 3, figsize=(14, 5.2), gridspec_kw={"wspace": 0.32})
lb = array_bounds(*mosaic_ll.shape, ll_transform)
ext_ll = [lb[0], lb[2], lb[1], lb[3]]
dxl = 111_320 * np.cos(np.radians(53.1)) * ll_transform.a
dyl = 111_320 * -ll_transform.e
for t, (arr, b) in pieces.items():
    draw_dem(axs[0], arr, dxl, dyl, [b[0], b[2], b[1], b[3]])
    axs[0].add_patch(Rectangle((b[0], b[1]), b[2] - b[0], b[3] - b[1], fill=False, ec=TILE_COLOURS[t], lw=2.2))
axs[0].set_title("1. Four windows, read from\nfour remote files (outlined)", fontsize=9)
draw_dem(axs[1], mosaic_ll, dxl, dyl, ext_ll)
axs[1].plot(aoi_lon, aoi_lat, color="black", lw=1.4, ls="--")
axs[1].set_title(f"2. rasterio.merge: one {mosaic_ll.shape[1]} x {mosaic_ll.shape[0]} px\ngrid in EPSG:4326 (dashed: AOI)",
                 fontsize=9)
for a in axs[:2]:
    a.set_xlim(lb[0], lb[2])
    a.set_ylim(lb[1], lb[3])
    a.set_aspect(1 / np.cos(np.radians(53.1)))
    a.set_xlabel("Longitude (°)")
axs[0].set_ylabel("Latitude (°)")
im = draw_dem(axs[2], dem, RES, RES, ext_bng)
axs[2].set_title(f"3. Reprojected and resampled:\n{width} x {height} px at {RES} m in EPSG:27700", fontsize=9)
km_axes(axs[2])
fig.colorbar(im, ax=axs, shrink=0.75, label="Elevation (m)", pad=0.02)
fig.savefig(OUT / "02-mosaic-reproject-steps.png")
plt.close(fig)

# 03: internal tiles and overview pyramid
fig, axs = plt.subplots(1, len(levels), figsize=(4.2 * len(levels), 4.6))
for ax, arr, f in zip(axs, levels, factors):
    draw_dem(ax, arr, RES * f, RES * f, [0, arr.shape[1], arr.shape[0], 0])
    nb = 0
    for r in range(0, arr.shape[0], 512):
        for c in range(0, arr.shape[1], 512):
            ax.add_patch(Rectangle((c, r), min(512, arr.shape[1] - c), min(512, arr.shape[0] - r), fill=False,
                                   ec="#d53e4f", lw=1.4))
            nb += 1
    level = "full resolution" if f == 1 else f"overview 1/{f}"
    ax.set_title(f"{level}: {arr.shape[1]} x {arr.shape[0]} px, {nb} block{'s' if nb > 1 else ''}", fontsize=9)
    ax.set_xticks([])
    ax.set_yticks([])
fig.suptitle("Inside the COG: 512 x 512 px internal tiles (red) at each level of the overview pyramid", fontsize=10)
fig.savefig(OUT / "03-internal-tiles-overviews.png")
plt.close(fig)

# 04: file size and write time for each variant
labels = ["plain GeoTIFF\n(uncompressed strips)"] + [f"{r.variant}\n{r.block} px blocks" for r in variants_df.itertuples()]
sizes = [mb(plain.stat().st_size)] + list(variants_df.MB)
wtimes = [np.nan] + list(variants_df.write_s)
colours = ["#9e9e9e", "#fdae61", "#3288bd", "#66c2a5", "#9e0142", "#5e4fa2", "#5e4fa2"]
fig, axs = plt.subplots(1, 2, figsize=(11, 4.8), sharey=True)
y = np.arange(len(labels))[::-1]
axs[0].barh(y, sizes, color=colours)
for yi, s in zip(y, sizes):
    axs[0].text(s + 0.2, yi, f"{s:.2f} MB", va="center", fontsize=8)
axs[0].set_yticks(y, labels, fontsize=8)
axs[0].set_xlabel("File size (MB)")
axs[0].set_xlim(0, max(sizes) * 1.25)
axs[0].set_title("File size, the same 2000 x 2000 float32 DEM", fontsize=9)
axs[1].barh(y[1:], wtimes[1:], color=colours[1:])
for yi, s in zip(y[1:], wtimes[1:]):
    axs[1].text(s + 0.01, yi, f"{s:.2f} s", va="center", fontsize=8)
axs[1].set_xlabel("Time to write with the GDAL COG driver (s)")
axs[1].set_xlim(0, np.nanmax(wtimes) * 1.3)
axs[1].set_title("Write time, including overviews", fontsize=9)
for ax in axs:
    ax.grid(axis="x", alpha=0.3)
fig.savefig(OUT / "04-compression-blocksize.png")
plt.close(fig)

# 05: bytes transferred for each read
fig, axs = plt.subplots(2, 1, figsize=(10, 9), gridspec_kw={"height_ratios": [1, 1.6], "hspace": 0.35})
r = remote.iloc[::-1]
axs[0].barh(np.arange(len(r)), r.MB, color=["#9e0142", "#f46d43", "#3288bd", "#66c2a5"])
for i, row in enumerate(r.itertuples()):
    axs[0].text(row.MB * 1.15, i, f"{row.MB:.2f} MB, {row.requests} req., {row.seconds:.2f} s", va="center", fontsize=8)
axs[0].set_yticks(np.arange(len(r)), r.job, fontsize=8)
axs[0].set_xscale("log")
axs[0].set_xlim(0.005, 200)
axs[0].set_xlabel("Data transferred (MB, log scale)")
axs[0].set_title(f"Copernicus tile N52_W004 on AWS S3 ({mb(big_size):.1f} MB), read over the internet", fontsize=9)
job_names = list(jobs)
file_names = [f[0] for f in files]
fcol = {"plain GeoTIFF": "#9e9e9e", "COG 256 px blocks": "#fdae61", "COG 512 px blocks": "#3288bd",
        "COG 1024 px blocks": "#5e4fa2"}
h = 0.8 / len(file_names)
for j, fname in enumerate(file_names):
    sub = local[local.file == fname].set_index("job").loc[job_names]
    ys = np.arange(len(job_names))[::-1] - (j - (len(file_names) - 1) / 2) * h
    axs[1].barh(ys, sub.MB, height=h, color=fcol.get(fname, "#666"), label=fname)
    for yi, row in zip(ys, sub.itertuples()):
        axs[1].text(row.MB * 1.15, yi, f"{row.MB:.2f} MB, {row.requests} req.", va="center", fontsize=7)
axs[1].set_yticks(np.arange(len(job_names))[::-1], job_names, fontsize=8)
axs[1].set_xscale("log")
axs[1].set_xlim(0.1, 200)
axs[1].set_xlabel("Data transferred (MB, log scale)")
axs[1].set_title("Our 2000 x 2000 px Eryri DEM, served locally with HTTP range support", fontsize=9)
axs[1].legend(fontsize=8, loc="upper right", framealpha=0.95)
for ax in axs:
    ax.grid(axis="x", alpha=0.3, which="major")
fig.savefig(OUT / "05-range-request-bytes.png")
plt.close(fig)

# 06: the final COG with the 10 km tile grid used in part 2
fig, ax = plt.subplots(figsize=(8, 7))
im = draw_dem(ax, dem, RES, RES, ext_bng)
for row in tiles_df.itertuples():
    ax.add_patch(Rectangle((row.easting, row.northing), 10_000, 10_000, fill=False, ec="#1a1a1a", lw=0.8))
    ax.text(row.easting + 500, row.northing + 9_500, row.tile, fontsize=7.5, va="top", color="#1a1a1a",
            bbox=dict(fc="white", ec="none", alpha=0.75, pad=1))
ax.add_patch(Rectangle(snow_win[:2], 2000, 2000, fill=False, ec="#d53e4f", lw=1.8))
ax.annotate("Yr Wyddfa\n2 km test window", (snow_win[2], snow_win[3]), xytext=(266_500, 362_500), fontsize=8,
            color="#9e0142", arrowprops=dict(arrowstyle="-", color="#9e0142"),
            bbox=dict(fc="white", ec="none", alpha=0.8, pad=1))
km_axes(ax)
fig.colorbar(im, ax=ax, shrink=0.75, label="Elevation (m)")
ax.set_title("Eryri DEM COG at 25 m, cut into 25 tiles on the OS 10 km grid")
fig.savefig(OUT / "06-eryri-cog-tile-grid.png")
plt.close(fig)
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("0*.png"))))