"""세류망 상·하류 분리와 1차 배수유역 산정.

도엽 하천중심선은 피처가 잘게 쪼개져 있고 지류가 본류 중간에 T자로 붙는다. 피처 단위로
자르면 상류망이 통째로 빠지므로, 노딩 → 도로 절단 → 끝점 그래프 확산으로 **이어진 망
전체**를 잡는다(2026-07-31 사용자 지시).

여기서 정해진 1차 배수유역의 bbox가 곧 격자 해석 범위가 된다.
표고 해석·격자 생성은 `B04_PreProcess_Engine_Watershed_Grid.py`가 맡는다.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any

import numpy as np
from scipy.interpolate import LinearNDInterpolator
from scipy.spatial import cKDTree
from shapely.geometry import LineString, MultiPolygon, Polygon, shape
from shapely.ops import substring, unary_union

from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import (
    ContourCloud,
    GridSpec,
    build_cell_mask,
    grid_spec_from_bounds,
    iter_linestrings,
)
from config.config_system import DRAINAGE_GRID_SIZE_M

logger = logging.getLogger(__name__)


# ── 세류망 상·하류 분리 ─────────────────────────────────────────────────────


@dataclass
class StreamSplit:
    """세류망을 도로 교차점에서 상·하류로 가른 결과. 검증 화면이 그대로 그린다.

    `upstream`은 **물이 흐르는 방향(상류 → 하류)으로 정렬**돼 있다. 마지막 좌표가 도로에
    가까운 끝이다. 격자 흐름에 세류 방향을 새겨 넣을 때 이 순서를 그대로 쓴다.
    """

    upstream: list[LineString] = field(default_factory=list)  # 채택 — 1차 영역의 기준
    downstream: list[LineString] = field(default_factory=list)  # 도로 아래로 이어진 망
    no_contact: int = 0  # 어느 쪽에도 이어지지 않아 제외한 조각 수


def split_streams_at_road(
    route_line: LineString,
    stream_features: list[dict[str, Any]],
    cloud: ContourCloud,
) -> StreamSplit:
    """세류망을 도로에서 끊고, 교차점 상류측으로 **이어진 망 전체**를 채택한다.

    도엽 하천중심선은 피처가 잘게 쪼개져 있고 지류가 본류 중간에 T자로 붙는다. 그래서
    피처 단위로 보면 상류망이 통째로 빠진다. 순서를 이렇게 잡는다:

      ① 세류선끼리 `unary_union`으로 노딩 — 중간에서 만나는 지류도 연결로 인식된다
      ② 도로 교차점에서 한 번 더 잘라 상·하류 조각을 물리적으로 분리한다
      ③ 끝점 그래프를 만들고, **도로 교차 노드는 통과하지 못하게** 막는다
      ④ 도로에 접한 조각을 교차점 표고와 비교해 상·하류 씨앗으로 정한다
      ⑤ 씨앗에서 퍼뜨려 이어진 망 전체를 채택 — 도로를 넘어가지 못하므로 상·하류가 섞이지 않는다
    """
    lines: list[LineString] = []
    for feature in stream_features:
        geometry = feature.get("geometry")
        if not geometry:
            continue
        try:
            parsed = shape(geometry)
        except Exception:  # noqa: BLE001
            continue
        lines.extend(line for line in iter_linestrings(parsed) if line.length > 0)
    if not lines:
        return StreamSplit()

    pieces, crossing_nodes = _cut_network_at_road(lines, route_line)
    if not pieces:
        return StreamSplit()

    node_edges: dict[tuple[float, float], list[int]] = {}
    ends: list[tuple[tuple[float, float], tuple[float, float]]] = []
    for index, piece in enumerate(pieces):
        head = _node_key(*piece.coords[0])
        tail = _node_key(*piece.coords[-1])
        ends.append((head, tail))
        node_edges.setdefault(head, []).append(index)
        node_edges.setdefault(tail, []).append(index)

    sampler = ElevationSampler(cloud)
    # 씨앗 조각 → 그 조각의 하류쪽 끝점(= 도로 교차 노드). 이 값이 물 흐름 방향의 기준이 된다.
    upper_seeds: dict[int, tuple[float, float]] = {}
    lower_seeds: dict[int, tuple[float, float]] = {}
    for index, piece in enumerate(pieces):
        touching = [node for node in ends[index] if node in crossing_nodes]
        if not touching:
            continue
        heights = sampler.at(np.array(touching, dtype=np.float64))
        crossing_node = touching[int(np.argmin(heights))]
        if _mean_elevation(piece, sampler) > float(np.min(heights)):
            upper_seeds[index] = crossing_node
        else:
            lower_seeds[index] = crossing_node

    upstream_flow = _spread_network(upper_seeds, ends, node_edges, crossing_nodes)
    downstream_flow = _spread_network(lower_seeds, ends, node_edges, crossing_nodes)
    upstream = set(upstream_flow)
    downstream = set(downstream_flow) - upstream
    logger.info(
        "배수유역: 세류 조각 %d개 → 상류망 %d개 채택 / 하류망 %d개 · 미연결 %d개 제외",
        len(pieces),
        len(upstream),
        len(downstream),
        len(pieces) - len(upstream) - len(downstream),
    )
    return StreamSplit(
        # 상류망은 물 흐름 방향(상류 → 하류)으로 뒤집어 둔다 — 격자 흐름 새김에 그대로 쓴다.
        upstream=[
            _oriented(pieces[index], upstream_flow[index], ends[index])
            for index in sorted(upstream)
        ],
        downstream=[pieces[index] for index in sorted(downstream)],
        no_contact=len(pieces) - len(upstream) - len(downstream),
    )


def _oriented(
    piece: LineString,
    downstream_node: tuple[float, float],
    piece_ends: tuple[tuple[float, float], tuple[float, float]],
) -> LineString:
    """조각을 하류쪽 끝이 마지막 좌표가 되도록 정렬한다."""
    head, _tail = piece_ends
    return LineString(list(piece.coords)[::-1]) if head == downstream_node else piece


def _cut_network_at_road(
    lines: list[LineString], route_line: LineString
) -> tuple[list[LineString], set[tuple[float, float]]]:
    """세류망을 노딩한 뒤 도로 교차점에서 자르고, 그 교차 노드를 함께 돌려준다."""
    noded = unary_union(lines)
    pieces: list[LineString] = []
    crossing_nodes: set[tuple[float, float]] = set()
    for piece in iter_linestrings(noded):
        if not piece.intersects(route_line):
            pieces.append(piece)
            continue
        hits = _intersection_points(piece.intersection(route_line))
        positions = sorted(
            {
                position
                for position in (piece.project(point) for point in hits)
                if 0.0 < position < piece.length
            }
        )
        for point in hits:
            crossing_nodes.add(_node_key(point.x, point.y))
        if not positions:
            # 끝점이 도로에 닿은 경우 — 자를 필요는 없고 그 끝점이 곧 교차 노드다.
            pieces.append(piece)
            continue
        bounds = [0.0, *positions, piece.length]
        for start, end in zip(bounds, bounds[1:]):
            if end - start <= 0:
                continue
            cut = _substring(piece, start, end)
            if cut is not None:
                pieces.append(cut)
    return pieces, crossing_nodes


def _spread_network(
    seeds: dict[int, tuple[float, float]],
    ends: list[tuple[tuple[float, float], tuple[float, float]]],
    node_edges: dict[tuple[float, float], list[int]],
    blocked: set[tuple[float, float]],
) -> dict[int, tuple[float, float]]:
    """씨앗 조각에서 끝점을 타고 퍼진다. 도로 교차 노드는 통과하지 않는다.

    조각마다 **어느 끝점을 통해 도달했는지**를 함께 기록한다. 그 끝점이 도로에 더 가까운
    쪽이므로 곧 그 조각의 하류 방향이다 — 세류망 전체의 물 흐름 방향이 이 한 번의 확산으로
    같이 정해진다.
    """
    downstream = dict(seeds)
    queue = list(seeds)
    while queue:
        index = queue.pop()
        for node in ends[index]:
            if node in blocked:
                continue
            for neighbour in node_edges.get(node, ()):
                if neighbour in downstream:
                    continue
                downstream[neighbour] = node
                queue.append(neighbour)
    return downstream


def _node_key(x: float, y: float) -> tuple[float, float]:
    """끝점 일치 판정용 좌표 키. 노딩 후에도 부동소수 오차가 남아 mm로 반올림한다."""
    return (round(float(x), 3), round(float(y), 3))


class ElevationSampler:
    """등고선 구름에서 임의 지점 표고를 읽는다 — 상·하류 판정 전용.

    최근접 등고선 정점만 쓰면 오차가 등고선 간격(주곡선 5m)만큼 나서, 계곡 교차점 표고가
    실제보다 한 등고선 위로 잡히고 상류 조각이 통째로 하류로 오판된다. TIN 선형보간을
    1차로 쓰고, TIN 밖(볼록껍질 외부)만 최근접 정점으로 메운다.
    """

    def __init__(self, cloud: ContourCloud) -> None:
        self._z = cloud.z
        if cloud.is_empty:
            self._interpolator = None
            self._tree = None
            return
        self._interpolator = LinearNDInterpolator(cloud.xy, cloud.z)
        self._tree = cKDTree(cloud.xy)

    def at(self, xy: np.ndarray) -> np.ndarray:
        """(N, 2) 좌표의 표고 (N,)."""
        if self._interpolator is None or self._tree is None:
            return np.zeros(xy.shape[0])
        values = np.asarray(self._interpolator(xy), dtype=np.float64)
        missing = ~np.isfinite(values)
        if missing.any():
            _, indices = self._tree.query(xy[missing])
            values[missing] = self._z[indices]
        return values


def _intersection_points(geometry: Any) -> list[Any]:
    if geometry.is_empty:
        return []
    if geometry.geom_type == "Point":
        return [geometry]
    if geometry.geom_type in {"MultiPoint", "GeometryCollection", "MultiLineString"}:
        points: list[Any] = []
        for part in geometry.geoms:
            points.extend(_intersection_points(part))
        return points
    if geometry.geom_type == "LineString":
        return [geometry.interpolate(0.5, normalized=True)]
    return []


def _substring(line: LineString, start: float, end: float) -> LineString | None:
    """선형 위 [start, end] 구간을 잘라낸다."""
    piece = substring(line, start, end)
    if piece.is_empty or piece.geom_type != "LineString" or piece.length <= 0:
        return None
    return piece


def _mean_elevation(line: LineString, sampler: ElevationSampler) -> float:
    """선을 10m 간격으로 훑은 평균 표고."""
    samples = max(2, int(line.length // 10.0) + 1)
    positions = np.linspace(0.0, line.length, samples)
    points = np.array([list(line.interpolate(position).coords)[0] for position in positions])
    return float(np.mean(sampler.at(points)))


# ── 1차 배수유역 ────────────────────────────────────────────────────────────


@dataclass
class PrimaryRegion:
    """1차 배수유역과 그 안에 생성된 격자. 검증 화면이 이 내용을 그대로 그린다."""

    split: StreamSplit
    # 상류 세류망 + 노선을 반경 버퍼해 합친 영역.
    area: Polygon | MultiPolygon | None
    spec: GridSpec
    radius_m: float
    # 1차 영역에 조금이라도 걸쳐 실제로 생성된 셀 (rows, cols) bool 마스크.
    cell_mask: np.ndarray | None = None
    # 1차 영역 밖으로 나간 노선 길이(m). 그 구간 사면은 해석에서 빠진다는 경고 지표.
    road_outside_m: float = 0.0

    @property
    def active_cells(self) -> int:
        return 0 if self.cell_mask is None else int(self.cell_mask.sum())


def build_primary_region(
    route_line: LineString,
    stream_features: list[dict[str, Any]],
    cloud: ContourCloud,
    radius_m: float,
    cell_m: float = DRAINAGE_GRID_SIZE_M,
) -> PrimaryRegion:
    """**상류 세류망 + 계획 노선**을 반경 버퍼한 범위 = 1차 배수유역, 그 안에 격자를 생성한다.

    노선 버퍼는 세류 교차가 없는 구간의 도로도 격자 안에 들어오게 한다 — 그래야 그 구간
    사면이 유역으로 잡힌다. 상류 세류망 선정이 정확해진 뒤 다시 포함했다(2026-07-31 사용자 지시).

    격자는 bbox를 통째로 채우지 않는다. **도로 시작점에 셀 모서리를 맞춘 뒤, 1차 영역에
    조금이라도 걸치는 셀만** 생성한다(2026-07-31 사용자 지시). bbox 전체를 쓰면 영역 밖
    빈 셀이 대부분이라 의미가 없고, 원점을 bbox 좌상단에 두면 영역이 조금만 변해도 격자가
    통째로 밀려 이전 결과와 셀이 대응되지 않는다.

    노선이 이 영역 밖으로 나가는 길이는 따로 재서 남긴다 — 그 구간은 도로 셀이 격자에
    없어 유역이 잡히지 않으므로 반경을 올릴지 판단하는 근거가 된다.
    """
    split = split_streams_at_road(route_line, stream_features, cloud)
    geometries = [route_line.buffer(radius_m)]
    geometries.extend(line.buffer(radius_m) for line in split.upstream)
    if not split.upstream:
        logger.warning("배수유역: 상류 세류망이 없어 노선 버퍼만으로 1차 영역을 잡습니다.")
    area = unary_union(geometries)
    x_min, y_min, x_max, y_max = area.bounds
    road_start = route_line.coords[0]
    spec = grid_spec_from_bounds(
        x_min, y_min, x_max, y_max, cell_m, anchor_xy=(float(road_start[0]), float(road_start[1]))
    )
    cell_mask = build_cell_mask(spec, area)
    outside = route_line.difference(area)
    road_outside_m = float(outside.length) if not outside.is_empty else 0.0
    active = int(cell_mask.sum())
    logger.info(
        "배수유역: 1차 영역 %.0f㎡ → 격자 %d×%d (%.2fm, 도로 시점 기준) 중 %d셀 생성 "
        "(bbox %d셀의 %.0f%%), 노선 이탈 %.0fm/%.0fm",
        area.area,
        spec.n_rows,
        spec.n_cols,
        spec.cell_m,
        active,
        spec.size,
        100.0 * active / max(spec.size, 1),
        road_outside_m,
        route_line.length,
    )
    return PrimaryRegion(
        split=split,
        area=area,
        spec=spec,
        radius_m=radius_m,
        cell_mask=cell_mask,
        road_outside_m=road_outside_m,
    )