"""Glacier velocity of Jakobshavn Isbrae (Sermeq Kujalleq), Greenland, from NASA ITS_LIVE in Python.

Finds the ITS_LIVE datacube that covers a point from the public catalogue, opens the image-pair
velocity datacube and the annual composites as Zarr straight from AWS S3 (no login), maps the
velocity field, traces a flowline by integrating the velocity vectors, extracts speed profiles along
it for 1985-2025, filters the image-pair time series at points on the flowline by error, time
separation and a rolling-median test, and computes strain rates and flow directions.
Run with:  python glacier_velocity_its_live.py
"""
import platform
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.pyplot as plt
import numpy as np
import pandas as pd
import requests
import xarray as xr
import zarr
from matplotlib.colors import LogNorm, TwoSlopeNorm
from pyproj import Transformer
from scipy import ndimage
from scipy.interpolate import RegularGridInterpolator
from shapely.geometry import Point

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"numpy {np.__version__} | pandas {pd.__version__} | scipy {version('scipy')} | matplotlib {matplotlib.__version__}")


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


def km_axes(ax) -> None:
    ax.xaxis.set_major_formatter(lambda val, _: f"{val / 1000:.0f}")
    ax.yaxis.set_major_formatter(lambda val, _: f"{val / 1000:.0f}")


# --- Step 1: find the datacube that covers the glacier ---------------------------------------
# ITS_LIVE splits the world into 100 km x 100 km Zarr datacubes in local polar stereographic or UTM grids.
# The catalogue is a GeoJSON file listing every cube's footprint (in lon/lat) and its URL.
LON, LAT = -49.14, 69.10            # a point on the main trunk of Jakobshavn Isbrae, ~25 km above the front
CATALOG = "https://its-live-data.s3.amazonaws.com/datacubes/catalog_v02.json"
cat_path = DATA / "catalog_v02.json"
if not cat_path.exists():
    cat_path.write_bytes(requests.get(CATALOG, timeout=300).content)
catalog = gpd.read_file(cat_path)
hit = catalog[catalog.contains(Point(LON, LAT))].iloc[0]
section("datacube lookup")
print(f"catalogue: {len(catalog)} datacubes")
print(f"cube:      {hit.zarr_url}")
print(f"composite: {hit.composite_zarr_url}")
print(f"EPSG:{hit.epsg}, {hit.granule_count} image pairs")

cube = xr.open_dataset(hit.zarr_url, engine="zarr", chunks={})          # lazy: nothing downloaded yet
comp = xr.open_dataset(hit.composite_zarr_url, engine="zarr", chunks={})
print(f"datacube: {cube.sizes['mid_date']} image pairs on a {cube.sizes['y']} x {cube.sizes['x']} grid "
      f"of {float(cube.x[1] - cube.x[0]):.0f} m pixels, chunks {cube.v.encoding['chunks']}")
print(f"mid_date range: {pd.Timestamp(cube.mid_date.values.min()):%Y-%m-%d} to "
      f"{pd.Timestamp(cube.mid_date.values.max()):%Y-%m-%d}")
print(f"composites: {comp.sizes['time']} annual mosaics, {pd.Timestamp(comp.time.values[0]).year}-"
      f"{pd.Timestamp(comp.time.values[-1]).year}")

# --- Step 2: velocity maps from the annual composites ----------------------------------------
section("annual composites")
years = pd.DatetimeIndex(comp.time.values).year
v0 = comp.v0.load()                                  # 2014-2024 climatological speed (time intercept 2018-01-01)
vx0, vy0 = comp.vx0.load(), comp.vy0.load()
v_amp, v_phase = comp.v_amp.load(), comp.v_phase.load()
landice = comp.landice.load()
x, y = comp.x.values, comp.y.values
ext = [x[0] - 60, x[-1] + 60, y[-1] - 60, y[0] + 60]
print(f"fastest climatological speed in the cube: {float(v0.max()):,.0f} m/yr")
print(f"pixels with v0 > 1 km/yr: {int((v0 > 1000).sum()):,} ({int((v0 > 1000).sum()) * 0.12 ** 2:,.0f} km2)")

