rustwx_stormscope/webmap_preview.py

#!/usr/bin/env python
"""Make a lightweight Leaflet web-map preview from StormScope .npz outputs."""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

import numpy as np
from PIL import Image

from .paths import default_outdir, default_webmap_dir


@dataclass(frozen=True)
class Product:
    slug: str
    field: str
    title: str
    default_opacity: float


PRODUCTS = {
    "refc": Product("refc", "refc", "Composite Reflectivity", 0.82),
    "ir": Product("ir", "ch_ir_1035", "IR Brightness Temperature", 0.64),
}

VIEWS = {
    "data": None,
    "conus": {"west": -126.0, "east": -66.0, "south": 23.0, "north": 50.0},
    "central_plains": {"west": -106.5, "east": -90.0, "south": 31.5, "north": 47.2},
}


REFC_LEVELS = np.array([5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75], dtype=np.float32)
REFC_COLORS = np.array(
    [
        (236, 243, 252, 0),
        (205, 222, 246, 125),
        (160, 194, 232, 155),
        (103, 157, 212, 185),
        (32, 102, 172, 205),
        (44, 162, 95, 215),
        (166, 217, 106, 220),
        (255, 240, 102, 225),
        (253, 174, 55, 232),
        (244, 109, 32, 238),
        (215, 48, 31, 242),
        (165, 0, 38, 246),
        (122, 50, 148, 248),
        (84, 39, 136, 250),
        (68, 68, 68, 252),
    ],
    dtype=np.uint8,
)

IR_STOPS = np.array([190, 205, 220, 235, 250, 265, 280, 300], dtype=np.float32)
IR_COLORS = np.array(
    [
        (62, 20, 115, 220),
        (44, 76, 160, 205),
        (32, 143, 195, 190),
        (99, 194, 174, 168),
        (226, 237, 154, 135),
        (250, 180, 92, 105),
        (126, 126, 126, 70),
        (255, 255, 255, 0),
    ],
    dtype=np.float32,
)


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Build a web-map overlay preview from StormScope .npz")
    ap.add_argument("--meta", default="", help="cycle metadata JSON; defaults to latest completed cycle")
    ap.add_argument("--outdir", default=str(default_outdir()))
    ap.add_argument("--output-dir", default=str(default_webmap_dir()))
    ap.add_argument("--products", default="refc,ir", help="comma-separated: refc,ir")
    ap.add_argument("--view", default="central_plains", choices=sorted(VIEWS), help="initial web-map view")
    ap.add_argument("--max-frames", type=int, default=12)
    ap.add_argument("--force", action="store_true")
    return ap.parse_args()


def read_json(path: Path) -> dict:
    return json.loads(path.read_text())


def parse_dt(value: str) -> datetime:
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


def latest_meta(outdir: Path) -> Path:
    candidates: list[tuple[datetime, Path]] = []
    for path in outdir.glob("stormscope_*_*_*_meta.json"):
        try:
            meta = read_json(path)
            if meta.get("complete", True):
                candidates.append((parse_dt(meta["init_time"]), path))
        except Exception:
            continue
    if not candidates:
        raise FileNotFoundError(f"no completed StormScope cycle metadata found in {outdir}")
    return sorted(candidates)[-1][1]


def normalize_lon(values: np.ndarray) -> np.ndarray:
    return np.where(values > 180.0, values - 360.0, values)


def bounds_from_npz(npz_path: Path) -> dict:
    with np.load(npz_path) as data:
        lats = data["lats"].astype(np.float32)
        lons = normalize_lon(data["lons"].astype(np.float32))
    finite = np.isfinite(lats) & np.isfinite(lons)
    return {
        "west": float(np.nanmin(lons[finite])),
        "east": float(np.nanmax(lons[finite])),
        "south": float(np.nanmin(lats[finite])),
        "north": float(np.nanmax(lats[finite])),
    }


def rgba_refc(values: np.ndarray) -> np.ndarray:
    rgba = np.zeros((*values.shape, 4), dtype=np.uint8)
    mask = np.isfinite(values) & (values >= REFC_LEVELS[0])
    if not np.any(mask):
        return rgba
    indices = np.searchsorted(REFC_LEVELS, values[mask], side="right") - 1
    indices = np.clip(indices, 0, len(REFC_COLORS) - 1)
    rgba[mask] = REFC_COLORS[indices]
    return rgba


def rgba_ir(values: np.ndarray) -> np.ndarray:
    rgba = np.zeros((*values.shape, 4), dtype=np.uint8)
    mask = np.isfinite(values)
    if not np.any(mask):
        return rgba
    clipped = np.clip(values[mask], IR_STOPS[0], IR_STOPS[-1])
    channels = []
    for channel in range(4):
        channels.append(np.interp(clipped, IR_STOPS, IR_COLORS[:, channel]))
    rgba[mask] = np.stack(channels, axis=1).clip(0, 255).astype(np.uint8)
    return rgba


