#!/usr/bin/env python

# SPDX-License-Identifier: GPL-3.0-or-later
# FileType: SOURCE
# SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences
# SPDX-FileCopyrightText: 2025, "Eva Boergens" at GFZ Potsdam

"""
Gaussian filtering on a sphere/ellipsoid.

Faithfully mirrors the kernel lookup/scaling in main_GRIDgaussfilter.f90
and uses an ellipsoidal reference surface with the same
anisotropy scaling. Builds a single sparse weight matrix (CSR) that can be
reused across many epochs — avoiding per-epoch recomputation.

Units & conventions (SI):
- All distances are **meters [m]**.
- Angular inputs are degrees; internal calculations convert to radians.

Available distance methods:
- "fortran": Legacy method with anisotropy scaling (ellipsoidal approximation)
- "vincenty": Exact ellipsoidal geodesic distance (most accurate)
- "spherical": Great-circle distance on a sphere (fastest)

Public API (SI):
- ellipsoidal_distance_method(lon_deg, lat_deg, gaussr_m, distance="fortran", mask=None) -> csr_matrix
- spherical_distance_method(lon_deg, lat_deg, gaussr_m, radius_m=None, mask=None) -> csr_matrix
- ellipsoidal_distance_method_cached(..., gaussr_m, ...) -> (csr_matrix, from_cache)
- spherical_distance_method_cached(...) -> (csr_matrix, from_cache)
- apply_filter(W, data_2d, mask=None, fill_value=None) -> np.ndarray
- apply_filter_time(W, data_3d, mask=None, fill_value=None) -> np.ndarray
- ellipsoid_radius(lat_rad) -> meters
- flattening() -> dimensionless
"""

import hashlib
import json
import math
from pathlib import Path
from typing import List, Tuple, Optional, Literal
from functools import lru_cache
from logging import Logger
import numpy as np
from scipy.sparse import csr_matrix
from geogravl3.datamodels.grids import Grid2DObject, Grid3DObject
from geogravl3.utils.global_constants import constants as cnsts
from geogravl3.utils.grid_utils import (
    anisotropy_scaled_distance_m as _anisotropy_scaled_distance_m,
    ellipsoid_radius as _ellipsoid_radius_m,
    periodic_lon_index as _periodic_lon_index,
    vincenty_distance as _vincenty_distance,
    spherical_distance as _spherical_distance,
)


# Helper
def _resolve_logger(logger: Optional[Logger]) -> Logger:
    import logging
    return logger if isinstance(logger, Logger) else logging.getLogger(__name__)


# constants
deg_2_rad = cnsts.deg_2_rad
rad_2_deg = cnsts.rad_2_deg
earths_semi_major_axis_m = cnsts.earths_radius_iers
earths_semi_minor_axis_m = earths_semi_major_axis_m * (1.0 - cnsts.earths_flattening_iers)
ntab = 100

# Ellipsoid in meters
radius_equator = earths_semi_major_axis_m
radius_polar = earths_semi_minor_axis_m

# Mean Earth radius for spherical calculations (good approximation)
# Using (2a + b)/3 gives a good mean radius for Earth's ellipsoid
earth_mean_radius_m = (2 * radius_equator + radius_polar) / 3

# IERS flattening, exposed as a function for test/API compatibility
_flattening = cnsts.earths_flattening_iers


# kernel: exact lookup-table logic from Fortran, in meters
def _kernel_table(gaussr_m: float) -> Tuple[np.ndarray, np.ndarray, float]:
    """Build the (xt, fx) lookup table in meters.

    Table definition:
      - rx = i / (2*ntab) for i=0..ntab, hence rx ∈ [0, 0.5]
      - fx = exp(-0.5 * (6*rx)^2)
      - gauss_halfwidth = sqrt((-2)*ln(0.5)) / 6
      - maxvaliddist    = gaussr * 0.5 / gauss_halfwidth
      - xt_m = rx * (2 * maxvaliddist)  (distances in m)
      - dxtab_m = maxvaliddist / ntab

    Returns:
        xt_m: distances (m) of length ntab+1.
        fx: weights corresponding to xt_m.
        dxtab_m: table spacing in m.
    """
    xt = np.zeros(ntab + 1, dtype=np.float64)
    fx = np.zeros(ntab + 1, dtype=np.float64)

    for i in range(ntab + 1):
        rx = i / (2.0 * ntab)  # range [0, 0.5]
        xh = -0.5 * (6.0 * rx) ** 2
        fx[i] = math.exp(xh)
        xt[i] = rx

    gauss_halfwidth = math.sqrt((-2.0) * math.log(0.5)) / 6.0
    maxvaliddist = gaussr_m * 0.5 / gauss_halfwidth  # meters
    xt_m = xt * (maxvaliddist * 2.0)
    dxtab_m = maxvaliddist / ntab

    return xt_m, fx, dxtab_m


def _kernel_weight(dist_m: float, xt_m: np.ndarray, fx: np.ndarray, dxtab_m: float) -> float:
    """Kernel weight via linear interpolation on a precomputed table (meters)."""
    itab = int(math.floor(dist_m / dxtab_m))
    if itab < ntab:
        base = fx[itab]
        slope = (fx[itab + 1] - fx[itab]) / dxtab_m
        return base + (dist_m - xt_m[itab]) * slope
    return 0.0


@lru_cache(maxsize=64)
def _kernel_table_cached(gaussr_m: float) -> Tuple[np.ndarray, np.ndarray, float]:
    """Cache the kernel lookup table for a given Gaussian radius (meters)."""
    return _kernel_table(gaussr_m)


def kernel_at_distance(dist_m: float, gaussr_m: float) -> float:
    """Return kernel weight at distance (meters) using a cached table."""
    xt_m, fx, dxtab_m = _kernel_table_cached(gaussr_m)
    return _kernel_weight(dist_m, xt_m, fx, dxtab_m)


