Repairing Invalid Catchment Geometries with make_valid

import logging

import geopandas as gpd
import numpy as np
from shapely import make_valid
from shapely.geometry import MultiPolygon, Polygon

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


def repair_catchments(
    gdf: gpd.GeoDataFrame,
    area_tolerance_pct: float = 0.01,
    keep_multipart: bool = True,
) -> gpd.GeoDataFrame:
    """
    Repair invalid catchment geometries, auditing every change.

    Parameters
    ----------
    gdf                : Catchment polygons.
    area_tolerance_pct : Area change above which a repair is reported as an error.
    keep_multipart     : Keep both lobes of a bow-tie as a MultiPolygon. Setting
                         this False keeps only the largest part, which is
                         sometimes what a downstream consumer needs — but it
                         DELETES area, so it is logged loudly.
    """
    out = gdf.copy()
    original_area = out.geometry.area.to_numpy()
    invalid = ~out.geometry.is_valid
    n_invalid = int(invalid.sum())
    log.info("%d of %d geometries invalid (%.1f %%)",
             n_invalid, len(out), 100.0 * n_invalid / max(1, len(out)))

    if n_invalid == 0:
        return out

    # --- Report WHY each one is invalid before repairing. The reason decides
    # whether the repair is safe: a self-intersection may hide real area,
    # a ring-order problem never does. ---
    from shapely.validation import explain_validity
    reasons = {}
    for idx in out.index[invalid]:
        reason = explain_validity(out.geometry.loc[idx]).split("[")[0].strip()
        reasons[reason] = reasons.get(reason, 0) + 1
    for reason, count in sorted(reasons.items(), key=lambda kv: -kv[1]):
        log.info("  %4d × %s", count, reason)

    repaired = out.geometry.apply(lambda g: make_valid(g) if g and not g.is_valid else g)

    # --- make_valid can return a GeometryCollection when the input contained
    # dangling lines as well as area. Keep only the areal parts; the lines
    # carry no catchment and would break every later area computation. ---
    def areal_only(g):
        if g is None or g.is_empty:
            return g
        if g.geom_type in ("Polygon", "MultiPolygon"):
            return g
        parts = [p for p in getattr(g, "geoms", []) if p.geom_type == "Polygon"]
        if not parts:
            log.warning("Repair produced no areal part for a geometry — dropped")
            return None
        return parts[0] if len(parts) == 1 else MultiPolygon(parts)

    repaired = repaired.apply(areal_only)

    if not keep_multipart:
        def largest_part(g):
            if g is not None and g.geom_type == "MultiPolygon":
                return max(g.geoms, key=lambda p: p.area)
            return g
        before = repaired.area.to_numpy()
        repaired = repaired.apply(largest_part)
        lost = float((before - repaired.area.to_numpy()).sum())
        if lost > 0:
            log.error("keep_multipart=False discarded %.1f m² of catchment area "
                      "across %d feature(s)", lost,
                      int((before - repaired.area.to_numpy() > 0).sum()))

    out = out.set_geometry(repaired)
    out = out[out.geometry.notna()]

    # --- The audit. A repair that changes area has edited the catchment, not
    # repaired it, and every downstream drainage area is now wrong. ---
    new_area = out.geometry.area.to_numpy()
    common = min(len(original_area), len(new_area))
    delta = new_area[:common] - original_area[:common]
    pct = 100.0 * delta / np.maximum(original_area[:common], 1e-9)
    worst = int(np.argmax(np.abs(pct))) if common else 0

    still_invalid = int((~out.geometry.is_valid).sum())
    log.info("After repair: %d invalid, total area change %.3f m² (%.5f %%)",
             still_invalid, float(delta.sum()),
             100.0 * float(delta.sum()) / float(original_area.sum()))
    if common and abs(pct[worst]) > area_tolerance_pct:
        log.error("Largest single-feature area change is %+.3f %% — above the "
                  "%.3f %% tolerance; inspect that feature before proceeding",
                  pct[worst], area_tolerance_pct)
    if still_invalid:
        log.error("%d geometry(ies) remain invalid after make_valid — these are "
                  "usually empty or degenerate rings", still_invalid)
    return out


def compare_repair_methods(gdf: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
    """
    Show, per feature, what buffer(0) would do that make_valid does not.

    Useful once, on a new data source, to decide whether the difference
    matters for that source.
    """
    invalid = gdf[~gdf.geometry.is_valid]
    rows = []
    for idx, geom in invalid.geometry.items():
        mv = make_valid(geom)
        b0 = geom.buffer(0)
        rows.append({
            "index": idx,
            "original_area": geom.area,
            "make_valid_area": mv.area,
            "buffer0_area": b0.area,
            "area_lost_by_buffer0": mv.area - b0.area,
            "make_valid_parts": len(getattr(mv, "geoms", [mv])),
        })
    df = gpd.pd.DataFrame(rows)
    if len(df):
        lost = df["area_lost_by_buffer0"]
        log.info("buffer(0) would lose area on %d of %d invalid features, "
                 "%.1f m² in total, worst %.1f m²",
                 int((lost > 1e-6).sum()), len(df), float(lost.sum()),
                 float(lost.max()))
    return df


# --- Example usage ---
# catchments = gpd.read_file("catchments.gpkg")
# print(compare_repair_methods(catchments))
# fixed = repair_catchments(catchments)