"""Cinematic 3D movie of the FKESSLER baroclinic-wave life cycle (days 0-10).

Software renderer, no GPU required:
  * globe        - per-pixel ray-cast unit sphere, Lambert + rim shading, texture
                   from the (amplitude-stretched) surface-pressure anomaly
  * jet streams  - Lagrangian particles advected through the real 232 hPa wind,
                   drawn as fading glow trails coloured by wind speed
  * rising air   - strongest ascending columns, drawn as vertical streaks whose
                   length tracks true -omega
  * condensation - cloud-liquid puffs at the true altitude of each column max,
                   colour ramping toward the Kessler autoconversion threshold
  * rain         - surface splashes sized by rain rate

All physics fields are the model's own daily means, linearly interpolated in
time for smooth motion. Only the pressure-anomaly TEXTURE is amplitude-stretched
(signed |x|^0.35); every readout prints true values.

Usage:  python render_movie.py [--preview N] [--fast]
"""
import argparse
import os
import sys

import numpy as np
import xarray as xr
from PIL import Image, ImageDraw, ImageFont
import matplotlib
matplotlib.use("Agg")
from matplotlib import cm, colors as mcolors

# ----------------------------------------------------------------- config
HIST_GLOB = "hpc/1f737b81-34c1-4b1b-ae39-9db1cbfd439a"
NDAY = 11                    # days 0..10 inclusive
FPS = 24
DURATION = 60.0              # seconds of movie
CL_THRESHOLD = 1000.0        # mg/kg, Kessler autoconversion
ALPHA_STRETCH = 0.35
OUT = "fkessler_physics_movie.mp4"

ap = argparse.ArgumentParser()
ap.add_argument("--preview", type=int, default=0, help="render N sample frames only")
ap.add_argument("--fast", action="store_true", help="half resolution, 12 fps")
ap.add_argument("--frames-dir", default="frames")
args = ap.parse_args()

W, H = (1600, 900) if not args.fast else (800, 450)
if args.fast:
    FPS = 12
NFRAME = int(DURATION * FPS)
SS = 1.0                     # sprite scale factor
FONT_DIR = os.path.join(os.path.dirname(matplotlib.__file__), "mpl-data", "fonts", "ttf")


UI = H / 900.0          # UI scale so the overlay matches any render size


def font(sz, bold=False):
    name = "DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf"
    return ImageFont.truetype(os.path.join(FONT_DIR, name), max(int(round(sz * UI)), 8))


def S(v):
    """scale a pixel offset by the UI factor"""
    return int(round(v * UI))


# ----------------------------------------------------------------- load data
import glob
files = sorted(glob.glob(os.path.join(HIST_GLOB, "*.cam.h0.*.nc")))
if not files:
    sys.exit(f"no history files under {HIST_GLOB}")
ds = xr.concat([xr.open_dataset(f, decode_timedelta=False) for f in files],
               dim="time", data_vars="all").isel(time=slice(0, NDAY))
lat = ds.lat.values
lon = ds.lon.values
lev = ds.lev.values
NLAT, NLON, NLEV = lat.size, lon.size, lev.size
kjet = int(np.abs(lev - 230).argmin())

PS = (ds.PS.values / 100.0).astype(np.float32)
OM = ds.OMEGA.values.astype(np.float32)
CLD = (ds.CLDLIQ.values * 1e6).astype(np.float32)
PRECT = (ds.PRECT.values * 86400 * 1000.0).astype(np.float32)
UJ = ds.U.values[:, kjet].astype(np.float32)
VJ = ds.V.values[:, kjet].astype(np.float32)

psa = PS - PS.mean(axis=(1, 2), keepdims=True)
PSA_DISP = np.sign(psa) * np.abs(psa) ** ALPHA_STRETCH
DISP_MAX = float(np.abs(PSA_DISP).max())
TRUE_AMP = np.abs(psa).max(axis=(1, 2))
CL_PEAK = CLD.max(axis=(1, 2, 3))
PR_PEAK = PRECT.max(axis=(1, 2))
ASC_PEAK = (-OM.min(axis=(1, 2, 3)))
NWET = (PRECT > 0.05).sum(axis=(1, 2))
RAIN_START = int(np.argmax(NWET > 0))          # first day with rain
print(f"days={NDAY} rain onset day {RAIN_START}; cloud peak {CL_PEAK.max():.0f} mg/kg "
      f"({100*CL_PEAK.max()/CL_THRESHOLD:.0f}% of threshold); wave {TRUE_AMP[1]:.2f}->{TRUE_AMP[-1]:.1f} hPa")


def lerp_day(arr, d):
    """linear interpolation of a (time,...) array at float day d"""
    d = float(np.clip(d, 0, NDAY - 1))
    i0 = int(np.floor(d))
    i1 = min(i0 + 1, NDAY - 1)
    f = d - i0
    if f == 0.0:
        return arr[i0]
    return (1.0 - f) * arr[i0] + f * arr[i1]


# ----------------------------------------------------------------- colour maps
def ramp(cmap_name, n=256):
    return (np.asarray(matplotlib.colormaps[cmap_name](np.linspace(0, 1, n)))[:, :3] * 255).astype(np.float32)