# ---- ellipsoid helpers (SI) ----
def ellipsoid_radius(lat_rad: float) -> float:
    """Return ellipsoid center-to-surface distance at latitude (meters)."""
    return float(_ellipsoid_radius_m(lat_rad))


def flattening() -> float:
    """Return geodetic flattening f = (a - b)/a used for anisotropy scaling."""
    return float(_flattening)


def _lat_bounds_indices(lat: np.ndarray, center_deg: float, halfspan_deg: float) -> Tuple[int, int]:
    """
    Return the start and end indices in a sorted latitude array within [center +/- halfspan].

    Handles both ascending and descending latitude arrays. If no latitudes fall
    strictly within the window, returns the index of the single closest latitude.

    Args:
        lat: 1D monotonic array of latitudes in degrees.
        center_deg: Center latitude of the search window (degrees).
        halfspan_deg: Half-width of the search window (degrees).

    Returns:
        Tuple[int, int]: (start_index, end_index) inclusive, clamped to array bounds.
    """
    asc = lat[1] > lat[0]
    if asc:
        j_start = int(np.searchsorted(lat, center_deg - halfspan_deg, side="left"))
        j_end = int(np.searchsorted(lat, center_deg + halfspan_deg, side="right") - 1)
    else:
        lat_rev = lat[::-1]
        j_start_r = int(np.searchsorted(lat_rev, center_deg - halfspan_deg, side="left"))
        j_end_r = int(np.searchsorted(lat_rev, center_deg + halfspan_deg, side="right") - 1)
        n = lat.size
        j_start = n - 1 - j_end_r
        j_end = n - 1 - j_start_r

    j_start = max(0, min(j_start, lat.size - 1))
    j_end = max(0, min(j_end, lat.size - 1))
    if j_end < j_start:
        j_start = j_end = int(np.argmin(np.abs(lat - center_deg)))
    return j_start, j_end


def _lon_neighborhood_indices(
    lon: np.ndarray,
    center_deg: float,
    halfspan_deg: float,
    is_global: bool,
    logger: Optional[Logger] = None,
) -> List[int]:
    """Return periodic-aware longitude indices within +/- halfspan_deg of center_deg.

    Args:
        lon: 1D longitudes in degrees, ascending, regular.
        center_deg: central longitude in degrees.
        halfspan_deg: half-span window in degrees.
        is_global: True if the longitudes wrap globally (enable periodic indexing).
        logger: optional Logger instance.

    Returns:
        Sorted unique list of column indices within the window.
    """
    _log = _resolve_logger(logger)

    nlon = lon.size
    if nlon < 2:
        _log.error("_lon_neighborhood_indices: need at least two longitudes; got %d.", nlon)
        raise ValueError("lon must contain at least two values.")

    dlon = (lon[-1] - lon[0]) / (nlon - 1)
    # index of the closest longitude to center_deg
    i_center = int(round((center_deg - lon[0]) / max(dlon, 1e-12)))

    steps = int(math.ceil(halfspan_deg / max(abs(dlon), 1e-12)))
    indices: List[int] = []

    for di in range(-steps, steps + 1):
        i = i_center + di
        ii = _periodic_lon_index(i, nlon, is_global)
        if ii is None:
            continue
        # shortest signed distance on the circle
        d = abs((lon[ii] - center_deg + 180.0) % 360.0 - 180.0)
        if d <= halfspan_deg + 1e-12:
            indices.append(ii)

    out = sorted(set(indices))
    return out


