Outlet Point Mapping & Validation

# conda (recommended — ensures GDAL/PROJ binary compatibility)
conda create -n hydro python=3.11
conda activate hydro
conda install -c conda-forge geopandas rasterio shapely scipy pyproj affine richdem

# or pip inside an activated venv
pip install geopandas rasterio shapely scipy pyproj affine

import logging
import geopandas as gpd
import rasterio
from pyproj import CRS

logger = logging.getLogger(__name__)

def load_and_align_inputs(outlet_path: str, stream_path: str, dem_path: str):
    """
    Load outlet points and stream network; reproject both to the DEM's CRS.
    Returns aligned GeoDataFrames and the open rasterio dataset for the DEM.
    """
    with rasterio.open(dem_path) as dem_ds:
        dem_crs = CRS.from_user_input(dem_ds.crs)
        logger.info("DEM CRS: %s  |  bounds: %s", dem_crs.to_epsg(), dem_ds.bounds)

    outlets = gpd.read_file(outlet_path)
    streams = gpd.read_file(stream_path)

    # Reproject to the DEM's projected CRS (required for metric distance calculations)
    outlets = outlets.to_crs(dem_crs)
    streams = streams.to_crs(dem_crs)
    logger.info("Loaded %d outlets and %d stream features; both reprojected to EPSG:%s",
                len(outlets), len(streams), dem_crs.to_epsg())

    # Build R-tree spatial index on stream geometries for fast proximity queries
    _ = streams.sindex  # geopandas builds this lazily; trigger it now
    logger.info("Stream spatial index built.")
    return outlets, streams

import numpy as np
from scipy.spatial import KDTree
import rasterio
from rasterio.transform import xy

def build_stream_kdtree(acc_path: str, min_accumulation: float):
    """
    Extract centroids of all accumulation cells exceeding min_accumulation
    and index them in a KD-tree for fast nearest-neighbour lookup.
    Returns (kdtree, world_coords array of shape [N, 2]).
    """
    with rasterio.open(acc_path) as acc_ds:
        acc = acc_ds.read(1)
        transform = acc_ds.transform
        nodata = acc_ds.nodata

    # Mask out nodata and sub-threshold cells
    valid_mask = (acc >= min_accumulation)
    if nodata is not None:
        valid_mask &= (acc != nodata)

    rows, cols = np.where(valid_mask)
    # Convert row/col indices to projected world coordinates (cell centres)
    xs, ys = xy(transform, rows, cols)
    coords = np.column_stack([xs, ys])

    tree = KDTree(coords)
    logger.info(
        "KD-tree built from %d stream cells (accumulation >= %.0f).",
        len(coords), min_accumulation
    )
    return tree, coords

from shapely.geometry import Point

# D8 direction → (delta_row, delta_col) lookup
D8_OFFSETS = {1: (0, 1), 2: (1, 1), 4: (1, 0), 8: (1, -1),
              16: (0, -1), 32: (-1, -1), 64: (-1, 0), 128: (-1, 1)}

def snap_outlet(outlet_xy, tree: KDTree, stream_coords,
                acc_arr, fdir_arr, transform,
                tolerance_m: float, min_accumulation: float,
                max_trace: int = 200):
    """
    Snap a single outlet (x, y in projected metres) to the nearest valid
    stream cell. Falls back to upstream trace when no candidate is found
    within tolerance_m.

    Returns dict with snapped coords, accumulation, snap_distance, method.
    """
    result = {"original_xy": outlet_xy, "snapped_xy": None,
              "accumulation": None, "snap_distance_m": None, "method": None}

    # --- Stage 1: KD-tree nearest neighbour within tolerance ---
    idxs = tree.query_ball_point(outlet_xy, r=tolerance_m)
    if idxs:
        # Pick the candidate with the highest accumulation value
        candidates = stream_coords[idxs]
        # Sample accumulation at each candidate
        acc_vals = [acc_arr[int(round((transform[5] - y) / (-transform[4]))),
                             int(round((x - transform[2]) / transform[0]))]
                    for x, y in candidates]
        best_idx = int(np.argmax(acc_vals))
        snapped_xy = tuple(candidates[best_idx])
        snap_dist = float(np.linalg.norm(np.array(snapped_xy) - np.array(outlet_xy)))

        # --- Stage 2: flow-path validity check ---
        row = int(round((transform[5] - snapped_xy[1]) / (-transform[4])))
        col = int(round((snapped_xy[0] - transform[2]) / transform[0]))
        if _flow_path_is_valid(fdir_arr, row, col, max_trace):
            result.update(snapped_xy=snapped_xy, accumulation=float(acc_vals[best_idx]),
                          snap_distance_m=snap_dist, method="stage1_kdtree")
            logger.debug("Outlet snapped via KD-tree: dist=%.1f m, acc=%.0f",
                         snap_dist, acc_vals[best_idx])
            return result

    # --- Fallback: trace upstream from raw cell to first qualifying cell ---
    raw_row = int(round((transform[5] - outlet_xy[1]) / (-transform[4])))
    raw_col = int(round((outlet_xy[0] - transform[2]) / transform[0]))
    snapped_xy, acc_val = _upstream_trace_fallback(
        acc_arr, fdir_arr, raw_row, raw_col, min_accumulation, max_trace, transform)

    if snapped_xy:
        snap_dist = float(np.linalg.norm(np.array(snapped_xy) - np.array(outlet_xy)))
        result.update(snapped_xy=snapped_xy, accumulation=acc_val,
                      snap_distance_m=snap_dist, method="upstream_trace_fallback")
        logger.warning("Outlet used fallback trace: dist=%.1f m, acc=%.0f", snap_dist, acc_val)
    else:
        logger.error("Could not snap outlet at %s — no valid stream cell found.", outlet_xy)

    return result