# Diverging map for the pressure anomaly with a DARK ocean-slate centre. A
# white-centred map (RdBu_r) blows the globe out to a featureless white ball on
# the low-anomaly early days, where |PS'| is near zero almost everywhere.
WAVE_LUT = (np.asarray(mcolors.LinearSegmentedColormap.from_list(
    "wave", ["#5ec8f2", "#2b7fc4", "#1b3f68", "#16233a", "#4a2033", "#a33b3f", "#e8703c"]
)(np.linspace(0, 1, 512)))[:, :3] * 255).astype(np.float32)
CLOUD_LUT = (np.asarray(mcolors.LinearSegmentedColormap.from_list(
    "cld", ["#f4f8fc", "#dde8f3", "#9fb8d4", "#5a7fb0", "#8e44ad"])(np.linspace(0, 1, 256)))[:, :3] * 255).astype(np.float32)
ASC_LUT = (np.asarray(mcolors.LinearSegmentedColormap.from_list(
    "asc", ["#ffe2b0", "#ffb347", "#ff7518", "#c62828"])(np.linspace(0, 1, 256)))[:, :3] * 255).astype(np.float32)
JET_LUT = (np.asarray(mcolors.LinearSegmentedColormap.from_list(
    "jetw", ["#2f3550", "#5d6a9e", "#a8b6e0", "#eef3ff"])(np.linspace(0, 1, 256)))[:, :3] * 255).astype(np.float32)


def lut_lookup(lut, vals, vmin, vmax):
    t = np.clip((np.asarray(vals, np.float32) - vmin) / max(vmax - vmin, 1e-9), 0, 1)
    idx = (t * (lut.shape[0] - 1)).astype(np.int32)
    return lut[idx]


# ----------------------------------------------------------------- geometry
def sph2cart(latd, lond, r=1.0):
    la = np.deg2rad(latd)
    lo = np.deg2rad(lond)
    return np.stack([r * np.cos(la) * np.cos(lo),
                     r * np.cos(la) * np.sin(lo),
                     r * np.sin(la)], axis=-1)


def norml(v):
    return v / np.maximum(np.linalg.norm(v, axis=-1, keepdims=True), 1e-12)


def altitude_of(press):
    """map pressure (hPa) to render radius"""
    return 1.02 + 0.34 * (1.0 - np.log10(np.clip(press, 5, 1000)) / 3.0)


class Camera:
    def __init__(self, eye, target=(0, 0, 0), up=(0, 0, 1), fov_deg=34.0):
        self.eye = np.asarray(eye, np.float64)
        self.f = norml(np.asarray(target, np.float64) - self.eye)
        up = np.asarray(up, np.float64)
        self.r = norml(np.cross(self.f, up))
        self.u = np.cross(self.r, self.f)
        self.tan_half = np.tan(np.deg2rad(fov_deg) / 2.0)
        self.aspect = W / H

    def project(self, pts):
        v = np.asarray(pts, np.float64) - self.eye
        z = v @ self.f
        x = v @ self.r
        y = v @ self.u
        z_safe = np.where(np.abs(z) < 1e-6, 1e-6, z)
        ndc_x = (x / z_safe) / (self.tan_half * self.aspect)
        ndc_y = (y / z_safe) / self.tan_half
        sx = (ndc_x + 1.0) * 0.5 * W
        sy = (1.0 - ndc_y) * 0.5 * H
        return sx, sy, z

    def hidden(self, pts):
        """True where the unit sphere occludes the point"""
        p = np.asarray(pts, np.float64)
        v = p - self.eye
        dist = np.linalg.norm(v, axis=-1)
        d = v / np.maximum(dist, 1e-12)[:, None]
        b = d @ self.eye
        c = float(self.eye @ self.eye) - 1.0
        disc = b * b - c
        t = np.where(disc > 0, -b - np.sqrt(np.maximum(disc, 0)), np.inf)
        return (disc > 0) & (t > 1e-4) & (t < dist - 1e-3)

    def px_per_unit(self, z):
        return (0.5 * H) / (self.tan_half * np.maximum(z, 1e-6))


# pixel ray grid (rebuilt per frame from camera basis)
_px = (np.arange(W, dtype=np.float32) + 0.5) / W * 2.0 - 1.0
_py = 1.0 - (np.arange(H, dtype=np.float32) + 0.5) / H * 2.0
NDCX, NDCY = np.meshgrid(_px, _py)