# ---- ellipsoidal method (with anisotropy) ----
def ellipsoidal_distance_method(
    lon_deg: np.ndarray,
    lat_deg: np.ndarray,
    gaussr_m: float,
    distance: Literal["fortran", "vincenty"] = "fortran",
    mask: np.ndarray | None = None,
    logger: Optional[Logger] = None,
) -> csr_matrix:
    """Construct a sparse filter matrix W such that y = W @ x (meters, SI).

    Args:
        lon_deg: 1D longitudes (degrees, regular spacing, ascending).
        lat_deg: 1D latitudes (degrees, regular spacing).
        gaussr_m: Gaussian half-width radius in **meters** (weight=0.5 at distance gaussr).
        distance: "fortran" uses spherical distance with anisotropy scaling;
                  "vincenty" uses ellipsoidal geodesic distance.
        mask: Optional boolean array of shape (nlat, nlon). True = valid cell.

    Returns:
        CSR sparse matrix with shape (nlat*nlon, nlat*nlon).
    """
    _log = _resolve_logger(logger)
    lon = np.asarray(lon_deg, dtype=np.float64)
    lat = np.asarray(lat_deg, dtype=np.float64)
    if mask is not None:
        mask = np.asarray(mask, dtype=bool)

    nlat = lat.size
    nlon = lon.size
    if nlat < 2 or nlon < 2:
        msg = (
            "ellipsoidal_distance_method: lon/lat require at least two points each "
            f"(got nlat={nlat}, nlon={nlon})."
        )
        _log.error(msg)
        raise ValueError(msg)

    if mask is not None and mask.shape != (nlat, nlon):
        _log.error(
            "ellipsoidal_distance_method: mask shape mismatch — expected (%d, %d), got %s.",
            nlat, nlon, tuple(mask.shape)
        )
        raise ValueError("mask must have shape (nlat, nlon) matching lat/lon.")

    lonstep = abs(lon[-1] - lon[0]) / (nlon - 1)
    is_global = nlon >= int(360.0 / max(lonstep, 1e-12)) - 1

    latweight = np.cos(lat * deg_2_rad)

    xt_m, fx, dxtab_m = _kernel_table(gaussr_m)

    # flattening via radii at equator/pole
    ra = radius_equator  # equatorial radius (IERS)
    rb = radius_polar  # polar radius (IERS)
    f_flat = (ra - rb) / ra
    rows: List[int] = []
    cols: List[int] = []
    data: List[float] = []

    sumwij = np.zeros((nlat, nlon), dtype=np.float64)

    for ilat in range(nlat):
        lat_src = float(lat[ilat])
        w_lat = float(latweight[ilat])
        radius_m = ellipsoid_radius(lat_src * deg_2_rad)
        if mask is not None and not mask[ilat, :].any():
            continue  # skip source row if all masked
            # choose distance function once per source-row (avoids branching in inner loop)
        if distance == "fortran":
            def _dist(lat_s, lon_s, lat_t, lon_t):
                return _anisotropy_scaled_distance_m(
                    lat_s, lon_s, lat_t, lon_t,
                    radius_m=radius_m,
                    flattening=f_flat,
                )
        elif distance == "vincenty":
            def _dist(lat_s, lon_s, lat_t, lon_t):
                # Use explicit, descriptive names for the ellipsoid axes.
                return _vincenty_distance(
                    lat_s, lon_s, lat_t, lon_t,
                    radius_equator=ra,   # equatorial radius
                    radius_polar=rb,   # polar radius
                )
        else:
            _log.error("Unsupported distance method '%s'; expected 'fortran' or 'vincenty'.", distance)
            raise ValueError('distance must be "fortran" or "vincenty"')

        # angular halfspans in degrees corresponding to kernel cutoff
        maxdist_m = float(xt_m[-1])
        fdlat = (maxdist_m / radius_m) * rad_2_deg
        coslat = max(math.cos(lat_src * deg_2_rad), 1e-12)
        fdlon = (maxdist_m / (radius_m * coslat)) * rad_2_deg

        j0, j1 = _lat_bounds_indices(lat, lat_src, fdlat)

        for ilon in range(nlon):
            lon_src = float(lon[ilon])
            i_indices = _lon_neighborhood_indices(lon, lon_src, fdlon, is_global, logger=_log)
            if mask is not None and not mask[ilat, ilon]:
                continue  # skip source if masked

            for j in range(j0, j1 + 1):
                lat_tgt = float(lat[j])

                for ii0 in i_indices:
                    if mask is not None and not mask[j, ii0]:
                        continue  # skip target if masked
                    lon_tgt = float(lon[ii0])

                    dist_m = _dist(lat_src, lon_src, lat_tgt, lon_tgt)

                    wij = _kernel_weight(dist_m, xt_m, fx, dxtab_m)
                    if wij <= 0.0:
                        continue

                    wij *= w_lat

                    sumwij[j, ii0] += wij

                    tgt = j * nlon + ii0
                    src = ilat * nlon + ilon

                    rows.append(tgt)
                    cols.append(src)
                    data.append(wij)

    W = csr_matrix((np.array(data), (np.array(rows), np.array(cols))),
                   shape=(nlat * nlon, nlat * nlon))

    # normalize rows
    sums = sumwij.reshape(-1)
    nz_rows = np.where(sums != 0.0)[0]
    if nz_rows.size:
        inv = np.zeros_like(sums)
        inv[nz_rows] = 1.0 / sums[nz_rows]
        W = csr_matrix(W.multiply(inv[:, None]))

    return W


# ---- spherical distance method ----
def spherical_distance_method(
    lon_deg: np.ndarray,
    lat_deg: np.ndarray,
    gaussr_m: float,
    radius_m: float | None = None,
    mask: np.ndarray | None = None,
    logger: Optional[Logger] = None,
) -> csr_matrix:
    """
    Construct a sparse filter matrix using spherical (great-circle) distance.

    This function uses great-circle distance on a sphere, which is significantly
    faster than ellipsoidal methods and suitable for global applications where
    the Earth's ellipsoidal shape is not critical.

    Args:
        lon_deg: 1D longitudes (degrees, regular spacing, ascending).
        lat_deg: 1D latitudes (degrees, regular spacing).
        gaussr_m: Gaussian half-width radius in **meters** (weight=0.5 at distance gaussr).
        radius_m: Sphere radius in meters. If None, uses mean Earth radius
                 ( (2*a + b)/3 ) which gives a good approximation.
        mask: Optional boolean array of shape (nlat, nlon). True = valid cell.
        logger: Optional Logger instance.

    Returns:
        CSR sparse matrix with shape (nlat*nlon, nlat*nlon).
    """
    _log = _resolve_logger(logger)
    lon = np.asarray(lon_deg, dtype=np.float64)
    lat = np.asarray(lat_deg, dtype=np.float64)
    if mask is not None:
        mask = np.asarray(mask, dtype=bool)

    nlat = lat.size
    nlon = lon.size
    if nlat < 2 or nlon < 2:
        msg = (
            "spherical_distance_method: lon/lat require at least two points each "
            f"(got nlat={nlat}, nlon={nlon})."
        )
        _log.error(msg)
        raise ValueError(msg)

    if mask is not None and mask.shape != (nlat, nlon):
        _log.error(
            "spherical_distance_method: mask shape mismatch — expected (%d, %d), got %s.",
            nlat, nlon, tuple(mask.shape)
        )
        raise ValueError("mask must have shape (nlat, nlon) matching lat/lon.")

    # Use mean Earth radius if not specified
    if radius_m is None:
        radius_m = earth_mean_radius_m
        _log.debug("Using mean Earth radius: %.3f km", radius_m / 1000.0)

    lonstep = abs(lon[-1] - lon[0]) / (nlon - 1)
    is_global = nlon >= int(360.0 / max(lonstep, 1e-12)) - 1

    latweight = np.cos(lat * deg_2_rad)

    xt_m, fx, dxtab_m = _kernel_table(gaussr_m)

    rows: List[int] = []
    cols: List[int] = []
    data: List[float] = []

    sumwij = np.zeros((nlat, nlon), dtype=np.float64)

    for ilat in range(nlat):
        lat_src = float(lat[ilat])
        w_lat = float(latweight[ilat])

        if mask is not None and not mask[ilat, :].any():
            continue  # skip source row if all masked

        # Define the spherical distance function for this source row
        def _spherical_dist(lat_s, lon_s, lat_t, lon_t):
            dist_m, _ = _spherical_distance(lat_s, lon_s, lat_t, lon_t, radius_m)
            return dist_m

        # angular halfspans in degrees corresponding to kernel cutoff
        maxdist_m = float(xt_m[-1])
        fdlat = (maxdist_m / radius_m) * rad_2_deg
        coslat = max(math.cos(lat_src * deg_2_rad), 1e-12)
        fdlon = (maxdist_m / (radius_m * coslat)) * rad_2_deg

        j0, j1 = _lat_bounds_indices(lat, lat_src, fdlat)

        for ilon in range(nlon):
            lon_src = float(lon[ilon])
            i_indices = _lon_neighborhood_indices(lon, lon_src, fdlon, is_global, logger=_log)
            if mask is not None and not mask[ilat, ilon]:
                continue  # skip source if masked

            for j in range(j0, j1 + 1):
                lat_tgt = float(lat[j])

                for ii0 in i_indices:
                    if mask is not None and not mask[j, ii0]:
                        continue  # skip target if masked
                    lon_tgt = float(lon[ii0])

                    dist_m = _spherical_dist(lat_src, lon_src, lat_tgt, lon_tgt)

                    wij = _kernel_weight(dist_m, xt_m, fx, dxtab_m)
                    if wij <= 0.0:
                        continue

                    wij *= w_lat

                    sumwij[j, ii0] += wij

                    tgt = j * nlon + ii0
                    src = ilat * nlon + ilon

                    rows.append(tgt)
                    cols.append(src)
                    data.append(wij)

    W = csr_matrix((np.array(data), (np.array(rows), np.array(cols))),
                   shape=(nlat * nlon, nlat * nlon))

    # normalize rows
    sums = sumwij.reshape(-1)
    nz_rows = np.where(sums != 0.0)[0]
    if nz_rows.size:
        inv = np.zeros_like(sums)
        inv[nz_rows] = 1.0 / sums[nz_rows]
        W = csr_matrix(W.multiply(inv[:, None]))

    return W