def field_rgba(product: Product, values: np.ndarray) -> np.ndarray:
    if product.slug == "refc":
        return rgba_refc(values)
    if product.slug == "ir":
        return rgba_ir(values)
    raise ValueError(f"unknown product {product.slug}")


def frame_label(meta: dict, index: int) -> str:
    step = int(meta.get("step_minutes") or 0)
    if step:
        lead = step * (index + 1)
        if lead % 60 == 0:
            return f"+{lead // 60}h"
        return f"+{lead}m"
    return f"f{index + 1:02d}"


def render_product_frame(npz_path: Path, product: Product, output: Path, force: bool) -> dict:
    if output.exists() and output.stat().st_size > 0 and not force and output.stat().st_mtime >= npz_path.stat().st_mtime:
        return {"path": output}
    output.parent.mkdir(parents=True, exist_ok=True)
    with np.load(npz_path) as data:
        values = data[product.field].astype(np.float32)
    rgba = field_rgba(product, values)
    Image.fromarray(rgba, "RGBA").save(output, optimize=False)
    return {"path": output}


def write_html(output_dir: Path) -> None:
    (output_dir / "index.html").write_text(
        """<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>StormScope Web Map Preview</title>
  <link rel="stylesheet" href="https://unpkg.com/leaflet@1.9.4/dist/leaflet.css">
  <style>
    * { box-sizing: border-box; }
    html, body, #map { height: 100%; margin: 0; }
    body { font: 14px/1.35 system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; color: #17201b; }
    #map { background: #eef2ef; }
    .panel {
      position: absolute;
      z-index: 700;
      left: 12px;
      right: 12px;
      top: 12px;
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 12px;
      padding: 10px 12px;
      background: rgba(255, 255, 255, .92);
      border: 1px solid #d5ddd7;
      border-radius: 8px;
      box-shadow: 0 8px 28px rgba(20, 30, 25, .12);
    }
    h1 { margin: 0; font-size: 16px; letter-spacing: 0; }
    .meta { color: #59655f; font-size: 12px; margin-top: 2px; }
    .controls { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; justify-content: end; }
    button, select {
      min-height: 32px;
      border: 1px solid #cbd5ce;
      border-radius: 6px;
      background: white;
      color: #17201b;
      font: inherit;
      font-weight: 650;
      padding: 0 10px;
    }
    input[type="range"] { width: min(42vw, 420px); }
    #stamp { min-width: 92px; font-weight: 750; text-align: center; }
    .leaflet-control-attribution { font-size: 10px; }
    @media (max-width: 720px) {
      .panel { display: grid; top: 8px; left: 8px; right: 8px; }
      .controls { justify-content: start; }
      input[type="range"] { width: 100%; }
    }
  </style>
</head>
<body>
  <div id="map"></div>
  <div class="panel">
    <div>
      <h1 id="title">StormScope</h1>
      <div id="meta" class="meta"></div>
    </div>
    <div class="controls">
      <select id="product"></select>
      <button id="prev" type="button">Prev</button>
      <input id="frame" type="range" min="0" max="0" value="0">
      <button id="next" type="button">Next</button>
      <button id="play" type="button">Play</button>
      <span id="stamp"></span>
    </div>
  </div>
  <script src="https://unpkg.com/leaflet@1.9.4/dist/leaflet.js"></script>
  <script src="manifest.js"></script>
  <script>
    const m = window.STORMSCOPE_WEBMAP;
    const bounds = [[m.bounds.south, m.bounds.west], [m.bounds.north, m.bounds.east]];
    const viewBounds = m.view_bounds ? [[m.view_bounds.south, m.view_bounds.west], [m.view_bounds.north, m.view_bounds.east]] : bounds;
    const map = L.map("map", { preferCanvas: true, zoomControl: false });
    L.tileLayer("https://tile.openstreetmap.org/{z}/{x}/{y}.png", {
      maxZoom: 12,
      attribution: "&copy; OpenStreetMap contributors"
    }).addTo(map);
    L.control.zoom({ position: "bottomleft" }).addTo(map);
    map.fitBounds(viewBounds, { padding: [18, 18] });

    const productSelect = document.getElementById("product");
    const slider = document.getElementById("frame");
    const stamp = document.getElementById("stamp");
    const playButton = document.getElementById("play");
    document.getElementById("title").textContent = `${m.mode} ${m.cycle}`;
    document.getElementById("meta").textContent = `${m.init_time} | ${m.resolution} | ${m.frame_count} frames`;
    Object.entries(m.products).forEach(([slug, product]) => {
      const option = document.createElement("option");
      option.value = slug;
      option.textContent = product.title;
      productSelect.appendChild(option);
    });

    let product = productSelect.value;
    let frame = 0;
    let timer = null;
    let overlay = null;

    function currentProduct() { return m.products[product]; }
    function currentFrame() { return currentProduct().frames[frame]; }

    function updateOverlay() {
      const p = currentProduct();
      slider.max = String(p.frames.length - 1);
      if (frame >= p.frames.length) frame = p.frames.length - 1;
      const f = currentFrame();
      slider.value = String(frame);
      stamp.textContent = f.label;
      if (!overlay) {
        overlay = L.imageOverlay(f.url, bounds, { opacity: p.opacity, interactive: false }).addTo(map);
      } else {
        overlay.setUrl(f.url);
        overlay.setOpacity(p.opacity);
      }
    }

    productSelect.addEventListener("change", () => {
      product = productSelect.value;
      frame = Math.min(frame, currentProduct().frames.length - 1);
      updateOverlay();
    });
    slider.addEventListener("input", () => { frame = Number(slider.value); updateOverlay(); });
    document.getElementById("prev").addEventListener("click", () => {
      frame = (frame + currentProduct().frames.length - 1) % currentProduct().frames.length;
      updateOverlay();
    });
    document.getElementById("next").addEventListener("click", () => {
      frame = (frame + 1) % currentProduct().frames.length;
      updateOverlay();
    });
    playButton.addEventListener("click", () => {
      if (timer) {
        clearInterval(timer);
        timer = null;
        playButton.textContent = "Play";
      } else {
        timer = setInterval(() => {
          frame = (frame + 1) % currentProduct().frames.length;
          updateOverlay();
        }, 650);
        playButton.textContent = "Pause";
      }
    });
    updateOverlay();
  </script>
</body>
</html>
""",
        encoding="utf-8",
    )