# --- Step 3: trace a flowline by integrating the velocity vectors ----------------------------
# Smooth the climatological vectors a little, then step along the unit flow direction (RK4) from a seed
# on the trunk: upstream until the ice is slow or we leave the cube, and downstream to the front.
valid = ndimage.gaussian_filter(np.isfinite(vx0.values).astype(float), 2)
sx = ndimage.gaussian_filter(np.nan_to_num(vx0.values), 2) / np.maximum(valid, 1e-6)
sy = ndimage.gaussian_filter(np.nan_to_num(vy0.values), 2) / np.maximum(valid, 1e-6)
sx[valid < 0.5] = np.nan
sy[valid < 0.5] = np.nan
grid = (y[::-1], x)                                   # RegularGridInterpolator needs ascending axes
ix = RegularGridInterpolator(grid, sx[::-1], bounds_error=False, fill_value=np.nan)
iy = RegularGridInterpolator(grid, sy[::-1], bounds_error=False, fill_value=np.nan)
spd = RegularGridInterpolator(grid, v0.values[::-1], bounds_error=False, fill_value=np.nan)


def unit(p):
    u_, v_ = ix([p[::-1]])[0], iy([p[::-1]])[0]
    n = np.hypot(u_, v_)
    return np.array([u_, v_]) / n if np.isfinite(n) and n > 0 else None


def trace(start, sign, step=120.0, max_km=150, min_speed=200.0):
    pts_ = [np.asarray(start, float)]
    for _ in range(int(max_km * 1000 / step)):
        p = pts_[-1]
        k1 = unit(p)
        k2 = unit(p + sign * step / 2 * k1) if k1 is not None else None
        k3 = unit(p + sign * step / 2 * k2) if k2 is not None else None
        k4 = unit(p + sign * step * k3) if k3 is not None else None
        if k4 is None:
            break
        q = p + sign * step / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
        if not (x[0] < q[0] < x[-1] and y[-1] < q[1] < y[0]) or not spd([q[::-1]])[0] >= min_speed:
            break
        pts_.append(q)
    return np.array(pts_)


lon_to = Transformer.from_crs("EPSG:4326", f"EPSG:{hit.epsg}", always_xy=True)
seed = np.array(lon_to.transform(LON, LAT))
up = trace(seed, -1)[::-1]
down = trace(seed, +1)
line = np.vstack([up, down[1:]])
# End the flowline where the static land-ice mask ends: beyond it is the ice melange in the fjord.
on_ice = landice.sel(x=xr.DataArray(line[:, 0]), y=xr.DataArray(line[:, 1]), method="nearest").values == 1
line = line[: np.flatnonzero(on_ice)[-1] + 1]
seg = np.hypot(*np.diff(line, axis=0).T)
dist_km = (np.r_[0, np.cumsum(seg)][-1] - np.r_[0, np.cumsum(seg)]) / 1000    # distance upstream of the end
section("flowline")
print(f"seed (lon {LON}, lat {LAT}) -> EPSG:{hit.epsg} x={seed[0]:.0f}, y={seed[1]:.0f}")
print(f"flowline: {len(line)} vertices, {dist_km[0]:.1f} km long; ice-mask edge at x={line[-1, 0]:.0f}, "
      f"y={line[-1, 1]:.0f}; upstream end at x={line[0, 0]:.0f}, y={line[0, 1]:.0f}")

POINTS_KM = {"A": 5, "B": 15, "C": 30, "D": 60}
pts = {k: line[np.argmin(np.abs(dist_km - d))] for k, d in POINTS_KM.items()}
for k, p in pts.items():
    print(f"point {k}: {POINTS_KM[k]:>2} km upstream, x={p[0]:.0f}, y={p[1]:.0f}, "
          f"v0={float(v0.sel(x=p[0], y=p[1], method='nearest')):,.0f} m/yr")