def _mask_crc(mask: np.ndarray | None) -> int:
    if mask is None:
        return 0
    packed = np.packbits(mask.astype(np.uint8), bitorder="little")
    return int(hashlib.sha256(packed.tobytes()).hexdigest(), 16) & ((1 << 32) - 1)


def _spherical_cache_key(
    lon: np.ndarray,
    lat: np.ndarray,
    gaussr_m: float,
    radius_m: float,
    mask_crc: int,
) -> str:
    """Build a stable SHA256 key for spherical method cache."""
    payload = {
        "method": "spherical",
        "shape": [int(lat.size), int(lon.size)],
        "lon0": float(lon[0]),
        "lonn": float(lon[-1]),
        "lat0": float(lat[0]),
        "latn": float(lat[-1]),
        "gaussr_m": float(gaussr_m),
        "radius_m": float(radius_m),
        "mask_crc32": int(mask_crc),
    }
    j = json.dumps(payload, sort_keys=True).encode("utf-8")
    return hashlib.sha256(j).hexdigest()


def _ellipsoidal_cache_key(
    lon: np.ndarray,
    lat: np.ndarray,
    gaussr_m: float,
    distance: Literal["fortran", "vincenty"],
    mask_crc: int,
) -> str:
    """Build a stable SHA256 key for ellipsoidal method cache."""
    payload = {
        "method": "ellipsoidal",
        "distance": distance,
        "shape": [int(lat.size), int(lon.size)],
        "lon0": float(lon[0]),
        "lonn": float(lon[-1]),
        "lat0": float(lat[0]),
        "latn": float(lat[-1]),
        "gaussr_m": float(gaussr_m),
        "radius_equator": float(radius_equator),
        "radius_polar": float(radius_polar),
        "mask_crc32": int(mask_crc),
    }
    j = json.dumps(payload, sort_keys=True).encode("utf-8")
    return hashlib.sha256(j).hexdigest()


def _cache_path(cache_dir: Path, key: str) -> Path:
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir / f"W_{key}.npz"


def _save_csr(path: Path, W: csr_matrix) -> None:
    np.savez_compressed(
        path,
        data=W.data,
        indices=W.indices,
        indptr=W.indptr,
        shape=np.array(W.shape, dtype=np.int64),
    )


def _load_csr(path: Path, logger: Optional[Logger] = None) -> csr_matrix:
    """
    Load a CSR matrix saved by `_save_csr` from a .npz file.

    Logs helpful diagnostics and validates presence of required arrays.

    Args:
        path: Path to the .npz file.
        logger: Optional logger to use for messages.

    Returns:
        A scipy.sparse.csr_matrix instance.
    """
    _log = _resolve_logger(logger)

    if not path.exists():
        _log.error("CSR load failed: file does not exist: %s", path)
        raise FileNotFoundError(f"CSR file not found: {path}")

    if path.suffix.lower() != ".npz":
        _log.warning("Expected .npz file for CSR load but got: %s", path)

    try:
        with np.load(path, allow_pickle=False) as npz:
            try:
                data = npz["data"]
                indices = npz["indices"]
                indptr = npz["indptr"]
                shape = tuple(npz["shape"].tolist())
            except KeyError as ke:
                _log.error("CSR load failed: missing key %s in %s", ke, path)
                raise

        W = csr_matrix((data, indices, indptr), shape=shape)  # type: ignore[arg-type]
        _log.debug("Loaded CSR: shape=%s nnz=%d from %s", W.shape, W.nnz, path)
        return W
    except Exception as e:
        _log.error("CSR load failed at %s: %s", path, e)
        raise


