Aislo/B05_wf2_Route/B05_wf2_Route_Engine_Watershed_Flow.py

"""배수유역 흐름 해석 — 도로 굽기 · 상류 추적(포인터 더블링) · 유역 폴리곤화.

핵심은 하나다: **물길을 따라가 도로에 닿는 셀만 유역이다.**
셀마다 D8 수신 셀을 따라가 종착점(root)을 구하고, 그 종착점이 도로 셀이면 활성이다.

이 방식은 능선을 따로 찾지 않는다. 능선 너머 셀의 물은 다른 계곡으로 빠져 도로에
닿지 못하므로 자동으로 비활성이 되고, 그 경계선이 곧 능선이다. 유역 안쪽 봉우리는
물이 결국 도로로 흘러 자동으로 포함된다.

종착점은 세부유역 라벨의 근거로도 그대로 쓴다 — 셀이 도달한 도로 셀이 정해지면
그 도로 셀을 담당하는 관이 곧 그 셀의 유역 번호다. 관 배치가 바뀌어도 격자 해석을
다시 돌릴 필요 없이 "도로 셀 → 관" 대응만 다시 계산하면 된다.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
from rasterio.features import rasterize, shapes
from rasterio.transform import from_origin
from scipy.spatial import cKDTree
from shapely.geometry import LineString, MultiPolygon, Polygon, shape
from shapely.ops import unary_union

from B05_wf2_Route.B05_wf2_Route_Engine_Watershed_Grid import GridSpec, TerrainGrid
from config.config_system import (
    DRAINAGE_MIN_BASIN_AREA_M2,
    DRAINAGE_POLYGON_SIMPLIFY_M,
    DRAINAGE_ROAD_WIDTH_M,
)

logger = logging.getLogger(__name__)

# 포인터 더블링 반복 상한. 한 번에 경로 길이가 2배가 되므로 2^40 스텝이면 어떤 격자도 덮는다.
_MAX_DOUBLING_ROUNDS = 40


@dataclass
class RoadRaster:
    """격자에 구운 도로. 도로 셀은 흐름을 흡수하는 종착점이 된다."""

    mask: np.ndarray  # (R, C) bool
    cell_index: np.ndarray  # (K,) int32 — 도로 셀의 평탄 인덱스
    chainage: np.ndarray  # (K,) float64 — 도로 셀의 누가거리(m)
    slot_of_cell: np.ndarray  # (R*C,) int32 — 도로 셀이면 K 내 위치, 아니면 −1

    @property
    def count(self) -> int:
        return int(self.cell_index.size)


@dataclass
class FlowResult:
    """상류 추적 결과."""

    root: np.ndarray  # (R*C,) int32 — 흐름 종착 셀의 평탄 인덱스
    road_slot: np.ndarray  # (R*C,) int32 — 도달한 도로 셀 슬롯, 도달 못하면 −1
    active: np.ndarray  # (R, C) bool — 도로에 물이 닿는 셀
    path_length: np.ndarray  # (R*C,) float32 — 종착점까지 물길 길이(m)
    strength: np.ndarray  # (K,) int64 — 도로 셀별 상류 셀 수(흐름 강도)


def grid_transform(spec: GridSpec):
    """rasterio 아핀 변환. 행 0이 북쪽(y_max)이다."""
    return from_origin(spec.x_min, spec.y_max, spec.cell_m, spec.cell_m)


# ── 도로 굽기 ───────────────────────────────────────────────────────────────


def rasterize_road(
    spec: GridSpec,
    route_line: LineString,
    width_m: float = DRAINAGE_ROAD_WIDTH_M,
) -> RoadRaster:
    """노선을 노폭만큼 두껍게 격자에 굽고, 각 도로 셀에 누가거리를 붙인다.

    폭을 주는 이유는 실제 노면이 물을 받기 때문이기도 하지만, 1셀 선으로 구우면 D8
    대각 이동이 도로를 건너뛰어 상류 물이 도로를 지나쳐 버리기 때문이다. 3셀 이상 두께면
    내리막 물길이 반드시 도로 셀을 한 번은 밟는다.
    """
    half_width = max(width_m / 2.0, spec.cell_m)
    burned = rasterize(
        [(route_line.buffer(half_width), 1)],
        out_shape=(spec.n_rows, spec.n_cols),
        transform=grid_transform(spec),
        fill=0,
        dtype="uint8",
        all_touched=True,
    ).astype(bool)

    slot_of_cell = np.full(spec.size, -1, dtype=np.int32)
    cell_index = np.flatnonzero(burned.reshape(-1)).astype(np.int32)
    if cell_index.size == 0:
        logger.warning("배수유역: 노선이 격자 범위 밖입니다 — 도로 셀 0개.")
        return RoadRaster(burned, cell_index, np.zeros(0), slot_of_cell)

    # 도로 셀 누가거리는 노선을 촘촘히 샘플링해 가장 가까운 샘플의 누가거리로 준다.
    step = max(spec.cell_m / 2.0, 0.25)
    positions = np.arange(0.0, route_line.length + step, step)
    samples = np.array([list(route_line.interpolate(p).coords)[0] for p in positions])
    rows = (cell_index // spec.n_cols).astype(np.float64)
    cols = (cell_index % spec.n_cols).astype(np.float64)
    centers = np.column_stack(
        (
            spec.x_min + (cols + 0.5) * spec.cell_m,
            spec.y_max - (rows + 0.5) * spec.cell_m,
        )
    )
    _, nearest = cKDTree(samples).query(centers)
    chainage = np.minimum(positions[nearest], route_line.length)
    slot_of_cell[cell_index] = np.arange(cell_index.size, dtype=np.int32)
    return RoadRaster(burned, cell_index, chainage, slot_of_cell)


# ── 상류 추적 ───────────────────────────────────────────────────────────────


def trace_flow(terrain: TerrainGrid, road: RoadRaster) -> FlowResult:
    """모든 셀의 물길 종착점을 구하고 도로 도달 여부(=유역 포함 여부)를 판정한다.

    포인터 더블링으로 한 번에 경로 길이를 2배씩 늘려 종착점을 찾는다. 채움·평탄해소를
    거친 표고에서는 흐름을 따라 표고가 단조 감소하므로 순환이 없고, 반복은 항상 끝난다.
    """
    spec = terrain.spec
    receiver = terrain.receiver.astype(np.int32, copy=True)
    step_length = terrain.step_length.astype(np.float32, copy=True)

    # 도로 셀은 흐름을 흡수한다 — 물이 도로에 닿으면 거기서 끝난다.
    receiver[road.cell_index] = road.cell_index
    step_length[road.cell_index] = 0.0

    jump = receiver
    path_length = step_length
    for _ in range(_MAX_DOUBLING_ROUNDS):
        next_jump = jump[jump]
        if np.array_equal(next_jump, jump):
            break
        path_length = path_length + path_length[jump]
        jump = next_jump

    road_slot = road.slot_of_cell[jump]
    active_flat = road_slot >= 0
    strength = (
        np.bincount(road_slot[active_flat], minlength=max(road.count, 1)).astype(np.int64)
        if road.count
        else np.zeros(0, dtype=np.int64)
    )
    logger.info(
        "배수유역: 활성 셀 %d / %d (도로 셀 %d)", int(active_flat.sum()), spec.size, road.count
    )
    return FlowResult(
        root=jump,
        road_slot=road_slot,
        active=active_flat.reshape(spec.n_rows, spec.n_cols),
        path_length=path_length,
        strength=strength,
    )


def border_contact(active: np.ndarray) -> dict[str, bool]:
    """활성 셀이 격자 최외곽에 닿은 방향. 전부 False면 유역이 능선 안에서 닫힌 것이다."""
    return {
        "north": bool(active[0, :].any()),
        "south": bool(active[-1, :].any()),
        "west": bool(active[:, 0].any()),
        "east": bool(active[:, -1].any()),
    }


# ── 폴리곤화 ────────────────────────────────────────────────────────────────


def polygonize_labels(
    spec: GridSpec,
    labels: np.ndarray,
    min_area_m2: float = DRAINAGE_MIN_BASIN_AREA_M2,
) -> dict[int, Polygon | MultiPolygon]:
    """라벨 격자를 라벨별 폴리곤으로 바꾼다. 음수 라벨은 배경으로 무시한다."""
    label_grid = np.ascontiguousarray(labels.reshape(spec.n_rows, spec.n_cols), dtype=np.int32)
    valid_mask = label_grid >= 0
    if not valid_mask.any():
        return {}
    collected: dict[int, list[Polygon]] = {}
    for geometry, value in shapes(
        label_grid, mask=valid_mask, transform=grid_transform(spec), connectivity=4
    ):
        polygon = shape(geometry)
        if polygon.is_empty or polygon.area < min_area_m2:
            continue
        collected.setdefault(int(value), []).append(polygon)

    merged: dict[int, Polygon | MultiPolygon] = {}
    for label, parts in collected.items():
        union = unary_union(parts)
        if union.is_empty:
            continue
        simplified = union.simplify(DRAINAGE_POLYGON_SIMPLIFY_M, preserve_topology=True)
        merged[label] = simplified if not simplified.is_empty else union
    return merged


def largest_ring(geometry: Polygon | MultiPolygon) -> list[tuple[float, float]]:
    """폴리곤(또는 멀티폴리곤)에서 가장 큰 조각의 외곽 링 좌표를 뽑는다."""
    if geometry.is_empty:
        return []
    if geometry.geom_type == "MultiPolygon":
        geometry = max(geometry.geoms, key=lambda part: part.area)
    return [(float(x), float(y)) for x, y in geometry.exterior.coords]


def outer_boundary(
    spec: GridSpec, active: np.ndarray, min_area_m2: float = DRAINAGE_MIN_BASIN_AREA_M2
) -> Polygon | MultiPolygon | None:
    """활성 셀 전체의 외곽 = 2차 전체 배수유역 경계.

    비활성 셀이 격자 최외곽에 띠로 완성되면 활성 영역이 그 안에 닫힌다. 그 닫힌 영역의
    바깥선이 곧 분수령이므로 능선을 따로 그릴 필요가 없다.
    """
    labels = np.where(active.reshape(-1), 0, -1).astype(np.int32)
    polygons = polygonize_labels(spec, labels, min_area_m2)
    return polygons.get(0)