# Speed along the flowline in every annual composite.
lx, ly = xr.DataArray(line[:, 0], dims="d"), xr.DataArray(line[:, 1], dims="d")
prof = comp.v.sel(x=lx, y=ly, method="nearest").load()
prof_err = comp.v_error.sel(x=lx, y=ly, method="nearest").load()
prof = prof.where(prof_err < 0.25 * prof)                # drop poorly constrained annual values
cover = prof.notnull().mean("d")
print("annual composite coverage along the flowline (fraction of vertices with a value):")
print(pd.Series(cover.values.round(2), index=years).to_frame("coverage").T.to_string())

# --- Figure 1: climatological velocity map with flowline -------------------------------------
ne_zip = DATA / "ne_50m_admin_0_countries.zip"
if not ne_zip.exists():
    ne_zip.write_bytes(requests.get("https://naciscdn.org/naturalearth/50m/cultural/ne_50m_admin_0_countries.zip",
                                    timeout=300).content)
greenland = gpd.read_file(ne_zip).query("ADMIN == 'Greenland'").to_crs(hit.epsg)

PCOL = {"A": "#d7191c", "B": "#fdae61", "C": "#abd9e9", "D": "#2c7bb6"}
fig, ax = plt.subplots(figsize=(8, 6.2))
im = ax.imshow(v0.where(v0 > 0), extent=ext, cmap="magma", norm=LogNorm(10, 15000))
ax.plot(line[:, 0], line[:, 1], color="white", lw=1.6)
for k, p in pts.items():
    ax.plot(*p, "o", ms=7, mfc=PCOL[k], mec="black")
    ax.annotate(k, p, xytext=(5, 6), textcoords="offset points", color="white", fontweight="bold")
ax.set_xlabel("x (km, EPSG:3413)"); ax.set_ylabel("y (km, EPSG:3413)"); km_axes(ax)
cb = fig.colorbar(im, ax=ax, shrink=0.85, pad=0.02); cb.set_label("speed, 2014–2024 climatology (m/yr, log scale)")
ax.set_title("Jakobshavn Isbræ (Sermeq Kujalleq): ITS_LIVE surface speed and flowline")
ins = ax.inset_axes([0.0, 0.62, 0.26, 0.38])
greenland.plot(ax=ins, color="#d9d9d9", edgecolor="#555555", lw=0.4)
ins.set_xlim(greenland.total_bounds[[0, 2]]); ins.set_ylim(greenland.total_bounds[[1, 3]])
ins.add_patch(plt.Rectangle((x[0], y[-1]), x[-1] - x[0], y[0] - y[-1], fc="none", ec="#d7191c", lw=1.5))
ins.set_xticks([]); ins.set_yticks([]); ins.set_xlabel(""); ins.set_ylabel(""); ins.set_facecolor("white")
fig.savefig(OUT / "01-jakobshavn-velocity-map.png"); plt.close(fig)

# --- Figure 2: annual maps and change ---------------------------------------------------------
MAP_YEARS = [2013, 2017, 2024]
yr = {y_: comp.v.sel(time=f"{y_}-01-01").load() for y_ in MAP_YEARS}
zoom = [-200000, -150000, -2295000, -2255000]
fig, axs = plt.subplots(2, 2, figsize=(10, 7.4), sharex=True, sharey=True, layout="constrained")
for ax, y_ in zip(axs.flat, MAP_YEARS):
    im = ax.imshow(yr[y_], extent=ext, cmap="magma", norm=LogNorm(100, 15000))
    ax.set_title(f"{y_} annual mean speed")