def ellipsoidal_distance_method_cached(
    lon_deg: np.ndarray,
    lat_deg: np.ndarray,
    gaussr_m: float,
    distance: Literal["fortran", "vincenty"] = "fortran",
    mask: np.ndarray | None = None,
    cache_dir: str | Path = ".gridgauss_cache",
    logger: Optional[Logger] = None,
) -> tuple[csr_matrix, bool]:
    """
    Return cached weights if available; otherwise build and cache (meters).

    Checks the provided cache_dir for a precomputed weight matrix matching
    the input geometry, parameters, and mask checksum.

    Args:
        lon_deg: 1D array of longitudes in degrees.
        lat_deg: 1D array of latitudes in degrees.
        gaussr_m: Gaussian filter radius (half-width) in meters.
        distance: Geodesic distance method ("fortran" or "vincenty").
        mask: Optional boolean mask of shape (nlat, nlon). True indicates
            a valid cell. Changes in the mask alter the cache key.
        cache_dir: Directory to store or retrieve cached .npz files.
        logger: Optional Logger instance for diagnostics.

    Returns:
        tuple[csr_matrix, bool]:
            - W: The sparse weight matrix (CSR format).
            - from_cache: True if loaded from disk, False if computed fresh.
    """
    _log = _resolve_logger(logger)
    lon = np.asarray(lon_deg, dtype=np.float64)
    lat = np.asarray(lat_deg, dtype=np.float64)
    mcrc = _mask_crc(mask)
    key = _ellipsoidal_cache_key(lon, lat, gaussr_m, distance, mcrc)
    cpath = _cache_path(Path(cache_dir), key)
    if cpath.exists():
        _log.debug("Ellipsoidal cache hit: %s", cpath)
        try:
            W = _load_csr(cpath)
            from_cache = True
        except Exception as e:
            _log.warning("Cache read failed at %s (%s); rebuilding.", cpath, e)
            W = ellipsoidal_distance_method(lon, lat, gaussr_m, distance=distance, mask=mask)
            _save_csr(cpath, W)
            from_cache = False
    else:
        _log.debug("Ellipsoidal cache miss: building weights (key=%s)", key)
        W = ellipsoidal_distance_method(lon, lat, gaussr_m, distance=distance, mask=mask)
        _save_csr(cpath, W)
        from_cache = False

    return W, from_cache


def spherical_distance_method_cached(
    lon_deg: np.ndarray,
    lat_deg: np.ndarray,
    gaussr_m: float,
    radius_m: float | None = None,
    mask: np.ndarray | None = None,
    cache_dir: str | Path = ".gridgauss_cache",
    logger: Optional[Logger] = None,
) -> tuple[csr_matrix, bool]:
    """
    Return cached spherical weights if available; otherwise build and cache.

    Args:
        lon_deg: 1D array of longitudes in degrees.
        lat_deg: 1D array of latitudes in degrees.
        gaussr_m: Gaussian filter radius (half-width) in meters.
        radius_m: Sphere radius in meters. If None, uses mean Earth radius.
        mask: Optional boolean mask of shape (nlat, nlon). True indicates a valid cell.
        cache_dir: Directory to store or retrieve cached .npz files.
        logger: Optional Logger instance for diagnostics.

    Returns:
        tuple[csr_matrix, bool]:
            - W: The sparse weight matrix (CSR format).
            - from_cache: True if loaded from disk, False if computed fresh.
    """
    _log = _resolve_logger(logger)
    lon = np.asarray(lon_deg, dtype=np.float64)
    lat = np.asarray(lat_deg, dtype=np.float64)

    # Use mean Earth radius if not specified
    if radius_m is None:
        radius_m = earth_mean_radius_m

    mcrc = _mask_crc(mask)
    key = _spherical_cache_key(lon, lat, gaussr_m, radius_m, mcrc)
    cpath = _cache_path(Path(cache_dir), key)

    if cpath.exists():
        _log.debug("Spherical cache hit: %s", cpath)
        try:
            W = _load_csr(cpath)
            from_cache = True
        except Exception as e:
            _log.warning("Cache read failed at %s (%s); rebuilding.", cpath, e)
            W = spherical_distance_method(lon, lat, gaussr_m, radius_m=radius_m, mask=mask)
            _save_csr(cpath, W)
            from_cache = False
    else:
        _log.debug("Spherical cache miss: building weights (key=%s)", key)
        W = spherical_distance_method(lon, lat, gaussr_m, radius_m=radius_m, mask=mask)
        _save_csr(cpath, W)
        from_cache = False

    return W, from_cache


