navi/backend/services/navi_offroute/friction.py

"""
Friction layer reader for OFFROUTE.

Reads friction values from the WorldCover friction VRT and resamples
to match the elevation grid dimensions.
"""
import os
from pathlib import Path
from typing import Tuple, Optional

import numpy as np

try:
    import rasterio
    from rasterio.windows import from_bounds
    from rasterio.enums import Resampling
except ImportError:
    raise ImportError("rasterio is required for friction layer support")

# Default path to the friction VRT (single source of truth); env-overridable.
DEFAULT_FRICTION_PATH = Path("/mnt/nav/worldcover/friction/friction_conus.vrt")


def friction_vrt_path() -> Path:
    """Friction VRT path, env-overridable via NAVI_OFFROUTE_FRICTION_VRT."""
    return Path(os.environ.get("NAVI_OFFROUTE_FRICTION_VRT", str(DEFAULT_FRICTION_PATH)))


class FrictionReader:
    """Reader for WorldCover friction raster."""

    def __init__(self, friction_path: Path = None):
        self.friction_path = Path(friction_path) if friction_path else friction_vrt_path()
        self._dataset = None
    
    def _open(self):
        """Lazy open the dataset."""
        if self._dataset is None:
            if not self.friction_path.exists():
                raise FileNotFoundError(f"Friction VRT not found at {self.friction_path}")
            self._dataset = rasterio.open(self.friction_path)
        return self._dataset
    
    def get_friction_grid(
        self,
        south: float,
        north: float,
        west: float,
        east: float,
        target_shape: Tuple[int, int]
    ) -> np.ndarray:
        """
        Get friction values for a bounding box, resampled to target shape.
        
        Args:
            south, north, west, east: Bounding box coordinates
            target_shape: (rows, cols) to resample to (matches elevation grid)
        
        Returns:
            np.ndarray of uint8 friction values, same shape as target_shape.
            Values: 10-40 = friction multiplier (divide by 10)
                    255 = impassable
                    0 = nodata (treat as impassable)
        """
        ds = self._open()
        
        # Create a window from the bounding box
        window = from_bounds(west, south, east, north, ds.transform)
        
        # Read with resampling to target shape
        # Use nearest neighbor for categorical data
        friction = ds.read(
            1,
            window=window,
            out_shape=target_shape,
            resampling=Resampling.nearest
        )
        
        return friction
    
    def sample_point(self, lat: float, lon: float) -> int:
        """Sample friction value at a single point."""
        ds = self._open()
        
        # Get pixel coordinates
        row, col = ds.index(lon, lat)
        
        # Check bounds
        if row < 0 or row >= ds.height or col < 0 or col >= ds.width:
            return 0  # Out of bounds = nodata
        
        # Read single pixel
        window = rasterio.windows.Window(col, row, 1, 1)
        value = ds.read(1, window=window)
        return int(value[0, 0])
    
    def close(self):
        """Close the dataset."""
        if self._dataset is not None:
            self._dataset.close()
            self._dataset = None


def friction_to_multiplier(friction: np.ndarray) -> np.ndarray:
    """
    Convert friction values to cost multipliers.
    
    Args:
        friction: uint8 array of friction values
        
    Returns:
        float32 array of multipliers.
        Values 10-40 become 1.0-4.0 (divide by 10).
        Values 0 or 255 become np.inf (impassable).
    """
    multiplier = friction.astype(np.float32) / 10.0
    
    # Mark impassable cells
    multiplier[friction == 0] = np.inf    # nodata
    multiplier[friction == 255] = np.inf  # water/impassable
    
    return multiplier


if __name__ == "__main__":
    print("Testing FrictionReader...")
    
    reader = FrictionReader()
    
    # Test point sampling - Murtaugh Lake (should be water = 255)
    lake_lat, lake_lon = 42.47, -114.15
    lake_friction = reader.sample_point(lake_lat, lake_lon)
    print(f"Murtaugh Lake ({lake_lat}, {lake_lon}): friction = {lake_friction}")
    print(f"  Expected: 255 (water/impassable)")
    
    # Test grid read for small bbox
    friction = reader.get_friction_grid(
        south=42.4, north=42.5, west=-114.2, east=-114.1,
        target_shape=(100, 100)
    )
    print(f"\nGrid test shape: {friction.shape}")
    print(f"Unique values: {np.unique(friction)}")
    print(f"Water cells (255): {np.sum(friction == 255)}")
    
    reader.close()
    print("\nFrictionReader test complete.")