fig.colorbar(im, ax=axs[0, :], shrink=0.8, label="speed (m/yr, log scale)")
im2 = axs[1, 1].imshow(yr[2017] - yr[2013], extent=ext, cmap="RdBu_r", norm=TwoSlopeNorm(0, -4000, 4000))
axs[1, 1].set_title("2017 minus 2013")
fig.colorbar(im2, ax=axs[1, :], shrink=0.8, label="speed change (m/yr)")
for ax in axs.flat:
    ax.set_xlim(zoom[:2]); ax.set_ylim(zoom[2:]); ax.plot(line[:, 0], line[:, 1], color="#bbbbbb", lw=0.8)
    km_axes(ax)
fig.supxlabel("x (km, EPSG:3413)"); fig.supylabel("y (km, EPSG:3413)")
fig.savefig(OUT / "02-annual-velocity-change.png"); plt.close(fig)

section("speed along the flowline in selected years (m/yr)")
rows = []
for y_ in [1985, 1995, 2000, 2003, 2007, 2013, 2015, 2017, 2018, 2019, 2021, 2022, 2024, 2025]:
    p = prof.sel(time=f"{y_}-01-01").values
    rows.append({"year": y_, **{f"{k} ({d} km)": p[np.argmin(np.abs(dist_km - d))] for k, d in POINTS_KM.items()},
                 "max on line": np.nanmax(p) if np.isfinite(p).any() else np.nan})
tab = pd.DataFrame(rows).set_index("year")
print(tab.round(0).to_string())

# --- Figure 3: flowline profiles ---------------------------------------------------------------
PROF_YEARS = [1995, 2000, 2003, 2007, 2013, 2017, 2019, 2024]
cmap = plt.get_cmap("viridis")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.2), width_ratios=[1.5, 1], layout="constrained")
for i, y_ in enumerate(PROF_YEARS):
    axs[0].plot(dist_km, prof.sel(time=f"{y_}-01-01"), color=cmap(i / (len(PROF_YEARS) - 1)), lw=1.3, label=str(y_))
axs[0].plot(dist_km, spd(line[:, ::-1]), color="black", lw=1, ls="--", label="2014–24 clim.")
axs[0].invert_xaxis(); axs[0].set_xlabel("distance upstream of the flowline end (km)")
axs[0].set_ylabel("annual mean speed (m/yr)"); axs[0].set_title("Speed along the flowline")
axs[0].legend(ncol=2, fontsize=7, frameon=False)
top = axs[0].get_ylim()[1]
for k, d in POINTS_KM.items():
    axs[0].axvline(d, color=PCOL[k], lw=0.8, ls=":")
    axs[0].text(d, top * 0.97, k, color="black", ha="center", va="top", fontweight="bold")
for k, d in POINTS_KM.items():
    s = prof.isel(d=np.argmin(np.abs(dist_km - d)))
    axs[1].plot(years, s, "o-", ms=3, color=PCOL[k], mec="black", mew=0.3, label=f"{k}: {d} km")
axs[1].set_xlabel("year"); axs[1].set_ylabel("annual mean speed (m/yr)"); axs[1].set_title("Annual speed at points A–D")
axs[1].legend(frameon=False, fontsize=8); axs[1].grid(alpha=0.3)
fig.savefig(OUT / "03-flowline-profiles.png"); plt.close(fig)

# --- Step 4: image-pair time series at points, with error filtering ----------------------------
section("image-pair time series at points")
VARS = ["v", "v_error", "date_dt", "mission_img1", "satellite_img1"]
series = {}
for k, p in pts.items():
    s = cube[VARS].sel(x=p[0], y=p[1], method="nearest").load().to_dataframe()
    s = s.dropna(subset=["v"]).sort_index()
    s["sensor"] = s.mission_img1.astype(str) + s.satellite_img1.astype(str)
    series[k] = s
    print(f"point {k}: {len(s):,} image pairs with a value, {s.index.min():%Y-%m-%d} to {s.index.max():%Y-%m-%d}")