def apply_filter(
    W: csr_matrix,
    data_2d: np.ndarray,
    mask: np.ndarray | None = None,
    fill_value: float | None = None,
    ignore_nans: bool = False,
    logger: Optional[Logger] = None,
) -> np.ndarray:
    """Apply precomputed weights to a single 2D (lat, lon) field.

    Args:
        W: CSR matrix returned by ``build_weights`` with shape (ny*nx, ny*nx).
        data_2d: Input array of shape (ny, nx).
        mask: Optional boolean array of shape (ny, nx). If provided with
            ``fill_value``, cells where ``mask`` is False are set to ``fill_value``
            after filtering.
        fill_value: Value to assign to masked output cells (for example ``np.nan``).
        ignore_nans: If True, use NaN-aware renormalized filtering. NaN values in the
            input are treated as missing data and excluded from the weighted average.
            The filter weights are dynamically renormalized for each output cell based
            on which input cells have valid (non-NaN) data. This prevents NaN propagation
            and avoids biasing results when some input cells are missing.

    Returns:
        np.ndarray: Filtered field of shape (ny, nx).
    """
    _log = _resolve_logger(logger)

    if not isinstance(W, csr_matrix):
        _log.error("W is not CSR; got %s", type(W))
        raise TypeError("W must be a scipy.sparse.csr_matrix")

    nlat_nlon = W.shape[0]
    try:
        nlat, nlon = data_2d.shape
    except Exception:
        _log.error("data_2d must be 2D (lat, lon); got shape=%s", getattr(data_2d, "shape", None))
        raise

    if nlat * nlon != nlat_nlon:
        _log.error("Shape mismatch: W expects %d elements but data_2d has %d (shape=%s).",
                   nlat_nlon, nlat * nlon, data_2d.shape)
        raise ValueError("Shape mismatch between W and data_2d.")

    if mask is not None:
        if mask.shape != (nlat, nlon):
            _log.error("Mask shape %s does not match data_2d shape %s.", mask.shape, (nlat, nlon))
            raise ValueError("mask shape must match data_2d shape.")
        if fill_value is None:
            _log.warning("mask provided without fill_value; masked cells will remain filtered values.")

    _log.debug("apply_filter: data shape=%s, W shape=%s, nnz=%d, dtype=%s",
               data_2d.shape, W.shape, W.nnz, data_2d.dtype)

    data_2d = np.asarray(data_2d, dtype=np.float64)

    if not ignore_nans:
        # Standard filtering: simple matrix-vector multiplication
        # This assumes all input values are valid (finite)
        x = data_2d.reshape(-1)
        y = W @ x
        out = np.asarray(y).reshape(nlat, nlon)

    else:
        # NaN-aware renormalized filtering
        x = data_2d
        valid = np.isfinite(x)
        x_flat = np.where(valid, x, 0.0).reshape(-1)
        valid_flat = valid.reshape(-1).astype(np.float64)

        numerator = W @ x_flat
        denominator = W @ valid_flat
        denominator = np.where(denominator == 0.0, np.nan, denominator)
        y = numerator / denominator
        out = np.asarray(y).reshape(nlat, nlon)

    if mask is not None and fill_value is not None:
        n_fill = int((~mask).sum())
        _log.debug("apply_filter: applying fill_value to %d masked cells.", n_fill)
        out = out.copy()
        out[~mask] = fill_value

    return out


def apply_filter_time(
    W: csr_matrix,
    data_3d: np.ndarray,
    mask: np.ndarray | None = None,
    fill_value: float | None = None,
    ignore_nans: bool = False,
    logger: Optional[Logger] = None,
) -> np.ndarray:
    """Apply a precomputed CSR weight matrix to a time stack of 2D (lat, lon) fields.

    Applies the same spatial filter to each time slice. Validates shapes and logs
    diagnostics for easier debugging.

    Args:
        W: CSR matrix returned by ``build_weights`` with shape (nlat*nlon, nlat*nlon).
        data_3d: Array with shape (t, nlat, nlon).
        mask: Optional boolean array with shape (nlat, nlon). If provided with
            ``fill_value``, masked output cells are set to ``fill_value`` in every frame.
        fill_value: Value to assign to masked output cells (for example ``np.nan``).
        ignore_nans: If True, use NaN-aware renormalized filtering. NaN values in the
            input are treated as missing data and excluded from the weighted average.
            The filter weights are dynamically renormalized for each output cell based
            on which input cells have valid (non-NaN) data. This prevents NaN propagation
            and avoids biasing results when some input cells are missing. Applied
            independently to each time slice.

    Returns:
        np.ndarray: Filtered array with shape (t, nlat, nlon).

    Raises:
        TypeError: If ``W`` is not a CSR matrix.
        ValueError: If shapes are incompatible.
    """
    _log = _resolve_logger(logger)

    if not isinstance(W, csr_matrix):
        _log.error("W is not CSR; got %s", type(W))
        raise TypeError("W must be a scipy.sparse.csr_matrix")

    if not hasattr(data_3d, "ndim") or data_3d.ndim != 3:
        _log.error("data_3d must be 3D (t, lat, lon); got shape=%s", getattr(data_3d, "shape", None))
        raise ValueError("data_3d must be a 3D array with shape (t, nlat, nlon).")

    t, nlat, nlon = data_3d.shape
    nlat_nlon = W.shape[0]
    if nlat * nlon != nlat_nlon:
        _log.error(
            "Shape mismatch: W expects %d elements per frame but data_3d has %d (nlat=%d, nlon=%d).",
            nlat_nlon, nlat * nlon, nlat, nlon,
        )
        raise ValueError("Shape mismatch between W and data_3d frames.")

    if mask is not None:
        if mask.shape != (nlat, nlon):
            _log.error("Mask shape %s does not match frame shape (%d, %d).", mask.shape, nlat, nlon)
            raise ValueError("mask shape must match (nlat, nlon).")
        if fill_value is None:
            _log.warning("mask provided without fill_value; masked cells will remain filtered values.")

    _log.debug(
        "apply_filter_time: data shape=%s, frames=%d, frame size=%dx%d, W shape=%s, nnz=%d, dtype=%s",
        data_3d.shape, t, nlat, nlon, W.shape, W.nnz, data_3d.dtype,
    )

    data_3d = np.asarray(data_3d, dtype=np.float64)

    if not ignore_nans:
        X = data_3d.reshape(t, nlat * nlon).T
        Y = W @ X
        out = np.asarray(Y.T).reshape(t, nlat, nlon)

    else:
        X = data_3d.reshape(t, nlat * nlon)
        out_flat = np.empty_like(X)

        for k in range(t):
            x = X[k]
            valid = np.isfinite(x)
            x_valid = np.where(valid, x, 0.0)
            numerator = W @ x_valid
            denominator = W @ valid.astype(np.float64)
            denominator = np.where(denominator == 0.0, np.nan, denominator)
            out_flat[k] = numerator / denominator

        out = out_flat.reshape(t, nlat, nlon)

    if mask is not None and fill_value is not None:
        n_fill_per_frame = int((~mask).sum())
        _log.debug(
            "apply_filter_time: applying fill_value to %d masked cells per frame (%d total across %d frames).",
            n_fill_per_frame, n_fill_per_frame * t, t,
        )
        out = out.copy()
        out[:, ~mask] = fill_value

    return out