def render_globe(cam, tex_rgb, light_dir):
    """ray-cast the unit sphere; returns (rgb float32 HxWx3, hit mask)"""
    dirs = (cam.f[None, None, :]
            + (NDCX * cam.tan_half * cam.aspect)[..., None] * cam.r[None, None, :]
            + (NDCY * cam.tan_half)[..., None] * cam.u[None, None, :])
    dirs /= np.maximum(np.linalg.norm(dirs, axis=-1, keepdims=True), 1e-12)
    b = dirs @ cam.eye
    c = float(cam.eye @ cam.eye) - 1.0
    disc = b * b - c
    hit = disc > 0.0
    out = np.zeros((H, W, 3), np.float32)

    # atmosphere glow for rays that miss (impact parameter just outside R=1)
    perp2 = np.maximum(c - b * b, 0.0)
    perp = np.sqrt(perp2)
    glow = np.exp(-np.clip(perp - 1.0, 0, None) * 11.0)
    glow = np.where(~hit, glow, 0.0)
    out += glow[..., None] * np.array([26.0, 48.0, 96.0], np.float32)

    if hit.any():
        t = -b - np.sqrt(np.maximum(disc, 0.0))
        t = np.where(t > 0, t, np.inf)
        vis = hit & np.isfinite(t)
        p = cam.eye[None, None, :] + t[..., None] * dirs
        n = p  # unit sphere -> normal == position
        latq = np.degrees(np.arcsin(np.clip(n[..., 2], -1, 1)))
        lonq = np.degrees(np.arctan2(n[..., 1], n[..., 0]))
        # bilinear texture sample (wrap lon, clamp lat)
        fi = np.interp(latq, lat, np.arange(NLAT, dtype=np.float32))
        i0 = np.clip(np.floor(fi).astype(np.int32), 0, NLAT - 2)
        wi = (fi - i0)[..., None]
        fj = np.mod(lonq, 360.0) / (360.0 / NLON)
        j0 = np.mod(np.floor(fj).astype(np.int32), NLON)
        j1 = np.mod(j0 + 1, NLON)
        wj = (fj - np.floor(fj))[..., None]
        tex = ((1 - wi) * ((1 - wj) * tex_rgb[i0, j0] + wj * tex_rgb[i0, j1])
               + wi * ((1 - wj) * tex_rgb[i0 + 1, j0] + wj * tex_rgb[i0 + 1, j1]))
        lam = np.clip(n @ light_dir, 0, 1)[..., None]
        view = -dirs
        rim = (1.0 - np.clip(np.sum(n * view, axis=-1), 0, 1))[..., None] ** 2.6
        # deeper terminator so the sphere reads as a lit 3-D body, plus a soft
        # specular highlight and a limb-darkening falloff
        shade = 0.30 + 0.95 * lam ** 0.9
        half = norml(light_dir + view)
        spec = np.clip(np.sum(n * half, axis=-1), 0, 1)[..., None] ** 60.0
        col = (tex * shade
               + rim * np.array([30.0, 60.0, 124.0], np.float32)
               + spec * np.array([26.0, 30.0, 38.0], np.float32))
        out = np.where(vis[..., None], col, out)
    return out, hit


# ----------------------------------------------------------------- glow splatting
_STAMP_CACHE = {}


def stamp(r):
    r = int(r)
    if r not in _STAMP_CACHE:
        d = np.arange(-r, r + 1, dtype=np.float32)
        dy, dx = np.meshgrid(d, d, indexing="ij")
        rr = np.sqrt(dx * dx + dy * dy) / max(r, 1)
        s = np.clip(1.0 - rr, 0, 1) ** 1.7
        _STAMP_CACHE[r] = s.astype(np.float32)
    return _STAMP_CACHE[r]


def splat(acc_rgb, acc_a, xs, ys, rs, cols, alphas):
    """additive soft-sprite splatting into colour+alpha accumulation buffers"""
    if len(xs) == 0:
        return
    xs = np.asarray(xs, np.float32)
    ys = np.asarray(ys, np.float32)
    rs = np.clip(np.round(np.asarray(rs)), 1, 26).astype(np.int32)
    cols = np.asarray(cols, np.float32).reshape(-1, 3)
    alphas = np.broadcast_to(np.asarray(alphas, np.float32), xs.shape).astype(np.float32)
    ix = np.round(xs).astype(np.int32)
    iy = np.round(ys).astype(np.int32)
    for r in np.unique(rs):
        m = rs == r
        st = stamp(r)
        k = st.shape[0]
        off = np.arange(-r, r + 1)
        oy, ox = np.meshgrid(off, off, indexing="ij")
        px = ix[m][:, None, None] + ox[None]
        py = iy[m][:, None, None] + oy[None]
        good = (px >= 0) & (px < W) & (py >= 0) & (py < H)
        wgt = alphas[m][:, None, None] * st[None]
        cc = cols[m][:, None, None, :] * np.ones((1, k, k, 1), np.float32)
        pxg, pyg, wg = px[good], py[good], wgt[good]
        np.add.at(acc_a, (pyg, pxg), wg)
        np.add.at(acc_rgb, (pyg, pxg), wg[:, None] * cc[good])


def composite(base, acc_rgb, acc_a):
    a = np.clip(acc_a, 0, 1)[..., None]
    col = acc_rgb / np.maximum(acc_a, 1e-6)[..., None]
    return base * (1 - a) + col * a


# ----------------------------------------------------------------- jet particles
RNG = np.random.default_rng(7)
NPART = 620
TRAIL = 16
R_JET = 1.155