def build_preview(meta_path: Path, outdir: Path, output_dir: Path, product_slugs: list[str], view: str, max_frames: int, force: bool) -> dict:
    meta = read_json(meta_path)
    files = list(meta.get("files", []))
    if max_frames > 0:
        files = files[:max_frames]
    if not files:
        raise ValueError(f"metadata has no files: {meta_path}")

    output_dir.mkdir(parents=True, exist_ok=True)
    assets = output_dir / "assets" / meta.get("cycle", "cycle")
    bounds = bounds_from_npz(outdir / files[0])
    manifest = {
        "schema": "rustwx-stormscope.webmap_preview.v1",
        "source_meta": str(meta_path),
        "cycle": meta.get("cycle"),
        "mode": "nearcast" if meta.get("mode") == "forecast" else meta.get("mode"),
        "resolution": meta.get("resolution"),
        "init_time": meta.get("init_time"),
        "bounds": bounds,
        "view": view,
        "view_bounds": VIEWS[view] or bounds,
        "frame_count": len(files),
        "products": {},
    }

    for slug in product_slugs:
        if slug not in PRODUCTS:
            raise ValueError(f"unknown product {slug!r}; available: {', '.join(PRODUCTS)}")
        product = PRODUCTS[slug]
        frames = []
        for index, filename in enumerate(files):
            npz_path = outdir / filename
            png_name = f"{Path(filename).stem}_{product.slug}.png"
            png_path = assets / product.slug / png_name
            render_product_frame(npz_path, product, png_path, force)
            valid_times = meta.get("valid_times") or []
            frames.append(
                {
                    "index": index,
                    "label": frame_label(meta, index),
                    "valid_time": valid_times[index] if index < len(valid_times) else None,
                    "url": png_path.relative_to(output_dir).as_posix(),
                    "source_npz": str(npz_path),
                }
            )
        manifest["products"][slug] = {
            "field": product.field,
            "title": product.title,
            "opacity": product.default_opacity,
            "frames": frames,
        }

    (output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    (output_dir / "manifest.js").write_text(
        "window.STORMSCOPE_WEBMAP = " + json.dumps(manifest, indent=2) + ";\n",
        encoding="utf-8",
    )
    write_html(output_dir)
    return manifest


def main() -> int:
    args = parse_args()
    outdir = Path(args.outdir)
    meta = Path(args.meta) if args.meta else latest_meta(outdir)
    products = [slug.strip() for slug in args.products.split(",") if slug.strip()]
    manifest = build_preview(meta, outdir, Path(args.output_dir), products, args.view, args.max_frames, args.force)
    print(json.dumps({"ok": True, "index": str(Path(args.output_dir) / "index.html"), "cycle": manifest["cycle"]}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())