def propagate_uncorrelated_variance(
    W: csr_matrix,
    sigma: np.ndarray,
    logger: Optional[Logger] = None,
) -> np.ndarray:
    """
    Propagate uncorrelated grid-point standard deviations through the filter.

    Args:
        W : CSR weight matrix of shape (nlat*nlon, nlat*nlon).
        sigma : Either:
            - 2D array (nlat, nlon) of standard deviations
            - 3D array (t, nlat, nlon) for time series

    Returns:
        Array of same dimensionality as sigma with filtered standard deviations.

    Formula (per grid cell):
        var_filtered[i] = sum_j( w_ij^2 * sigma_j^2 ) / sum_j( w_ij^2 )
    """
    sigma = np.asarray(sigma, dtype=np.float64)

    if sigma.ndim == 2:
        sigma_flat = sigma.reshape(-1)
        sigma2 = sigma_flat**2

        W2 = W.multiply(W)
        num = W2 @ sigma2
        denom = np.asarray(W2.sum(axis=1)).reshape(-1)
        denom = np.where(denom == 0.0, np.nan, denom)

        var = num / denom
        return np.sqrt(var).reshape(sigma.shape)

    elif sigma.ndim == 3:
        t, nlat, nlon = sigma.shape
        sigma2 = (sigma.reshape(t, nlat*nlon))**2

        W2 = W.multiply(W)

        out = np.empty_like(sigma)
        for i in range(t):
            num = W2 @ sigma2[i]
            denom = np.asarray(W2.sum(axis=1)).reshape(-1)
            denom = np.where(denom == 0.0, np.nan, denom)
            out[i] = np.sqrt(num/denom).reshape(nlat, nlon)

        return out

    else:
        msg = "sigma must be 2D (nlat, nlon) or 3D (t, nlat, nlon)."
        _log = _resolve_logger(logger)
        _log.error(msg)
        raise ValueError(msg)


def propagate_full_covariance(
    W: csr_matrix,
    Sigma_xx: np.ndarray,
    logger: Optional[Logger] = None,
) -> np.ndarray:
    """
    Propagate a full covariance matrix or a time stack of covariance matrices.

    Computes:
        Sigma_out = W * Sigma_xx * W^T (for 2D input)
        Sigma_out[t] = W * Sigma_xx[t] * W^T (for 3D input)

    Args:
        W : csr_matrix, shape (N, N)
        Sigma_xx : ndarray, shape (N, N) or (t, N, N)

    Returns:
        ndarray: propagated covariance, shape (N, N) or (t, N, N)
    """
    _log = _resolve_logger(logger)

    if not isinstance(W, csr_matrix):
        msg = "W must be a CSR matrix."
        _log.error(msg)
        raise TypeError(msg)

    if Sigma_xx.ndim == 2:
        N = W.shape[0]
        if Sigma_xx.shape != (N, N):
            msg = (
                f"Sigma_xx must have shape ({N}, {N}) for 2D input, "
                f"got {Sigma_xx.shape}."
            )
            _log.error(msg)
            raise ValueError(msg)

        temp = W @ Sigma_xx
        out = temp @ W.T
        if hasattr(out, "toarray"):
            out = out.toarray()
        return out

    elif Sigma_xx.ndim == 3:
        t, N1, N2 = Sigma_xx.shape
        N = W.shape[0]

        if N1 != N or N2 != N:
            msg = (
                f"Sigma_xx must have shape (t, {N}, {N}) for 3D input, "
                f"got {Sigma_xx.shape}."
            )
            _log.error(msg)
            raise ValueError(msg)

        out = np.empty((t, N, N), dtype=float)

        for i in range(t):
            temp = W @ Sigma_xx[i]
            cov_i = temp @ W.T
            if hasattr(cov_i, "toarray"):
                cov_i = cov_i.toarray()
            out[i] = cov_i

        return out

    else:
        msg = (
            "Sigma_xx must be 2D (N, N) or 3D (t, N, N)."
        )
        _log.error(msg)
        raise ValueError(msg)


def _get_mask_if_present(
    grid: Grid2DObject | Grid3DObject,
    nlat: int,
    nlon: int,
    logger: Optional[Logger] = None,
) -> np.ndarray | None:
    """Return a boolean mask (nlat, nlon) if present on the grid; otherwise None."""
    _log = _resolve_logger(logger)
    m = getattr(grid, "mask", None)
    if m is None:
        return None
    m = np.asarray(m, dtype=bool)
    if m.shape != (nlat, nlon):
        _log.error(
            "_get_mask_if_present: mask shape mismatch; expected (%d, %d), got %s.",
            nlat, nlon, tuple(m.shape),
        )
        raise ValueError("mask must have shape (nlat, nlon) matching lat/lon.")
    return m