def clean(s: pd.DataFrame, max_dt=48, max_rel_err=0.15, window="45D", n_mad=3.0) -> pd.DataFrame:
    """Flag short, well-constrained pairs whose speed is close to the local rolling median."""
    s = s[s.index >= "2014-01-01"].copy()
    s["keep_dt"] = s.date_dt <= max_dt
    s["keep_err"] = s.v_error <= max_rel_err * s.v
    s1 = s[s.keep_dt & s.keep_err]
    med = s1.v.rolling(window, center=True, min_periods=5).median()
    mad = (s1.v - med).abs().rolling(window, center=True, min_periods=5).median() * 1.4826
    ok = (s1.v - med).abs() <= n_mad * np.maximum(mad, 0.02 * med)
    s["keep_mad"] = False
    s.loc[s1.index[ok.values], "keep_mad"] = True
    s["keep"] = s.keep_dt & s.keep_err & s.keep_mad
    return s


filt = {k: clean(s) for k, s in series.items()}
summary = pd.DataFrame({k: {"pairs since 2014": len(f), "dt <= 48 d": int(f.keep_dt.sum()),
                            "+ error <= 15%": int((f.keep_dt & f.keep_err).sum()),
                            "+ rolling MAD": int(f.keep.sum()),
                            "kept %": 100 * f.keep.mean()} for k, f in filt.items()}).T
print(summary.round(1).to_string())
print("\nsensors among kept pairs at point B:")
print(filt["B"][filt["B"].keep].sensor.value_counts().to_string())

# --- Figure 4: what filtering does at point B ----------------------------------------------------
f = filt["B"]
fig, axs = plt.subplots(2, 1, figsize=(10, 6), sharex=True, layout="constrained", height_ratios=[2, 1])
rej = f[~f.keep]
axs[0].scatter(rej.index, rej.v, s=4, c="#bbbbbb", label=f"rejected ({len(rej):,})")
SENS_COL = {"S1A": "#2c7bb6", "S1B": "#abd9e9", "S2A": "#1a9641", "S2B": "#a6d96a", "L8": "#d7191c", "L9": "#fdae61", "L7": "#7b3294"}
kept = f[f.keep]
for sname, g in kept.groupby("sensor"):
    axs[0].scatter(g.index, g.v, s=5, c=SENS_COL.get(sname, "black"), label=f"{sname} ({len(g):,})")
axs[0].set_ylim(0, np.nanpercentile(f.v, 99.5) * 1.1)
axs[0].set_ylabel("speed (m/yr)")
axs[0].set_title(f"Point B ({POINTS_KM['B']} km upstream): every image pair since 2014, kept and rejected")
axs[0].legend(ncol=4, fontsize=7, markerscale=2.5, frameon=False, loc="upper left")
axs[1].scatter(f.index, f.date_dt, s=3, c=np.where(f.keep, "#2c7bb6", "#bbbbbb"))
axs[1].axhline(48, color="#d7191c", lw=1, ls="--")
axs[1].text(f.index.min(), 56, "48-day limit", color="#d7191c", fontsize=8)
axs[1].set_yscale("log"); axs[1].set_ylabel("days between images")
axs[1].xaxis.set_major_locator(mdates.YearLocator(2))
fig.savefig(OUT / "04-image-pair-filtering.png"); plt.close(fig)

