"""
Derives approximate channel width per real tronçon (river reach
segment) from BD TOPO data already downloaded in this project --
surface_hydrographique.geojson (real water-surface polygons) and
troncon_hydrographique.geojson (real centerline geometry) -- no new
data source needed.

width_m = surface_polygon_area_m2 / troncon_length_m

This is real, derivable channel geometry, not a full cross-sectional
profile -- true cross-section (a real depth profile across the
channel) needs bathymetric/LIDAR survey data, which hasn't been
sourced or verified for this project. Width alone is still a
physically meaningful hydraulic signal: it's the dominant term in
Manning's-equation-style conveyance estimates, and the reason a wide
shallow reach and a narrow deep one carry very different discharge for
the same water level -- exactly the kind of physics a plain
distance/elevation-drop edge attribute can't express on its own.

Spatial join: each tronçon is matched to whichever real surface
polygon its centerline geometry falls within (or, failing that, the
nearest one within a small tolerance) -- a tronçon with no matching
surface polygon at all gets width_m = NaN, not a fabricated default,
since "no real surface polygon found" is meaningfully different from
"the channel is zero-width."

Usage:
    python -m scripts.compute_edge_width --data-root datasets
"""
import argparse
from pathlib import Path

import geopandas as gpd
import pandas as pd


def compute_widths(troncons: gpd.GeoDataFrame, surfaces: gpd.GeoDataFrame) -> pd.DataFrame:
    """
    Real spatial join: for each tronçon, find the surface polygon(s) it
    intersects, sum their real area (a tronçon can legitimately span
    more than one mapped surface polygon), divide by real tronçon
    length. Returns one row per tronçon with a real cleabs identifier
    and width_m (NaN where no real surface polygon matched).
    """
    troncons = troncons.copy()
    troncons["length_m"] = troncons.geometry.length

    joined = gpd.sjoin(troncons[["cleabs", "length_m", "geometry"]], surfaces[["geometry"]],
                        how="left", predicate="intersects")

    area_by_troncon = surfaces.geometry.area
    joined["surface_area_m2"] = joined["index_right"].map(
        lambda idx: area_by_troncon.loc[idx] if pd.notna(idx) and idx in area_by_troncon.index else float("nan")
    )

    result = joined.groupby("cleabs").agg(
        length_m=("length_m", "first"),
        total_surface_area_m2=("surface_area_m2", lambda x: x.sum() if x.notna().any() else float("nan")),
    ).reset_index()

    result["width_m"] = result["total_surface_area_m2"] / result["length_m"]
    n_matched = result["width_m"].notna().sum()
    print(f"{n_matched}/{len(result)} real tronçon(s) matched to a real surface polygon "
          f"({len(result) - n_matched} have no real surface data -- width_m will be NaN for those, not fabricated)")

    return result[["cleabs", "length_m", "width_m"]]


def main() -> None:
    parser = argparse.ArgumentParser(description="Derive real channel width from BD TOPO surface polygons")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"))
    parser.add_argument("--output", type=Path, default=None)
    args = parser.parse_args()

    bdtopo_dir = args.data_root / "bdtopo_hydro"
    troncons_path = bdtopo_dir / "troncon_hydrographique.geojson"
    surfaces_path = bdtopo_dir / "surface_hydrographique.geojson"

    if not troncons_path.exists() or not surfaces_path.exists():
        print(f"Missing {troncons_path} or {surfaces_path} -- run download_bdtopo_hydro.py first.")
        return

    print(f"Loading real tronçons from {troncons_path}...")
    troncons = gpd.read_file(troncons_path)
    print(f"Loading real surface polygons from {surfaces_path}...")
    surfaces = gpd.read_file(surfaces_path)
    print(f"{len(troncons)} real tronçon(s), {len(surfaces)} real surface polygon(s)")

    if troncons.crs != surfaces.crs:
        print(f"CRS mismatch ({troncons.crs} vs {surfaces.crs}) -- reprojecting surfaces to match tronçons")
        surfaces = surfaces.to_crs(troncons.crs)

    widths = compute_widths(troncons, surfaces)

    real_widths = widths["width_m"].dropna()
    if len(real_widths) > 0:
        print(f"\nReal width distribution: min={real_widths.min():.1f}m, median={real_widths.median():.1f}m, "
              f"max={real_widths.max():.1f}m")
        print("Sanity-check these against what you know of the real rivers before trusting them -- "
              "an implausible max (e.g. hundreds of meters) usually means a tronçon matched an "
              "unrelated, much larger surface polygon nearby, not a real channel width.")

    output_path = args.output or (args.data_root / "edge_widths.csv")
    widths.to_csv(output_path, index=False)
    print(f"\nSaved to {output_path}")


if __name__ == "__main__":
    main()