def apply_gaussian_filter_grid(
    grid: Grid2DObject | Grid3DObject,
    *,
    gaussr_km: float | None = None,
    gaussr_m: float | None = None,
    distance: Literal["spherical", "vincenty", "fortran"] = "spherical",
    radius_m: float | None = None,  # For spherical method
    cache_dir: str | Path | None = ".gridgauss_cache",
    fill_value: float | None = None,
    in_place: bool = False,
    full_covariance: np.ndarray | None = None,
    sigma: np.ndarray | None = None,
    ignore_nans: bool = False,
    logger: Optional[Logger] = None,
) -> Grid2DObject | Grid3DObject:
    """
    Apply a Gaussian spatial filter to a Grid2DObject or Grid3DObject.

    Supports three distance methods:
    - "spherical": Great-circle distance on a sphere (fastest)
    - "vincenty": Exact ellipsoidal geodesic distance (most accurate)
    - "fortran": Legacy method with anisotropy scaling (ellipsoidal approximation)

    Args:
        grid: Grid2DObject or Grid3DObject.
        gaussr_km: Gaussian kernel half-width in kilometers.
        gaussr_m: Legacy meters argument; used only if gaussr_km is None.
        distance: Distance method - "spherical", "vincenty", or "fortran".
        radius_m: Sphere radius in meters (only for distance="spherical").
        cache_dir: Directory for caching the sparse weight matrix.
        fill_value: Value assigned to masked target cells.
        in_place: Whether to modify the input grid.
        logger: Optional logger instance.
        full_covariance: Full covariance matrix (num_cells x num_cells).

    Returns:
        A filtered Grid2DObject or Grid3DObject, possibly with:
            .covariance  (full propagated covariance matrix)
    """
    _log = _resolve_logger(logger)

    # Normalize kernel half-width
    if gaussr_km is not None and gaussr_m is not None:
        msg = "Both gaussr_km and gaussr_m were provided; using gaussr_km and ignoring gaussr_m."
        _log.warning(msg)
    if gaussr_km is not None:
        gaussr_m_eff = float(gaussr_km) * 1_000.0
    elif gaussr_m is not None:
        msg = "gaussr_m is deprecated; pass gaussr_km instead."
        _log.warning(msg)
        gaussr_m_eff = float(gaussr_m)
    else:
        msg = "Missing kernel half-width: provide gaussr_km (preferred) or gaussr_m."
        _log.error(msg)
        raise ValueError(msg)

    if not np.isfinite(gaussr_m_eff) or gaussr_m_eff <= 0:
        msg = f"Invalid kernel half-width: gaussr_m={gaussr_m_eff}."
        _log.error(msg)
        raise ValueError(msg)

    # Validate distance method
    valid_distances = ["spherical", "vincenty", "fortran"]
    distance = str(distance).strip().lower()
    if distance not in valid_distances:
        msg = f"Unsupported distance '{distance}'; expected one of: {valid_distances}"
        _log.error(msg)
        raise ValueError(msg)

    # Geometry
    lon = np.asarray(grid.lon, float)
    lat = np.asarray(grid.lat, float)
    nlat, nlon = lat.size, lon.size

    # Data and mask
    data = np.asarray(grid.grid)
    if data.ndim not in (2, 3):
        msg = 'grid.grid must be 2D (nlat, nlon) or 3D (ntime, nlat, nlon).'
        _log.error(msg)
        raise ValueError(msg)
    mask = _get_mask_if_present(grid, nlat, nlon, logger=_log)

    # Validate sigma shape if provided
    if sigma is not None:
        if data.ndim == 2 and sigma.shape != (nlat, nlon):
            msg = "sigma must have shape (nlat, nlon) for 2D grids."
            _log.error(msg)
            raise ValueError(msg)
        if data.ndim == 3 and sigma.shape != (data.shape[0], nlat, nlon):
            msg = "sigma must have shape (ntime, nlat, nlon) for 3D grids."
            _log.error(msg)
            raise ValueError(msg)

    # Build weights based on selected distance method
    if distance == "spherical":
        if cache_dir is None:
            W = spherical_distance_method(
                lon, lat, gaussr_m_eff, radius_m=radius_m, mask=mask, logger=_log
            )
        else:
            W, _ = spherical_distance_method_cached(
                lon, lat, gaussr_m_eff, radius_m=radius_m, mask=mask,
                cache_dir=cache_dir, logger=_log
            )
    else:  # distance in ["vincenty", "fortran"]
        # Both vincenty and fortran are handled by ellipsoidal_distance_method
        if cache_dir is None:
            W = ellipsoidal_distance_method(
                lon, lat, gaussr_m_eff,
                distance=distance,  # Pass through "vincenty" or "fortran"
                mask=mask,
                logger=_log
            )
        else:
            W, _ = ellipsoidal_distance_method_cached(
                lon, lat, gaussr_m_eff,
                distance=distance,  # Pass through "vincenty" or "fortran"
                mask=mask,
                cache_dir=cache_dir,
                logger=_log
            )

    # Rest of the function remains the same...
    # Apply filter (2D case)
    _log_obj = _resolve_logger(logger)
    if data.ndim == 2:
        filtered = apply_filter(
            W,
            data,
            mask=mask,
            fill_value=fill_value,
            ignore_nans=ignore_nans,
            logger=_log,
        )

        if in_place:
            grid.grid = filtered

            if sigma is not None:
                grid.sigma = propagate_uncorrelated_variance(W, sigma, logger=_log)
            else:
                grid.sigma = None

            if full_covariance is not None:
                grid.covariance = propagate_full_covariance(W, full_covariance, logger=_log)
            else:
                grid.covariance = None

            return grid
        else:
            out = type(grid)(
                logger=_log_obj,
                grid=filtered,
                lon=lon,
                lat=lat,
            )

        if sigma is not None:
            out.sigma = propagate_uncorrelated_variance(W, sigma, logger=_log)

        if full_covariance is not None:
            out.covariance = propagate_full_covariance(W, full_covariance, logger=_log)

        return out

    # 3D case
    filtered3d = apply_filter_time(
        W,
        data,
        mask=mask,
        fill_value=fill_value,
        ignore_nans=ignore_nans,
        logger=_log
    )
    dates = getattr(grid, "dates", None)

    if in_place:
        grid.grid = filtered3d

        if sigma is not None:
            grid.sigma = propagate_uncorrelated_variance(W, sigma, logger=_log)
        else:
            grid.sigma = None

        if full_covariance is not None:
            grid.covariance = propagate_full_covariance(W, full_covariance, logger=_log)
        else:
            grid.covariance = None

        return grid

    try:
        out = type(grid)(
            logger=_log_obj,
            grid=filtered3d,
            dates=dates,
            lon=lon,
            lat=lat,
        )
    except Exception:
        class _Grid3D:
            def __init__(self, grid, dates, lon, lat):
                self.grid = grid
                self.dates = dates
                self.lon = lon
                self.lat = lat
        out = _Grid3D(grid=filtered3d, dates=dates, lon=lon, lat=lat)

    if sigma is not None:
        out.sigma = propagate_uncorrelated_variance(W, sigma, logger=_log)

    if full_covariance is not None:
        out.covariance = propagate_full_covariance(W, full_covariance, logger=_log)

    return out
