Aislo/B04_PreProcess/B04_PreProcess_Engine_Watershed_Descent.py

"""등고선 기반 흐름 방향 — 높은 등고 라인에서 낮은 등고 라인으로 내려가며 방향을 세운다.

기존 방식(등고선 → TIN 보간 → 지표면 기울기 → D8)은 보간면이 만든 가짜 웅덩이와 평탄
삼각형 때문에 흐름이 중간에서 끊겼다. 실데이터에서 채움·평탄해소를 거치고도 싱크가 수천 개
남았고, 그 싱크 하나가 상류 유역 전체를 통째로 삼켰다.

여기서는 **보간면을 거치지 않는다.** 등고선을 격자에 직접 굽고, 셀마다 "내가 속한 등고
라인보다 한 단 낮은 등고 라인"이 어디인지를 찾아 그쪽으로 방향을 준다(2026-07-31 사용자 지시).

  ① 등고선을 격자에 굽는다 — 셀이 어느 표고의 라인 위인지 기록
  ② 셀마다 가장 가까운 등고 라인을 찾아 그 표고를 '밴드'로 삼는다
  ③ 표고가 높은 밴드부터 내려오며, 각 밴드에서 **한 단 낮은 등고 라인까지의 거리**를 잰다
  ④ 위치에너지 = 밴드 순위 × 큰 수 + 그 거리
  ⑤ 수신 셀 = 위치에너지가 더 낮은 8이웃 중 화살표 방향에 가장 가까운 셀

④의 위치에너지는 흐름을 따라 반드시 감소한다. 그래서 **순환도 웅덩이도 원리적으로 생기지
않는다** — 채움이나 평탄면 해소가 아예 필요 없다.

⑤ 덕분에 화면 화살표(32방위)와 실제 추적 경로가 항상 같은 방향을 가리킨다. 예전에는
화살표는 기울기, 추적은 D8이라 서로 어긋나 눈으로 검증할 수가 없었다.
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from typing import Any

import numpy as np
from rasterio.features import rasterize
from scipy.ndimage import distance_transform_edt
from shapely.geometry import shape

from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import (
    AZIMUTH_INVALID,
    AZIMUTH_SINK,
    AZIMUTH_STEPS,
    GridSpec,
    _feature_elevation,
    grid_transform,
    iter_linestrings,
)
from config.config_system import DRAINAGE_CONTOUR_MIN_LENGTH_M

logger = logging.getLogger(__name__)

# 8이웃 (행 증분, 열 증분).
_NEIGHBOURS = (
    (-1, -1),
    (-1, 0),
    (-1, 1),
    (0, -1),
    (0, 1),
    (1, -1),
    (1, 0),
    (1, 1),
)


@dataclass
class ContourDescent:
    """등고선에서 직접 세운 흐름 방향 격자."""

    spec: GridSpec
    band_elevation: np.ndarray  # (R, C) float32 — 셀이 속한 등고 라인 표고, 무효는 NaN
    valid: np.ndarray  # (R, C) bool — 방향을 세운 셀
    receiver: np.ndarray  # (R*C,) int32 — 다음 셀, 최하단 밴드는 자기 자신
    step_length: np.ndarray  # (R*C,) float32
    azimuth: np.ndarray  # (R*C,) int16 — 32방위 코드(32=제자리, 33=무효)
    levels: list[float]  # 사용된 등고 표고(내림차순)


def rasterize_contours(
    spec: GridSpec,
    contour_features: list[dict[str, Any]],
    elevation_floor_m: float | None = None,
) -> tuple[np.ndarray, list[float]]:
    """등고선을 격자에 굽는다. 셀마다 그 위를 지나는 등고 라인의 표고(없으면 NaN)."""
    by_level: dict[float, list[Any]] = {}
    for feature in contour_features:
        geometry = feature.get("geometry")
        if not geometry:
            continue
        elevation = _feature_elevation(feature.get("properties") or {})
        if elevation is None:
            continue
        if elevation_floor_m is not None and elevation < elevation_floor_m:
            continue
        try:
            parsed = shape(geometry)
        except Exception:  # noqa: BLE001
            continue
        for line in iter_linestrings(parsed):
            if line.length < DRAINAGE_CONTOUR_MIN_LENGTH_M:
                continue
            by_level.setdefault(float(elevation), []).append(line)

    burned = np.full((spec.n_rows, spec.n_cols), np.nan, dtype=np.float32)
    levels = sorted(by_level, reverse=True)
    transform = grid_transform(spec)
    for elevation in levels:
        stamp = rasterize(
            [(line, 1) for line in by_level[elevation]],
            out_shape=(spec.n_rows, spec.n_cols),
            transform=transform,
            fill=0,
            dtype="uint8",
            all_touched=True,
        ).astype(bool)
        # 낮은 표고부터 덮어써야 겹치는 셀이 낮은 라인으로 남는다 — 물은 낮은 쪽으로 간다.
        burned[stamp] = elevation
    logger.info(
        "배수유역: 등고 라인 %d단(%.0f~%.0fm)을 격자에 굽어 %d셀",
        len(levels),
        levels[-1] if levels else 0.0,
        levels[0] if levels else 0.0,
        int(np.isfinite(burned).sum()),
    )
    return burned, levels


def build_contour_descent(
    spec: GridSpec,
    contour_features: list[dict[str, Any]],
    domain: np.ndarray | None = None,
    elevation_floor_m: float | None = None,
) -> ContourDescent:
    """등고선만으로 셀별 흐름 방향을 세운다. 보간면을 만들지 않는다."""
    rows, cols = spec.n_rows, spec.n_cols
    burned, levels = rasterize_contours(spec, contour_features, elevation_floor_m)
    empty = ContourDescent(
        spec=spec,
        band_elevation=np.full((rows, cols), np.nan, dtype=np.float32),
        valid=np.zeros((rows, cols), dtype=bool),
        receiver=np.arange(spec.size, dtype=np.int32),
        step_length=np.zeros(spec.size, dtype=np.float32),
        azimuth=np.full(spec.size, AZIMUTH_INVALID, dtype=np.int16),
        levels=levels,
    )
    if len(levels) < 2:
        logger.warning("배수유역: 등고 라인이 2단 미만이라 방향을 세울 수 없습니다.")
        return empty

    on_contour = np.isfinite(burned)
    # ② 셀마다 가장 가까운 등고 라인의 표고 = 그 셀의 밴드.
    _, (near_row, near_col) = distance_transform_edt(~on_contour, return_indices=True)
    band_elevation = burned[near_row, near_col].astype(np.float32)
    inside = domain if domain is not None else np.ones((rows, cols), dtype=bool)
    band_elevation = np.where(inside, band_elevation, np.nan)

    # ③④ 높은 밴드부터 내려오며 한 단 낮은 등고 라인까지의 거리와 목표 셀을 구한다.
    distance = np.full((rows, cols), np.inf, dtype=np.float32)
    target_row = np.zeros((rows, cols), dtype=np.int32)
    target_col = np.zeros((rows, cols), dtype=np.int32)
    band_rank = np.full((rows, cols), -1, dtype=np.int32)
    for rank, elevation in enumerate(levels[:-1]):
        members = inside & (band_elevation == elevation)
        if not members.any():
            continue
        lower = on_contour & (burned < elevation)
        if not lower.any():
            continue
        step_distance, (step_row, step_col) = distance_transform_edt(~lower, return_indices=True)
        distance[members] = step_distance[members].astype(np.float32)
        target_row[members] = step_row[members]
        target_col[members] = step_col[members]
        band_rank[members] = len(levels) - 1 - rank  # 높을수록 큰 값

    # 최하단 밴드는 더 내려갈 등고 라인이 없다. 무효로 빼지 않고 **정지 셀**로 남긴다 —
    # 그래야 그 자리가 도로·세류선이면 적색으로 잡히고, 아니면 물이 고이는 지점으로 보인다.
    lowest = inside & (band_elevation == levels[-1]) & (band_rank < 0)
    if lowest.any():
        distance[lowest] = 0.0
        band_rank[lowest] = 0

    valid = band_rank >= 0
    if not valid.any():
        logger.warning("배수유역: 하강 방향을 세운 셀이 없습니다.")
        return empty

    # 밴드가 하나 낮아지면 위치에너지가 반드시 떨어지도록 거리 최대치보다 큰 간격을 준다.
    finite = distance[valid & np.isfinite(distance)]
    span = (float(finite.max()) if finite.size else 1.0) + 2.0
    distance[valid & ~np.isfinite(distance)] = 0.0
    potential = np.where(valid, band_rank.astype(np.float64) * span + distance, np.inf)

    receiver, step_length, azimuth = _route_by_potential(
        spec, potential, valid, target_row, target_col
    )
    logger.info(
        "배수유역: 등고선 하강 방향 %d셀 (밴드 %d단), 최하단 정지 %d셀",
        int(valid.sum()),
        int(band_rank[valid].max() - band_rank[valid].min() + 1),
        int((azimuth == AZIMUTH_SINK).sum()),
    )
    return ContourDescent(
        spec=spec,
        band_elevation=np.where(valid, band_elevation, np.nan).astype(np.float32),
        valid=valid,
        receiver=receiver,
        step_length=step_length,
        azimuth=azimuth,
        levels=levels,
    )


def _route_by_potential(
    spec: GridSpec,
    potential: np.ndarray,
    valid: np.ndarray,
    target_row: np.ndarray,
    target_col: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """위치에너지가 낮은 8이웃 중 **화살표 방향에 가장 가까운** 셀을 수신 셀로 고른다.

    화살표는 "한 단 낮은 등고 라인 쪽"을 가리키는 연속 방위이고, 수신 셀은 그 방위에 가장
    가까운 이웃이다. 그래서 화면 화살표와 실제 추적 경로가 어긋나지 않는다.
    위치에너지가 더 낮은 이웃만 후보로 두므로 순환이 생기지 않는다.
    """
    rows, cols = spec.n_rows, spec.n_cols
    grid_row, grid_col = np.meshgrid(np.arange(rows), np.arange(cols), indexing="ij")
    # 목표(한 단 낮은 등고 라인 위의 셀)를 향하는 연속 방위.
    aim_row = (target_row - grid_row).astype(np.float64)
    aim_col = (target_col - grid_col).astype(np.float64)
    aim_norm = np.hypot(aim_row, aim_col)
    aim_norm[aim_norm == 0.0] = 1.0
    aim_row /= aim_norm
    aim_col /= aim_norm

    padded = np.full((rows + 2, cols + 2), np.inf)
    padded[1:-1, 1:-1] = potential
    flat_index = np.arange(spec.size, dtype=np.int32).reshape(rows, cols)
    padded_index = np.full((rows + 2, cols + 2), -1, dtype=np.int32)
    padded_index[1:-1, 1:-1] = flat_index

    best_score = np.full((rows, cols), -np.inf)
    receiver = flat_index.copy()
    step = np.zeros((rows, cols), dtype=np.float32)
    centre = potential
    for row_shift, col_shift in _NEIGHBOURS:
        neighbour = padded[
            1 + row_shift : 1 + row_shift + rows, 1 + col_shift : 1 + col_shift + cols
        ]
        length = math.hypot(row_shift, col_shift)
        # 방위 일치도(코사인 유사도)가 클수록 좋은 후보다.
        score = (aim_row * row_shift + aim_col * col_shift) / length
        better = valid & np.isfinite(neighbour) & (neighbour < centre) & (score > best_score)
        if not better.any():
            continue
        best_score = np.where(better, score, best_score)
        neighbour_index = padded_index[
            1 + row_shift : 1 + row_shift + rows, 1 + col_shift : 1 + col_shift + cols
        ]
        receiver = np.where(better, neighbour_index, receiver)
        step = np.where(better, np.float32(length * spec.cell_m), step)

    moved = receiver != flat_index
    delta_row = (receiver // cols - flat_index // cols).astype(np.float64)
    delta_col = (receiver % cols - flat_index % cols).astype(np.float64)
    angle = np.arctan2(delta_row, delta_col)
    code = np.rint(angle / (2.0 * math.pi / AZIMUTH_STEPS)).astype(np.int16) % AZIMUTH_STEPS
    azimuth = np.where(moved, code, AZIMUTH_SINK)
    azimuth = np.where(valid, azimuth, AZIMUTH_INVALID)
    return receiver.reshape(-1), step.reshape(-1), azimuth.reshape(-1).astype(np.int16)