# --- Figure 5: seasonal and multi-year speed -----------------------------------------------------
section("seasonal cycle and multi-year change (filtered pairs, 2014 onwards)")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.4), width_ratios=[1.7, 1], layout="constrained")
seas_rows = []
for k, f in filt.items():
    g = f[f.keep]
    axs[0].scatter(g.index, g.v, s=2, color=PCOL[k], alpha=0.25)
    m = g.v.resample("MS").median()
    axs[0].plot(m.index + pd.Timedelta(days=14), m, color=PCOL[k], lw=1.4, label=f"{k}: {POINTS_KM[k]} km")
    # Seasonal cycle: each monthly median relative to a centred 12-month mean, so the multi-year trend drops out,
    # then the median of those departures for each calendar month.
    trend = m.rolling(13, center=True, min_periods=7).mean()
    rel = (100 * (m / trend - 1)).groupby(m.index.month).median()
    axs[1].plot(rel.index, rel, "o-", color=PCOL[k], mec="black", mew=0.3, ms=4, label=k)
    ann = g.v.groupby(g.index.year).median()
    seas_rows.append({"point": k, "km": POINTS_KM[k], "pairs": len(g), "peak month": int(rel.idxmax()),
                      "min month": int(rel.idxmin()), "seasonal range %": rel.max() - rel.min(),
                      "2014 median": ann.get(2014, np.nan), "2017 median": ann.get(2017, np.nan),
                      "2024 median": ann.get(2024, np.nan)})
axs[0].set_ylabel("speed (m/yr)"); axs[0].set_title("Filtered image pairs (dots) and monthly medians (lines)")
axs[0].legend(frameon=False, fontsize=8, ncol=4, loc="upper left"); axs[0].grid(alpha=0.3)
axs[0].set_ylim(0, None)
axs[1].axhline(0, color="black", lw=0.6)
axs[1].set_xticks(range(1, 13), list("JFMAMJJASOND")); axs[1].set_ylabel("departure from the point's mean (%)")
axs[1].set_title("Median seasonal cycle, 2014–2025"); axs[1].legend(frameon=False, fontsize=8); axs[1].grid(alpha=0.3)
fig.savefig(OUT / "05-seasonal-multiyear-speed.png"); plt.close(fig)
seas = pd.DataFrame(seas_rows).set_index("point")
print(seas.round(1).to_string())

# Cross-check: calendar-year medians of the filtered pairs against the annual composites.
b = filt["B"][filt["B"].keep]
bm = b.v.resample("MS").median()
chk = pd.DataFrame({"pairs median": b.v.groupby(b.index.year).median(),
                    "mean of monthly medians": bm.groupby(bm.index.year).mean(),
                    "months": bm.groupby(bm.index.year).count(),
                    "composite": prof.isel(d=np.argmin(np.abs(dist_km - POINTS_KM["B"]))).to_series().set_axis(years)})
print("\npoint B, calendar-year median of filtered pairs vs annual composite (m/yr):")
print(chk.dropna().round(0).astype(int).T.to_string())

print("\nITS_LIVE climatological seasonal fit (2014-2024) at the points:")
for k, p in pts.items():
    a = float(v_amp.sel(x=p[0], y=p[1], method="nearest"))
    ph = float(v_phase.sel(x=p[0], y=p[1], method="nearest"))
    vv = float(v0.sel(x=p[0], y=p[1], method="nearest"))
    print(f"  {k}: v0 {vv:,.0f} m/yr, amplitude {a:,.0f} m/yr ({100 * a / vv:.0f}%), peak on day {ph:.0f} "
          f"(~{(pd.Timestamp('2018-01-01') + pd.Timedelta(days=ph - 1)):%d %b})")