def _flow_path_is_valid(fdir_arr, row, col, max_steps):
    """Trace D8 downstream; return True if path exits without hitting sink or loop."""
    visited = set()
    for _ in range(max_steps):
        cell = (row, col)
        if cell in visited:
            return False  # loop detected
        visited.add(cell)
        d = int(fdir_arr[row, col])
        if d not in D8_OFFSETS:
            return False  # sink or nodata
        dr, dc = D8_OFFSETS[d]
        row += dr
        col += dc
        if row < 0 or col < 0 or row >= fdir_arr.shape[0] or col >= fdir_arr.shape[1]:
            return True  # reached raster edge — valid drainage outlet
    return True  # trace exhausted without finding a sink


def _upstream_trace_fallback(acc_arr, fdir_arr, row, col,
                              min_acc, max_steps, transform):
    """
    Walk upstream (reverse D8 directions) from raw cell until an accumulation-
    qualifying cell is found. Returns (world_xy, acc_value) or (None, None).
    """
    REVERSE_D8 = {v: k for k, v in D8_OFFSETS.items()}
    # Invert: we need cells whose downstream neighbour is the current cell
    # Simple approach: search a small window for cells flowing into (row, col)
    for step in range(max_steps):
        if acc_arr[row, col] >= min_acc:
            from rasterio.transform import xy as _xy
            wx, wy = _xy(transform, row, col)
            return (wx, wy), float(acc_arr[row, col])
        # Move to the neighbour with the highest accumulation
        best_acc, best_row, best_col = -1, row, col
        for dr, dc in D8_OFFSETS.values():
            nr, nc = row + dr, col + dc
            if 0 <= nr < acc_arr.shape[0] and 0 <= nc < acc_arr.shape[1]:
                if acc_arr[nr, nc] > best_acc:
                    best_acc, best_row, best_col = acc_arr[nr, nc], nr, nc
        if best_row == row and best_col == col:
            break  # no improvement — genuine flat/sink area
        row, col = best_row, best_col
    return None, None

def check_topological_consistency(snapped_results: list, fdir_arr, acc_arr,
                                   transform, max_trace: int = 500):
    """
    For each snapped result, verify the flow path reaches the raster edge
    (or a downstream outlet) without entering a sink. Flags inconsistent records.
    """
    flagged = []
    for rec in snapped_results:
        if rec["snapped_xy"] is None:
            rec["topology_valid"] = False
            flagged.append(rec["original_xy"])
            continue

        x, y = rec["snapped_xy"]
        row = int(round((transform[5] - y) / (-transform[4])))
        col = int(round((x - transform[2]) / transform[0]))
        valid = _flow_path_is_valid(fdir_arr, row, col, max_trace)
        rec["topology_valid"] = valid
        if not valid:
            flagged.append(rec["original_xy"])
            logger.warning("Topological inconsistency at snapped point %s", rec["snapped_xy"])

    logger.info("Topology check complete. %d / %d outlets flagged.",
                len(flagged), len(snapped_results))
    return snapped_results, flagged

import json
import logging
from datetime import datetime, timezone
from pathlib import Path

import geopandas as gpd
import numpy as np
import rasterio
from rasterio.transform import xy as rasterio_xy
from scipy.spatial import KDTree
from shapely.geometry import Point

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")


