Choosing a Projected CRS for Basin-Scale Hydrology

import logging

import geopandas as gpd
import numpy as np
from pyproj import CRS
from pyproj.aoi import AreaOfInterest
from pyproj.database import query_utm_crs_info

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


def recommend_crs(catchment_path: str, purpose: str = "area") -> dict:
    """
    Recommend a projected CRS for a catchment, and say why.

    Parameters
    ----------
    catchment_path : Catchment polygon in any CRS.
    purpose        : "area" for contributing area and volumes, "shape" for
                     slope and flow direction, "distance" for travel time.
    """
    gdf = gpd.read_file(catchment_path).to_crs(4326)
    minx, miny, maxx, maxy = gdf.total_bounds
    lon_c, lat_c = (minx + maxx) / 2.0, (miny + maxy) / 2.0
    span_deg = max(maxx - minx, maxy - miny)
    span_km = span_deg * 111.0 * np.cos(np.radians(lat_c))

    utm_zone = int((lon_c + 180) // 6) + 1
    zone_of_min = int((minx + 180) // 6) + 1
    zone_of_max = int((maxx + 180) // 6) + 1
    straddles = zone_of_min != zone_of_max

    log.info("Extent %.1f × %.1f degrees (~%.0f km), centre (%.3f, %.3f)",
             maxx - minx, maxy - miny, span_km, lat_c, lon_c)
    if straddles:
        log.warning("Basin straddles UTM zones %d–%d — a single zone will "
                    "distort one side", zone_of_min, zone_of_max)

    # --- Under about 400 km and inside one zone, UTM is the pragmatic choice:
    # it is conformal, widely understood, and its area error is negligible. ---
    if not straddles and span_km < 400:
        candidates = query_utm_crs_info(
            datum_name="WGS 84",
            area_of_interest=AreaOfInterest(minx, miny, maxx, maxy),
        )
        crs = CRS.from_epsg(candidates[0].code) if candidates else None
        rationale = (f"single UTM zone {utm_zone}, extent {span_km:.0f} km — "
                     "conformal, negligible area error at this size")
    elif purpose == "area":
        # A custom Albers with standard parallels at one sixth and five sixths
        # of the latitude range minimises distortion over the extent.
        sp1 = miny + (maxy - miny) / 6.0
        sp2 = maxy - (maxy - miny) / 6.0
        crs = CRS.from_proj4(
            f"+proj=aea +lat_1={sp1:.4f} +lat_2={sp2:.4f} "
            f"+lat_0={lat_c:.4f} +lon_0={lon_c:.4f} "
            "+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
        rationale = ("custom Albers equal-area fitted to the extent — areas "
                     "preserved, which is what the analysis measures")
    else:
        crs = CRS.from_proj4(
            f"+proj=tmerc +lat_0={lat_c:.4f} +lon_0={lon_c:.4f} +k=0.9996 "
            "+x_0=500000 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
        rationale = ("custom transverse Mercator centred on the basin — "
                     "conformal across a zone-straddling extent")

    log.info("Recommended: %s", rationale)
    return {"crs": crs, "rationale": rationale, "span_km": span_km,
            "straddles_utm_zone": straddles}


def verify_area(catchment_path: str, crs, known_area_km2: float,
                tolerance_pct: float = 1.0) -> dict:
    """
    Check a projection against an independently known area.

    A systematic offset here is a projection problem. Scattered differences
    across several features are delineation differences, which is a different
    investigation entirely.
    """
    gdf = gpd.read_file(catchment_path).to_crs(crs)
    computed = float(gdf.geometry.area.sum()) / 1e6
    diff_pct = 100.0 * (computed - known_area_km2) / known_area_km2

    level = log.info if abs(diff_pct) <= tolerance_pct else log.error
    level("Area in this CRS: %.3f km² vs known %.3f km² (%+.2f %%)",
          computed, known_area_km2, diff_pct)
    if abs(diff_pct) > tolerance_pct:
        log.error("Beyond the %.1f %% tolerance — check the projection before "
                  "assuming the delineation is wrong", tolerance_pct)
    return {"computed_km2": computed, "known_km2": known_area_km2,
            "diff_pct": diff_pct, "within_tolerance": abs(diff_pct) <= tolerance_pct}


def distortion_across_extent(crs, bounds_4326, samples: int = 5) -> dict:
    """
    Sample the area scale factor across the extent, so the distortion is a
    measured number rather than an assumption.
    """
    from shapely.geometry import box

    minx, miny, maxx, maxy = bounds_4326
    factors = []
    for lon in np.linspace(minx, maxx, samples):
        for lat in np.linspace(miny, maxy, samples):
            d = 0.01  # a small test cell, about 1 km
            cell = gpd.GeoSeries([box(lon, lat, lon + d, lat + d)], crs=4326)
            true_m2 = float(cell.to_crs(
                f"+proj=aeqd +lat_0={lat} +lon_0={lon} +datum=WGS84 +units=m"
            ).area.iloc[0])
            proj_m2 = float(cell.to_crs(crs).area.iloc[0])
            factors.append(proj_m2 / true_m2)

    arr = np.asarray(factors)
    log.info("Area scale factor across the extent: min %.5f, max %.5f, "
             "spread %.3f %%", arr.min(), arr.max(),
             100.0 * (arr.max() - arr.min()))
    return {"min": float(arr.min()), "max": float(arr.max()),
            "spread_pct": float(100.0 * (arr.max() - arr.min()))}


# --- Example usage ---
# rec = recommend_crs("basin.gpkg", purpose="area")
# verify_area("basin.gpkg", rec["crs"], known_area_km2=1482.0)