# --- Step 5: strain rates and flow direction -------------------------------------------------------
section("strain rates (2014-2024 climatology)")
dx, dy = float(x[1] - x[0]), float(y[1] - y[0])            # dy is negative: y decreases down the rows
# Derivatives amplify noise, so smooth more strongly first (sigma 5 pixels = 600 m, a few ice thicknesses).
w = ndimage.gaussian_filter(np.isfinite(vx0.values).astype(float), 5)
u = ndimage.gaussian_filter(np.nan_to_num(vx0.values), 5) / np.maximum(w, 1e-6)
v = ndimage.gaussian_filter(np.nan_to_num(vy0.values), 5) / np.maximum(w, 1e-6)
mask = (w > 0.9) & (landice.values == 1)
dudy, dudx = np.gradient(u, dy, dx)
dvdy, dvdx = np.gradient(v, dy, dx)
exy = 0.5 * (dudy + dvdx)
speed = np.hypot(u, v)
e_long = (u ** 2 * dudx + u * v * (dudy + dvdx) + v ** 2 * dvdy) / speed ** 2     # along-flow strain rate, 1/yr
e_eff = np.sqrt(0.5 * (dudx ** 2 + dvdy ** 2 + (dudx + dvdy) ** 2) + exy ** 2)     # effective strain rate, 1/yr (ezz from incompressibility)
for arr in (e_long, e_eff):
    arr[~mask | ~(speed >= 50)] = np.nan
along = RegularGridInterpolator(grid, e_long[::-1], bounds_error=False, fill_value=np.nan)(line[:, ::-1])
eff_line = RegularGridInterpolator(grid, e_eff[::-1], bounds_error=False, fill_value=np.nan)(line[:, ::-1])
for k, d in POINTS_KM.items():
    i = np.argmin(np.abs(dist_km - d))
    print(f"point {k}: longitudinal strain rate {along[i]:+.3f} /yr, effective {eff_line[i]:.3f} /yr")
print(f"max longitudinal stretching on the flowline: {np.nanmax(along):.2f} /yr at {dist_km[np.nanargmax(along)]:.1f} km")
print(f"max compression on the flowline: {np.nanmin(along):.2f} /yr at {dist_km[np.nanargmin(along)]:.1f} km")

fig, axs = plt.subplots(1, 2, figsize=(11, 4.4), sharex=True, sharey=True, layout="constrained")
im = axs[0].imshow(e_long, extent=ext, cmap="RdBu_r", norm=TwoSlopeNorm(0, -0.6, 0.6))
fig.colorbar(im, ax=axs[0], shrink=0.85, label="longitudinal strain rate (1/yr)\nred = stretching, blue = compression")
axs[0].set_title("Strain rate along flow")
im = axs[1].imshow(np.where(mask, speed, np.nan), extent=ext, cmap="magma", norm=LogNorm(50, 15000))
fig.colorbar(im, ax=axs[1], shrink=0.85, label="speed (m/yr, log scale)")
st = 25
X, Y = np.meshgrid(x[::st], y[::st])
U, V = u[::st, ::st], v[::st, ::st]
S = np.hypot(U, V)
q = (S > 100) & mask[::st, ::st]
axs[1].quiver(X[q], Y[q], (U / S)[q], (V / S)[q], color="white", scale=32, width=0.003, headwidth=4)
axs[1].set_title("Flow direction (unit arrows every 3 km)")
for ax in axs:
    ax.plot(line[:, 0], line[:, 1], color="black", lw=0.8)
    ax.set_xlim(-205000, -120000); ax.set_ylim(-2300000, -2245000); km_axes(ax)
fig.supxlabel("x (km, EPSG:3413)"); fig.supylabel("y (km, EPSG:3413)")
fig.savefig(OUT / "06-strain-rate-flow-direction.png"); plt.close(fig)

# --- Step 6: save results -------------------------------------------------------------------------
out = pd.DataFrame({"dist_km": dist_km, "x": line[:, 0], "y": line[:, 1], "v0": spd(line[:, ::-1]), "e_long": along})
out = out.join(pd.DataFrame(prof.values.T, columns=[f"v_{y_}" for y_ in years]))
out.to_csv(OUT / "flowline_profiles.csv", index=False, float_format="%.3f")
pd.concat({k: f[f.keep][["v", "v_error", "date_dt", "sensor"]] for k, f in filt.items()}, names=["point"]).to_csv(
    OUT / "point_timeseries_filtered.csv", float_format="%.2f")
print("\nwrote output/flowline_profiles.csv and output/point_timeseries_filtered.csv")