def validate_outlet_points(
    outlet_path: str,
    dem_path: str,
    acc_path: str,
    fdir_path: str,
    stream_path: str,
    output_path: str,
    tolerance_m: float = 90.0,
    min_accumulation: float = 1000.0,
    max_trace: int = 300,
) -> gpd.GeoDataFrame:
    """
    Full outlet-point validation pipeline.

    Parameters
    ----------
    outlet_path     : Path to raw outlet points (CSV / GeoJSON / shapefile).
    dem_path        : Hydrologically conditioned DEM GeoTIFF.
    acc_path        : Flow-accumulation GeoTIFF (float64, same grid as DEM).
    fdir_path       : D8 flow-direction GeoTIFF (uint8, power-of-two encoding).
    stream_path     : Vector stream network (GeoPackage or shapefile).
    output_path     : Output GeoJSON path for validated outlets.
    tolerance_m     : Maximum snapping distance in projected metres.
    min_accumulation: Minimum upstream cell count to qualify as a stream cell.
    max_trace       : Maximum D8 steps for flow-path and fallback traces.

    Returns
    -------
    GeoDataFrame of validated outlets with audit columns.
    """
    logger.info("=== Outlet Validation Pipeline START ===")
    logger.info("Inputs: outlets=%s  acc_threshold=%.0f  tolerance=%.0fm",
                outlet_path, min_accumulation, tolerance_m)

    # 1. Load and align inputs
    with rasterio.open(dem_path) as dem_ds:
        dem_crs = dem_ds.crs
        dem_transform = dem_ds.transform

    outlets_gdf = gpd.read_file(outlet_path).to_crs(dem_crs)
    if "outlet_id" not in outlets_gdf.columns:
        outlets_gdf["outlet_id"] = outlets_gdf.index.astype(str)
        logger.warning("No 'outlet_id' column found; using row index as identifier.")

    # 2. Load raster arrays (windowed reads are viable for large DEMs; full read here)
    with rasterio.open(acc_path) as acc_ds:
        acc_arr = acc_ds.read(1).astype(np.float64)
        acc_transform = acc_ds.transform
        acc_nodata = acc_ds.nodata

    with rasterio.open(fdir_path) as fdir_ds:
        fdir_arr = fdir_ds.read(1).astype(np.int32)

    # 3. Build KD-tree from stream cells above threshold
    valid_mask = acc_arr >= min_accumulation
    if acc_nodata is not None:
        valid_mask &= acc_arr != acc_nodata
    rows_valid, cols_valid = np.where(valid_mask)
    xs, ys = rasterio_xy(acc_transform, rows_valid, cols_valid)
    stream_coords = np.column_stack([xs, ys])
    tree = KDTree(stream_coords)
    logger.info("Stream KD-tree: %d qualifying cells.", len(stream_coords))

    # 4. Snap each outlet
    results = []
    for _, row_data in outlets_gdf.iterrows():
        oid = row_data["outlet_id"]
        geom = row_data.geometry
        outlet_xy = (geom.x, geom.y)
        snap_result = snap_outlet(
            outlet_xy, tree, stream_coords, acc_arr, fdir_arr,
            acc_transform, tolerance_m, min_accumulation, max_trace
        )
        snap_result["outlet_id"] = oid
        results.append(snap_result)

    # 5. Topological consistency
    results, flagged = check_topological_consistency(
        results, fdir_arr, acc_arr, acc_transform, max_trace
    )

    # 6. Assemble output GeoDataFrame
    rows_out = []
    for rec in results:
        if rec["snapped_xy"] is None:
            geom_out = None
        else:
            geom_out = Point(rec["snapped_xy"])
        rows_out.append({
            "outlet_id": rec["outlet_id"],
            "geometry": geom_out,
            "original_x": rec["original_xy"][0],
            "original_y": rec["original_xy"][1],
            "snap_distance_m": rec.get("snap_distance_m"),
            "accumulation": rec.get("accumulation"),
            "snap_method": rec.get("method"),
            "topology_valid": rec.get("topology_valid", False),
            "validated_at": datetime.now(timezone.utc).isoformat(),
            "tolerance_m": tolerance_m,
            "min_accumulation": min_accumulation,
        })

    out_gdf = gpd.GeoDataFrame(rows_out, crs=dem_crs)
    valid_count = out_gdf["topology_valid"].sum()
    logger.info("Validation complete: %d / %d outlets fully valid.",
                valid_count, len(out_gdf))

    # 7. Export
    out_gdf.to_file(output_path, driver="GeoJSON")
    logger.info("Validated outlets written to %s", output_path)
    logger.info("=== Outlet Validation Pipeline END ===")
    return out_gdf


# Example invocation
if __name__ == "__main__":
    result = validate_outlet_points(
        outlet_path="data/raw_gauges.geojson",
        dem_path="data/dem_conditioned.tif",
        acc_path="data/flow_accumulation.tif",
        fdir_path="data/flow_direction_d8.tif",
        stream_path="data/nhd_streams.gpkg",
        output_path="outputs/validated_outlets.geojson",
        tolerance_m=90.0,
        min_accumulation=1000.0,
    )
    print(result[["outlet_id", "snap_distance_m", "accumulation", "topology_valid"]])

import matplotlib.pyplot as plt
out_gdf["snap_distance_m"].hist(bins=40)
plt.axvline(90, color="red", linestyle="--", label="Tolerance boundary")
plt.xlabel("Snap distance (m)")
plt.title("Outlet snap displacement distribution")
plt.legend()
plt.tight_layout()
plt.savefig("outputs/snap_displacement.png", dpi=150)

n_before = len(outlets_gdf)
outlets_gdf = outlets_gdf.drop_duplicates(subset=["outlet_id"], keep="first")
logger.info("Dropped %d duplicate outlet IDs.", n_before - len(outlets_gdf))

if rec.get("terrain_type") == "endorheic":
    rec["topology_valid"] = True  # internal drainage is expected
    logger.info("Outlet %s classified endorheic; topology check skipped.", rec["outlet_id"])

import os
os.environ["GDAL_DISABLE_READDIR_ON_OPEN"] = "EMPTY_DIR"
import rasterio