app/outputs/spatial_web.py

"""Web-friendly spatial data export.

Turns the in-memory `SpatialData` produced by EO products into a
downsampled grid JSON structure that MapLibre can render directly, so
users see the actual indicator overlay on the results dashboard instead
of just a tinted AOI rectangle.
"""
from __future__ import annotations

import logging
import math
from typing import Any

import numpy as np

from app.eo_products.base import SpatialData

logger = logging.getLogger(__name__)

# Target grid size for the web overlay. Small enough that MapLibre can
# render as GeoJSON polygons without lagging, large enough to still
# look like a heatmap rather than a mosaic.
TARGET_GRID_SIZE = 96


def raster_to_grid_dict(
    raster_path: str,
    *,
    band: int = 1,
    spatial: SpatialData,
    status_value: str,
) -> dict[str, Any] | None:
    """Read a raster file and return a grid-format spatial JSON dict.

    The returned dict matches the existing "grid" spatial JSON contract
    consumed by the frontend (`map_type`, `data`, `lats`, `lons`, ...)
    and additionally embeds explicit color `stops` derived from the
    product's matplotlib colormap name so the browser doesn't need to
    know anything about colormap conventions.

    Returns None on any read/processing failure — the caller should
    fall back to the status-only overlay in that case.
    """
    try:
        import rasterio
    except ImportError:
        logger.warning("raster_to_grid_dict: rasterio not available")
        return None

    try:
        with rasterio.open(raster_path) as src:
            band = min(band or 1, src.count)
            data = src.read(band).astype(np.float32)
            nodata = src.nodata
            bounds = src.bounds
    except Exception as exc:
        logger.warning("raster_to_grid_dict: cannot open %s (%s)", raster_path, exc)
        return None

    if data.size == 0:
        logger.warning("raster_to_grid_dict: empty data array for %s", raster_path)
        return None

    if nodata is not None:
        data = np.where(data == nodata, np.nan, data)

    # Downsample to a web-friendly size using block-mean, ignoring NaNs.
    h, w = data.shape
    scale = max(math.ceil(max(h, w) / TARGET_GRID_SIZE), 1)
    if scale > 1:
        data = _block_reduce_nanmean(data, scale)

    # Rasterio reads row 0 as the northernmost row; the frontend grid
    # renderer expects row 0 at the south edge. Flip once here so the
    # polygon construction in map.js lines up with the data.
    data = np.flipud(data)

    ny, nx = data.shape
    lons = np.linspace(bounds.left, bounds.right, nx + 1)
    lats = np.linspace(bounds.bottom, bounds.top, ny + 1)

    # Determine value range, preferring product-supplied vmin/vmax.
    vmin: float
    vmax: float
    if spatial.vmin is not None:
        vmin = float(spatial.vmin)
    else:
        valid = data[np.isfinite(data)]
        vmin = float(np.nanmin(valid)) if valid.size else float("nan")
    if spatial.vmax is not None:
        vmax = float(spatial.vmax)
    else:
        valid = data[np.isfinite(data)]
        vmax = float(np.nanmax(valid)) if valid.size else float("nan")

    if not math.isfinite(vmin) or not math.isfinite(vmax):
        # All values were NaN — still render an empty grid rather than
        # dropping the indicator entirely, so the frontend can at least
        # show the legend and the user knows the indicator was processed.
        logger.warning(
            "raster_to_grid_dict: all values NaN for %s; rendering empty grid",
            raster_path,
        )
        vmin, vmax = 0.0, 1.0
    elif vmin == vmax:
        vmin -= 0.5
        vmax += 0.5

    stops = _colormap_stops(spatial.colormap or "RdYlGn", vmin, vmax, n=5)

    # Replace NaN with None so JSON.dump yields nulls the frontend can skip.
    data_list: list[list[float | None]] = [
        [None if not math.isfinite(v) else round(float(v), 4) for v in row]
        for row in data.tolist()
    ]

    return {
        "map_type": "grid",
        "status": status_value,
        "data": data_list,
        "lats": [round(float(x), 6) for x in lats.tolist()],
        "lons": [round(float(x), 6) for x in lons.tolist()],
        "label": spatial.label,
        "colormap": spatial.colormap,
        "vmin": round(float(vmin), 4),
        "vmax": round(float(vmax), 4),
        "stops": stops,
    }


def _block_reduce_nanmean(data: np.ndarray, factor: int) -> np.ndarray:
    """Reduce `data` by `factor` using NaN-aware block mean."""
    h, w = data.shape
    new_h = h // factor
    new_w = w // factor
    if new_h == 0 or new_w == 0:
        return data
    trimmed = data[: new_h * factor, : new_w * factor]
    blocks = trimmed.reshape(new_h, factor, new_w, factor)
    import warnings
    with np.errstate(invalid="ignore"), warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return np.nanmean(blocks, axis=(1, 3))


def _colormap_stops(
    cmap_name: str, vmin: float, vmax: float, *, n: int = 5
) -> list[dict[str, Any]]:
    """Return n evenly-spaced (value, hex color) stops from a matplotlib cmap.

    Falls back to a neutral grey ramp if matplotlib can't resolve the
    colormap name.
    """
    try:
        import matplotlib.cm as mcm
        import matplotlib.colors as mcolors

        cmap = mcm.get_cmap(cmap_name)
    except Exception:
        cmap = None

    stops: list[dict[str, Any]] = []
    for i in range(n):
        t = i / (n - 1)
        value = vmin + t * (vmax - vmin)
        if cmap is None:
            # Fallback grey ramp: dark → light
            shade = int(230 - t * 160)
            color = f"#{shade:02x}{shade:02x}{shade:02x}"
        else:
            rgba = cmap(t)
            color = mcolors.to_hex(rgba)
        stops.append({"value": round(float(value), 4), "color": color})
    return stops