def seed_particles(n):
    la = np.concatenate([RNG.uniform(26, 62, n // 2), RNG.uniform(-62, -26, n - n // 2)])
    lo = RNG.uniform(0, 360, n)
    return la.astype(np.float32), lo.astype(np.float32)


# fixed random subsample of the near-uniform tropical background cloud layer
# (a grid stride would show up as a visible lattice on the rendered sphere)
BG_KEEP = np.random.default_rng(11).random((NLAT, NLON)) < 0.16

P_LAT, P_LON = seed_particles(NPART)
P_AGE = RNG.uniform(0, 1, NPART).astype(np.float32)
P_HIST_LAT = np.tile(P_LAT[:, None], (1, TRAIL))
P_HIST_LON = np.tile(P_LON[:, None], (1, TRAIL))


def bilin_field(F, latq, lonq):
    fi = np.interp(latq, lat, np.arange(NLAT, dtype=np.float32))
    i0 = np.clip(np.floor(fi).astype(np.int32), 0, NLAT - 2)
    wi = fi - i0
    fj = np.mod(lonq, 360.0) / (360.0 / NLON)
    j0 = np.mod(np.floor(fj).astype(np.int32), NLON)
    j1 = np.mod(j0 + 1, NLON)
    wj = fj - np.floor(fj)
    return ((1 - wi) * ((1 - wj) * F[i0, j0] + wj * F[i0, j1])
            + wi * ((1 - wj) * F[i0 + 1, j0] + wj * F[i0 + 1, j1]))


def advance_particles(U, V, dt_s):
    """advect the jet tracers through the real wind field"""
    global P_LAT, P_LON, P_AGE, P_HIST_LAT, P_HIST_LON
    Rearth = 6.371e6
    u = bilin_field(U, P_LAT, P_LON)
    v = bilin_field(V, P_LAT, P_LON)
    coslat = np.maximum(np.cos(np.deg2rad(P_LAT)), 0.15)
    P_LAT = np.clip(P_LAT + np.degrees(v * dt_s / Rearth), -87.0, 87.0)
    P_LON = np.mod(P_LON + np.degrees(u * dt_s / (Rearth * coslat)), 360.0)
    P_HIST_LAT = np.roll(P_HIST_LAT, 1, axis=1)
    P_HIST_LON = np.roll(P_HIST_LON, 1, axis=1)
    P_HIST_LAT[:, 0] = P_LAT
    P_HIST_LON[:, 0] = P_LON
    P_AGE += dt_s / 9.0e5
    dead = P_AGE > 1.0
    if dead.any():
        nl, no = seed_particles(int(dead.sum()))
        P_LAT[dead], P_LON[dead] = nl, no
        P_AGE[dead] = 0.0
        P_HIST_LAT[dead] = nl[:, None]
        P_HIST_LON[dead] = no[:, None]
    return np.sqrt(u * u + v * v)


# ----------------------------------------------------------------- camera path
def smooth(t):
    t = np.clip(t, 0, 1)
    return t * t * (3 - 2 * t)


def ease_seq(x, keys):
    """piecewise-smooth interpolation through (pos, value) keyframes"""
    xs = np.array([k[0] for k in keys], float)
    vs = np.array([k[1] for k in keys], float)
    if x <= xs[0]:
        return vs[0]
    if x >= xs[-1]:
        return vs[-1]
    i = int(np.searchsorted(xs, x) - 1)
    f = smooth((x - xs[i]) / (xs[i + 1] - xs[i]))
    return vs[i] * (1 - f) + vs[i + 1] * f


# day as a function of movie progress: linger on the nonlinear phase
DAY_KEYS = [(0.00, 0.0), (0.08, 0.0), (0.30, 4.0), (0.50, 6.5),
            (0.66, 8.0), (0.82, 9.0), (0.95, 10.0), (1.00, 10.0)]
# Elevation stays over the northern storm track (where the wave grows and rains)
EL_KEYS = [(0.0, 12.0), (0.30, 26.0), (0.62, 36.0), (0.85, 32.0), (1.0, 20.0)]
# Distance keeps the globe and its marker shell (max radius ~1.36) inside frame
DIST_KEYS = [(0.0, 4.95), (0.12, 4.55), (0.45, 3.95), (0.70, 3.55), (0.88, 3.75), (1.0, 4.35)]

# The wave itself propagates eastward ~12 deg/day. Track the longitude of the
# dominant NH midlatitude pressure anomaly so the growing system stays facing
# the camera; this makes the orbit physically motivated rather than arbitrary.
_nh = (lat > 25) & (lat < 75)
_flon = []
for _t in range(NDAY):
    _f = np.abs(psa[_t][_nh])
    _i, _j = np.unravel_index(int(_f.argmax()), _f.shape)
    _flon.append(float(lon[_j]))
FEAT_LON = np.degrees(np.unwrap(np.radians(np.array(_flon))))
print(f"wave longitude {FEAT_LON[0]:.0f} -> {FEAT_LON[-1]:.0f} deg "
      f"({(FEAT_LON[-1]-FEAT_LON[0])/(NDAY-1):.0f} deg/day eastward)")


def azimuth_at(day, x):
    """point the camera at the tracked wave, with a slow extra drift for parallax"""
    return float(np.interp(day, np.arange(NDAY), FEAT_LON)) + 26.0 * x - 12.0

STAGES = [
    (0.00, "An atmosphere at rest", "a 0.8 hPa ripple on two westerly jets"),
    (0.16, "Baroclinic instability", "the ripple feeds on the jet's temperature contrast"),
    (0.34, "Exponential growth", "wave amplitude e-folds about every 1.7 days"),
    (0.52, "Warm moist air rises", "ascent strengthens by two orders of magnitude"),
    (0.64, "Cooling and condensation", "cloud liquid water appears where air is lifted"),
    (0.76, "Nearing the rain threshold", "Kessler autoconversion needs about 1 g/kg"),
    (0.86, "Rain falls beneath the updraughts", "cyclones and fronts in the storm track"),
]


def stage_at(x):
    s = STAGES[0]
    for k in STAGES:
        if x >= k[0]:
            s = k
    return s


# ----------------------------------------------------------------- overlay
F_TITLE = font(30, True)
F_SUB = font(17)
F_STAGE = font(26, True)
F_STAGE2 = font(17)
F_LBL = font(15, True)
F_SM = font(14)
F_TINY = font(12)
F_BIG = font(44, True)

KEY = [("rising air", "#ff8c2b"), ("condensation / cloud water", "#cfe0f2"),
       ("cloud-water threshold reached", "#8e44ad"), ("precipitation", "#5b9bff"),
       ("jet-stream wind", "#a8b6e0")]


def draw_overlay(img, x, day, cam, feats):
    d = ImageDraw.Draw(img, "RGBA")
    PAD = S(26)
    # ---- top banner
    d.rectangle([0, 0, W, S(88)], fill=(8, 10, 18, 208))
    d.text((PAD, S(14)), "A baroclinic wave grows and starts to rain", font=F_TITLE, fill=(242, 245, 252))
    d.text((PAD + S(2), S(56)), "CESM 2.1.5 · CAM with Kessler warm-rain physics · FKESSLER f19_f19_mg17",
           font=F_SUB, fill=(150, 162, 190))
    # ---- day readout (right of banner)
    dtxt = f"day {day:4.1f}"
    dw = d.textlength(dtxt, font=F_BIG)
    d.text((W - PAD - dw, S(20)), dtxt, font=F_BIG, fill=(236, 240, 250))

    # ---- stage caption
    st = stage_at(x)
    d.rectangle([0, H - S(112), W, H], fill=(8, 10, 18, 208))
    d.text((PAD, H - S(100)), st[1], font=F_STAGE, fill=(240, 244, 252))
    d.text((PAD + S(2), H - S(66)), st[2], font=F_STAGE2, fill=(152, 166, 196))

    # ---- timeline (bottom-right of the caption bar)
    x0 = int(W * 0.52)
    x1 = W - PAD
    yb = H - S(74)
    d.line([x0, yb, x1, yb], fill=(70, 78, 104), width=max(S(3), 1))
    for dd in range(NDAY):
        px = x0 + (x1 - x0) * dd / (NDAY - 1)
        d.line([px, yb - S(5), px, yb + S(5)], fill=(96, 106, 138), width=max(S(2), 1))
        if dd % 2 == 0:
            d.text((px - S(4), yb + S(8)), str(dd), font=F_TINY, fill=(120, 132, 164))
    if RAIN_START <= NDAY - 1:
        prx = x0 + (x1 - x0) * RAIN_START / (NDAY - 1)
        d.line([prx, yb - S(16), prx, yb - S(6)], fill=(91, 155, 255), width=max(S(2), 1))
        d.text((prx - S(28), yb - S(32)), "rain begins", font=F_TINY, fill=(140, 176, 255))
    pnow = x0 + (x1 - x0) * day / (NDAY - 1)
    d.line([x0, yb, pnow, yb], fill=(200, 90, 70), width=max(S(3), 1))
    d.ellipse([pnow - S(6), yb - S(6), pnow + S(6), yb + S(6)], fill=(224, 106, 82))
    d.text((x0, yb - S(30)), "model day", font=F_TINY, fill=(120, 132, 164))

    # ---- key (upper left, under banner)
    ky = S(104)
    for name, hexc in KEY:
        rgb = tuple(int(hexc[i:i + 2], 16) for i in (1, 3, 5))
        d.ellipse([PAD, ky + S(4), PAD + S(11), ky + S(15)], fill=rgb + (255,))
        d.text((PAD + S(20), ky), name, font=F_SM, fill=(196, 206, 228))
        ky += S(23)

    # ---- cloud-water gauge vs the Kessler threshold (right edge)
    gw, gh = S(24), S(210)
    gx = W - PAD - gw
    gy = S(158)
    cl_now = feats["cl_peak"]
    frac = float(np.clip(cl_now / CL_THRESHOLD, 0, 1))
    lab = "cloud liquid water"
    lw = d.textlength(lab, font=F_LBL)
    d.text((gx + gw - lw, gy - S(52)), lab, font=F_LBL, fill=(206, 224, 242))
    fh = int(gh * frac)
    for i in range(max(fh, 1)):
        t = i / max(gh - 1, 1)
        rgb = tuple(int(v) for v in lut_lookup(CLOUD_LUT, [t], 0, 1)[0])
        d.line([gx + S(2), gy + gh - i, gx + gw - S(2), gy + gh - i], fill=rgb)
    d.rectangle([gx, gy, gx + gw, gy + gh], outline=(96, 108, 140), width=max(S(2), 1))
    # threshold line at the top of the gauge
    d.line([gx - S(8), gy, gx + gw, gy], fill=(184, 124, 224), width=max(S(3), 1))
    for txt, dy, f_, col in (("1000 mg/kg rain threshold", -S(34), F_SM, (198, 148, 232)),
                             ("not reached in days 0–10", -S(19), F_TINY, (160, 120, 200))):
        tw = d.textlength(txt, font=f_)
        d.text((gx + gw - tw, gy + dy), txt, font=f_, fill=col)
    # a tick at the fill level (no moving text, which would collide with the
    # fixed threshold label as the gauge fills)
    d.line([gx - S(7), gy + gh - fh, gx, gy + gh - fh], fill=(226, 236, 248), width=max(S(2), 1))
    vtxt = f"{cl_now:,.0f} mg/kg"
    tw = d.textlength(vtxt, font=F_LBL)
    d.text((gx + gw - tw, gy + gh + S(8)), vtxt, font=F_LBL, fill=(226, 236, 248))
    ptxt = f"{100*frac:.0f}% of threshold"
    tw = d.textlength(ptxt, font=F_TINY)
    d.text((gx + gw - tw, gy + gh + S(28)), ptxt, font=F_TINY, fill=(170, 186, 214))

    # ---- live physics numbers (right edge, under the gauge)
    by = gy + gh + S(56)
    rows = [("peak ascent", f"{feats['asc_peak']:.3f} Pa/s", (255, 150, 70)),
            ("wave amplitude", f"{feats['amp']:.2f} hPa", (226, 120, 100)),
            ("peak rain", f"{feats['pr_peak']:.1f} mm/day", (110, 170, 255)),
            ("raining cells", f"{feats['nwet']:d}", (110, 170, 255))]
    for nm, val, rgb in rows:
        tw = d.textlength(nm, font=F_TINY)
        d.text((W - PAD - tw, by), nm, font=F_TINY, fill=(150, 162, 190))
        tw = d.textlength(val, font=F_LBL)
        d.text((W - PAD - tw, by + S(15)), val, font=F_LBL, fill=rgb)
        by += S(40)

    # ---- data-driven callouts, de-overlapped: each label box is pushed
    # vertically until it clears every box already placed this frame
    placed = []
    for sx, sy, text, hexc in feats["callouts"]:
        rgb = tuple(int(hexc[i:i + 2], 16) for i in (1, 3, 5))
        tw = d.textlength(text, font=F_LBL)
        bh = S(26)
        tx = min(max(sx + S(48), S(180)), W - S(300) - tw)
        ty = min(max(sy - S(42), S(200)), H - S(190))
        for _ in range(24):
            box = (tx - S(7), ty - S(4), tx + tw + S(9), ty + S(21))
            clash = any(not (box[2] < o[0] or box[0] > o[2] or box[3] < o[1] or box[1] > o[3])
                        for o in placed)
            if not clash:
                break
            ty += bh
            if ty > H - S(190):        # ran out of room below: go above the anchor
                ty = max(sy - S(42) - bh * (len(placed) + 1), S(196))
        placed.append((tx - S(7), ty - S(4), tx + tw + S(9), ty + S(21)))
        d.line([sx, sy, tx - S(6), ty + S(9)], fill=rgb + (215,), width=max(S(2), 1))
        d.ellipse([sx - S(4), sy - S(4), sx + S(4), sy + S(4)], outline=rgb + (255,), width=max(S(2), 1))
        d.rectangle([tx - S(7), ty - S(4), tx + tw + S(9), ty + S(21)], fill=(10, 12, 20, 215))
        d.text((tx, ty), text, font=F_LBL, fill=rgb)

    # ---- footnote
    d.text((PAD, H - S(136)), "surface colour = pressure anomaly, amplitude-stretched (signed |x|^0.35) so the weak early wave is visible; "
                              "printed values are true model output", font=F_TINY, fill=(118, 130, 160))
    return img


# ----------------------------------------------------------------- per-frame
LIGHT = norml(np.array([0.55, 0.45, 0.70]))
os.makedirs(args.frames_dir, exist_ok=True)


def build_frame(k):
    x = k / max(NFRAME - 1, 1)
    day = ease_seq(x, DAY_KEYS)
    az = np.deg2rad(azimuth_at(day, x))
    el = np.deg2rad(ease_seq(x, EL_KEYS))
    dist = ease_seq(x, DIST_KEYS)
    eye = np.array([dist * np.cos(el) * np.cos(az), dist * np.cos(el) * np.sin(az), dist * np.sin(el)])
    cam = Camera(eye, fov_deg=34.0)

    # interpolated fields
    psa_d = lerp_day(PSA_DISP, day)
    om_d = lerp_day(OM, day)
    cld_d = lerp_day(CLD, day)
    pr_d = lerp_day(PRECT, day)
    u_d = lerp_day(UJ, day)
    v_d = lerp_day(VJ, day)

    tex = lut_lookup(WAVE_LUT, psa_d.ravel(), -DISP_MAX, DISP_MAX).reshape(NLAT, NLON, 3)
    base, _ = render_globe(cam, tex, LIGHT)

    acc_rgb = np.zeros((H, W, 3), np.float32)
    acc_a = np.zeros((H, W), np.float32)

    # ---- jet trails
    spd = advance_particles(u_d, v_d, dt_s=1500.0 if not args.fast else 3000.0)
    tl_lat = P_HIST_LAT.ravel()
    tl_lon = P_HIST_LON.ravel()
    fade = np.tile((1.0 - np.arange(TRAIL) / TRAIL) ** 1.5, (NPART, 1)).ravel()
    life = np.repeat(np.clip(np.sin(np.pi * np.clip(P_AGE, 0, 1)) * 1.5, 0, 1), TRAIL)
    pts = sph2cart(tl_lat, tl_lon, R_JET)
    sx, sy, z = cam.project(pts)
    vis = (z > 0) & ~cam.hidden(pts) & (sx > -30) & (sx < W + 30) & (sy > -30) & (sy < H + 30)
    if vis.any():
        sp_rep = np.repeat(spd, TRAIL)[vis]
        cols = lut_lookup(JET_LUT, sp_rep, 6.0, 34.0)
        rr = np.clip(cam.px_per_unit(z[vis]) * 0.011, 1.0, 6.0)
        splat(acc_rgb, acc_a, sx[vis], sy[vis], rr, cols, 0.42 * fade[vis] * life[vis])

    # ---- rising air: strongest ascending columns, vertical streaks
    band = (lev > 150) & (lev < 900)
    omb = np.where(band[:, None, None], om_d, 0.0)
    nsel = 620
    flat = omb.ravel()
    idx = np.argpartition(flat, nsel)[:nsel]
    kk, ii, jj = np.unravel_index(idx, omb.shape)
    strength = -flat[idx]
    keep = strength > 1e-4
    kk, ii, jj, strength = kk[keep], ii[keep], jj[keep], strength[keep]
    asc_peak_now = float(strength.max()) if strength.size else 0.0
    if strength.size:
        top_r = altitude_of(lev[kk])
        cols = lut_lookup(ASC_LUT, strength, 0.0, max(ASC_PEAK.max(), 1e-6))
        # jitter within the grid cell: markers on exact cell centres read as a
        # regular lattice rather than a physical field
        jz = np.random.default_rng(1000 + int(day * 37)).uniform(-0.5, 0.5, (2, ii.size))
        la_j = lat[ii] + jz[0] * 1.9
        lo_j = lon[jj] + jz[1] * 2.5
        # 4-point vertical streak from just above the surface up to the parcel
        for frac, aw in ((0.15, 0.30), (0.45, 0.45), (0.75, 0.70), (1.0, 1.0)):
            rr_ = 1.02 + (top_r - 1.02) * frac
            p = sph2cart(la_j, lo_j, rr_)
            sx, sy, z = cam.project(p)
            vv = (z > 0) & ~cam.hidden(p)
            if vv.any():
                sz = np.clip(cam.px_per_unit(z[vv]) * 0.016 * (0.5 + 0.5 * frac), 1.0, 9.0)
                a = 0.30 * aw * np.clip(strength[vv] / max(asc_peak_now, 1e-9), 0.15, 1.0)
                splat(acc_rgb, acc_a, sx[vv], sy[vv], sz, cols[vv], a)

    # ---- condensation / cloud water
    cmax = cld_d.max(axis=0)
    karg = cld_d.argmax(axis=0)
    sig = cmax > 40.0
    bg = (cmax > 0.5) & ~sig
    # random (not grid-stride) thinning of the uniform background layer, so no
    # visible lattice appears; every significant cell is always kept
    dec = BG_KEEP
    ii, jj = np.where(sig | (bg & dec))
    cl_peak_now = float(cmax.max())
    if ii.size:
        vals = cmax[ii, jj]
        pr_ = altitude_of(lev[karg[ii, jj]])
        jz = np.random.default_rng(2000 + int(day * 37)).uniform(-0.5, 0.5, (2, ii.size))
        p = sph2cart(lat[ii] + jz[0] * 1.9, lon[jj] + jz[1] * 2.5, pr_)
        sx, sy, z = cam.project(p)
        vv = (z > 0) & ~cam.hidden(p)
        if vv.any():
            cols = lut_lookup(CLOUD_LUT, vals[vv], 0.0, CL_THRESHOLD)
            big = np.clip(vals[vv] / CL_THRESHOLD, 0, 1)
            sz = np.clip(cam.px_per_unit(z[vv]) * (0.010 + 0.030 * big ** 0.5), 1.0, 22.0)
            a = 0.20 + 0.62 * big ** 0.45
            splat(acc_rgb, acc_a, sx[vv], sy[vv], sz, cols, a)

    # ---- precipitation
    ii, jj = np.where(pr_d > 0.05)
    pr_peak_now = float(pr_d.max())
    nwet_now = int(ii.size)
    if ii.size:
        vals = pr_d[ii, jj]
        jz = np.random.default_rng(3000 + int(day * 37)).uniform(-0.5, 0.5, (2, ii.size))
        la_r = lat[ii] + jz[0] * 1.9
        lo_r = lon[jj] + jz[1] * 2.5
        for rr_, aw, szf in ((1.006, 1.0, 0.020), (1.030, 0.55, 0.013), (1.055, 0.30, 0.009)):
            p = sph2cart(la_r, lo_r, rr_)
            sx, sy, z = cam.project(p)
            vv = (z > 0) & ~cam.hidden(p)
            if vv.any():
                sz = np.clip(cam.px_per_unit(z[vv]) * szf * (0.4 + 0.6 * np.clip(vals[vv] / max(PR_PEAK.max(), 1e-9), 0, 1) ** 0.4), 1.0, 16.0)
                cols = np.tile(np.array([[92.0, 155.0, 255.0]], np.float32), (vv.sum(), 1))
                splat(acc_rgb, acc_a, sx[vv], sy[vv], sz, cols, 0.55 * aw)

    frame = composite(base, acc_rgb, acc_a)

    # ---- data-driven callouts anchored on real features
    callouts = []

    def anchor(field2d, text, hexc, thresh):
        if float(np.nanmax(field2d)) < thresh:
            return
        i0, j0 = np.unravel_index(int(np.nanargmax(field2d)), field2d.shape)
        p = sph2cart(np.array([lat[i0]]), np.array([lon[j0]]), 1.05)
        sxx, syy, zz = cam.project(p)
        if zz[0] > 0 and not cam.hidden(p)[0] and 90 < sxx[0] < W - 90 and 130 < syy[0] < H - 150:
            callouts.append((float(sxx[0]), float(syy[0]), text, hexc))

    if x > 0.40:
        anchor(-om_d.min(axis=0), "rising air", "#ff8c2b", 0.01)
    if x > 0.58 and cl_peak_now > 40:
        anchor(cmax, f"condensation · {cl_peak_now:,.0f} mg/kg", "#cfe0f2", 40)
    if cl_peak_now > 0.80 * CL_THRESHOLD:
        anchor(cmax, f"{100*cl_peak_now/CL_THRESHOLD:.0f}% of cloud-water threshold", "#b07cd0", 0.8 * CL_THRESHOLD)
    if pr_peak_now > 0.05:
        anchor(pr_d, f"precipitation · {pr_peak_now:.0f} mm/day", "#5b9bff", 0.05)

    img = Image.fromarray(np.clip(frame, 0, 255).astype(np.uint8), "RGB")
    feats = dict(cl_peak=cl_peak_now, asc_peak=asc_peak_now, pr_peak=pr_peak_now,
                 nwet=nwet_now, amp=float(np.interp(day, np.arange(NDAY), TRUE_AMP)),
                 callouts=callouts)
    draw_overlay(img, x, day, cam, feats)
    return img


# ----------------------------------------------------------------- drive
if args.preview:
    picks = np.linspace(0, NFRAME - 1, args.preview).astype(int)
    for n, k in enumerate(picks):
        # advance particle state up to k so trails look developed
        build_frame(k).save(f"preview_frame_{n:02d}.png")
    print("previews:", [f"preview_frame_{n:02d}.png" for n in range(len(picks))])
    sys.exit(0)

def reset_particles():
    """restore the deterministic initial particle state"""
    global P_LAT, P_LON, P_AGE, P_HIST_LAT, P_HIST_LON, RNG
    RNG = np.random.default_rng(7)
    P_LAT, P_LON = seed_particles(NPART)
    P_AGE = RNG.uniform(0, 1, NPART).astype(np.float32)
    P_HIST_LAT = np.tile(P_LAT[:, None], (1, TRAIL))
    P_HIST_LON = np.tile(P_LON[:, None], (1, TRAIL))


def advect_only(k):
    """replay just the (cheap) particle advection for frame k, no rendering"""
    x = k / max(NFRAME - 1, 1)
    day = ease_seq(x, DAY_KEYS)
    advance_particles(lerp_day(UJ, day), lerp_day(VJ, day),
                      dt_s=1500.0 if not args.fast else 3000.0)


def render_chunk(bounds):
    """render frames [k0,k1); replays advection from 0 so trails stay continuous"""
    k0, k1 = bounds
    reset_particles()
    for k in range(k0):
        advect_only(k)
    n = 0
    for k in range(k0, k1):
        path = os.path.join(args.frames_dir, f"f{k:05d}.png")
        if os.path.exists(path):      # resumable: skip frames already rendered
            continue
        build_frame(k).save(path)
        n += 1
    return n


if __name__ == "__main__":
    import time
    from multiprocessing import Pool

    nproc = max(1, min(os.cpu_count() or 1, 4))
    edges = np.linspace(0, NFRAME, nproc + 1).astype(int)
    chunks = [(int(edges[i]), int(edges[i + 1])) for i in range(nproc) if edges[i + 1] > edges[i]]
    print(f"rendering {NFRAME} frames at {W}x{H} on {len(chunks)} worker(s)", flush=True)
    t0 = time.time()
    if len(chunks) == 1:
        render_chunk(chunks[0])
    else:
        with Pool(len(chunks)) as pool:
            for done in pool.imap_unordered(render_chunk, chunks):
                print(f"  chunk of {done} frames done, {time.time()-t0:.0f}s elapsed", flush=True)
    made = len(glob.glob(os.path.join(args.frames_dir, "f*.png")))
    print(f"rendered {made}/{NFRAME} frames in {(time.time()-t0)/60:.1f} min")