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

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

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

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

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

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

from __future__ import annotations

import hashlib
import json
import logging
import math
import time
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
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]  # 사용된 등고 표고(내림차순)


_LINES_CACHE: list[tuple[Any, int, float | None, dict[float, list[Any]]]] = []


def _lines_by_level(
    contour_features: list[dict[str, Any]], elevation_floor_m: float | None
) -> dict[float, list[Any]]:
    """등고선 피처를 표고별 선 묶음으로 푼다 — **격자와 무관**하므로 한 번만 푼다.

    확장 회차마다 다시 부르는데 피처 4,200개를 매번 `shape()` 로 푸는 비용이 그대로
    붙었다. 같은 목록·같은 하한이면 그대로 돌려준다(목록 객체를 함께 들고 있어 id 가
    다른 목록에 재사용되지 않는다).
    """
    for holder, count, floor, cached in _LINES_CACHE:
        if (
            holder is contour_features
            and count == len(contour_features)
            and floor == elevation_floor_m
        ):
            return cached
    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)
    _LINES_CACHE.append((contour_features, len(contour_features), elevation_floor_m, by_level))
    del _LINES_CACHE[:-2]
    return by_level


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 = _lines_by_level(contour_features, elevation_floor_m)
    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


# 밴드별 하강거리 EDT 를 **쓸 자리 둘레**에서만 돌리기 위한 값들.
#
# 왜(2026-09-07 실측) — 확장 회차마다 방향장을 다시 만드는데, 그 안의 밴드별 EDT 가
# 격자 전체를 단계마다 훑어 확장 95.8s 중 50.4s 를 썼다. 정작 쓰는 값은 그 단계의
# 밴드 셀뿐이다. 다만 `analyze_domain` 이 `domain` 을 안 넘겨 밴드가 **격자 전역에
# 흩어지므로** 바운딩박스 하나로는 안 좁아진다. 그래서 격자를 타일로 나눠 **셀이 있는
# 타일만** 그 둘레 `margin` 까지 잘라 EDT 를 돌린다.
#
# 창 밖에 더 가까운 등고선이 있을 수 있으면(창 안 최대 거리가 여유에 닿거나 창에 낮은
# 라인이 없으면) **그 타일만** 창을 4배로 넓혀 다시 잰다 — 근사가 아니라 같은 값을 싸게
# 구하는 것. 단계 전체를 격자 전체로 되돌리면 이득이 사라진다(그 방식일 때 폴백 25/91).
#
# 고른 근거(745×1035 격자·166단 중 EDT 도는 91단, 전체격자 방식 5.2s):
#   **타일 128 여유 32 → 2.1s(넓힘 15)**  타일 128 여유 24 → 1.9s(넓힘 52)
#   타일 128 여유 48 → 2.5s(넓힘 7)       타일 256 여유 32 → 2.7s(넓힘 15)
#   타일 512 여유 32 → 4.5s(넓힘 15)
# 여유 24 가 0.2s 빠르지만 넓힘이 52회로 지형에 예민해 32 로 뒀다. 한 단 아래 등고선은
# 주곡선 5m·셀 1m 에서 대개 수십 칸 안이다. 넓힘 횟수는 계측 줄에 찍힌다.
DESCENT_WINDOW_MARGIN_CELLS = 32
DESCENT_TILE_CELLS = 128


def _distance_to_lower(
    members: np.ndarray,
    lower: np.ndarray,
    margin: int | None = None,
    tile: int | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
    """`members` 셀에서 한 단 낮은 등고 라인까지의 거리와 목표 셀(**전체 격자 좌표**).

    네 번째 값은 창을 넓혀 다시 잰 타일 수다(0 = 전부 첫 창에서 끝남). 반환 순서는
    `members` 의 행우선 순서 — 호출부가 `distance[members] = ...` 로 그대로 넣는다.
    """
    margin = DESCENT_WINDOW_MARGIN_CELLS if margin is None else margin
    tile = DESCENT_TILE_CELLS if tile is None else tile
    rows, cols = members.shape
    out_distance = np.zeros((rows, cols), dtype=np.float64)
    out_row = np.zeros((rows, cols), dtype=np.int32)
    out_col = np.zeros((rows, cols), dtype=np.int32)
    widened = 0

    for row_start in range(0, rows, tile):
        row_stop = min(rows, row_start + tile)
        for col_start in range(0, cols, tile):
            col_stop = min(cols, col_start + tile)
            tile_members = members[row_start:row_stop, col_start:col_stop]
            if not tile_members.any():
                continue
            # 첫 창에서 안 닿으면 **그 타일만** 창을 넓혀 다시 잰다 — 단계 전체를 격자
            # 전체로 되돌리면 이득이 사라진다(실측 폴백 25/91).
            reach = margin
            attempt = 0
            while True:
                win_row0 = max(0, row_start - reach)
                win_row1 = min(rows, row_stop + reach)
                win_col0 = max(0, col_start - reach)
                win_col1 = min(cols, col_stop + reach)
                whole = win_row0 == 0 and win_col0 == 0 and win_row1 == rows and win_col1 == cols
                window_lower = lower[win_row0:win_row1, win_col0:win_col1]
                if not window_lower.any():
                    if whole:
                        return _full_distance_to_lower(members, lower)
                    reach *= 4
                    attempt += 1
                    widened += 1
                    continue
                step, (step_row, step_col) = distance_transform_edt(
                    ~window_lower, return_indices=True
                )
                in_window = np.zeros(window_lower.shape, dtype=bool)
                in_window[
                    row_start - win_row0 : row_stop - win_row0,
                    col_start - win_col0 : col_stop - win_col0,
                ] = tile_members
                hit_rows, hit_cols = np.nonzero(in_window)
                values = step[hit_rows, hit_cols]
                # 타일 둘레로 `reach` 를 뒀으므로, 거리가 그보다 짧으면 창 밖에 더
                # 가까운 것은 있을 수 없다. 닿으면 창을 넓혀 다시 잰다.
                if not whole and values.size and float(values.max()) >= reach:
                    reach *= 4
                    attempt += 1
                    widened += 1
                    continue
                break
            global_rows = hit_rows + win_row0
            global_cols = hit_cols + win_col0
            out_distance[global_rows, global_cols] = values
            out_row[global_rows, global_cols] = step_row[hit_rows, hit_cols] + win_row0
            out_col[global_rows, global_cols] = step_col[hit_rows, hit_cols] + win_col0

    return out_distance[members], out_row[members], out_col[members], widened


def _full_distance_to_lower(
    members: np.ndarray, lower: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, bool]:
    """격자 전체 EDT — 창으로 못 믿을 때만 쓰는 폴백."""
    step, (step_row, step_col) = distance_transform_edt(~lower, return_indices=True)
    return step[members], step_row[members], step_col[members], True


def build_contour_descent(
    spec: GridSpec,
    contour_features: list[dict[str, Any]],
    domain: np.ndarray | None = None,
    elevation_floor_m: float | None = None,
) -> ContourDescent:
    """등고선만으로 셀별 흐름 방향을 세운다. 보간면을 만들지 않는다."""
    from B03_FileInput.B03_FileInput_Service_Chain import _log_steps

    marks = [("시작", time.perf_counter())]
    rows, cols = spec.n_rows, spec.n_cols
    burned, levels = rasterize_contours(spec, contour_features, elevation_floor_m)
    marks.append(("등고선 굽기(rasterize_contours)", time.perf_counter()))
    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)
    marks.append(("밴드 표고(EDT 1회)", time.perf_counter()))
    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)
    window_widened = 0
    edt_steps = 0
    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, widened = _distance_to_lower(members, lower)
        window_widened += widened
        edt_steps += 1
        distance[members] = step_distance.astype(np.float32)
        target_row[members] = step_row
        target_col[members] = step_col
        band_rank[members] = len(levels) - 1 - rank  # 높을수록 큰 값

    # 최하단 밴드는 더 내려갈 등고 라인이 없다. 무효로 빼지 않고 **정지 셀**로 남긴다 —
    # 그래야 그 자리가 도로·세류선이면 적색으로 잡히고, 아니면 물이 고이는 지점으로 보인다.
    marks.append(
        (
            f"밴드별 하강거리(EDT {edt_steps}단 · 창 넓힘 {window_widened}회)",
            time.perf_counter(),
        )
    )
    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)

    marks.append(("위치에너지 조립", time.perf_counter()))
    receiver, step_length, azimuth = _route_by_potential(
        spec, potential, valid, target_row, target_col
    )
    marks.append(("흐름 경로(_route_by_potential)", time.perf_counter()))
    _log_steps(f"하강 방향장({rows}×{cols})", marks)
    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,
    )


# ── 흐름 방향 저장본(347-1) ─────────────────────────────────────────────
#
# 방향장은 등고선 · 남긴 등고 단 · 격자(원점 · 셀 · 범위)만으로 정해진다. 원점이 좌표 정수 m 라
# 노선이 저장 범위 안이면 읽어서 잘라 쓴다. 파일 이름 = descent_<입력 지문>_<범위 지문>.npz ·
# 종류마다 최근 2벌만 남긴다. 계산 규칙이 바뀌면 `DESCENT_FILE_VERSION` 을 올린다.
DESCENT_FILE_VERSION = 1
_DESCENT_KEEP = 2


def descent_input_key(
    contour_key: str,
    contour_features: list[dict[str, Any]],
    elevation_floor_m: float | None,
    cell_m: float,
) -> str:
    """방향장 입력 지문 — 도엽 등고선 지문 + 남긴 등고 단 + 셀 + 계산 규칙 값.

    남긴 단은 속성 표고만 본다 — 같은 도엽 · 같은 단 집합이면 남는 선도 같다(길이 거름은 결정적).
    """
    levels = sorted(
        {
            elevation
            for feature in contour_features
            if feature.get("geometry")
            and (elevation := _feature_elevation(feature.get("properties") or {})) is not None
            and (elevation_floor_m is None or elevation >= elevation_floor_m)
        }
    )
    text = json.dumps(
        [
            DESCENT_FILE_VERSION,
            contour_key,
            levels,
            cell_m,
            DRAINAGE_CONTOUR_MIN_LENGTH_M,
            DESCENT_WINDOW_MARGIN_CELLS,
            DESCENT_TILE_CELLS,
        ]
    )
    return hashlib.sha1(text.encode(), usedforsecurity=False).hexdigest()[:16]


def iter_saved_descents(directory: Path, input_key: str) -> Iterator[ContourDescent]:
    """같은 입력 지문으로 저장한 방향장 — 최근 것부터. 못 읽는 파일은 건너뜀."""
    paths = sorted(
        directory.glob(f"descent_{input_key}_*.npz"),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )
    for path in paths:
        try:
            with np.load(path, allow_pickle=False) as data:
                x_min, y_max, cell_m, n_rows, n_cols = data["spec"].tolist()
                yield ContourDescent(
                    spec=GridSpec(
                        float(x_min), float(y_max), float(cell_m), int(n_rows), int(n_cols)
                    ),
                    band_elevation=data["band_elevation"],
                    valid=data["valid"],
                    receiver=data["receiver"],
                    step_length=data["step_length"],
                    azimuth=data["azimuth"],
                    levels=data["levels"].tolist(),
                )
        except (OSError, KeyError, ValueError):
            logger.warning("배수유역: 흐름 방향 저장본을 읽지 못해 건너뜁니다 (%s).", path)


def save_descent(directory: Path, input_key: str, descent: ContourDescent) -> None:
    """방향장을 저장하고 최근 2벌만 남긴다. 실패는 로그만(계산은 이어감)."""
    spec = descent.spec
    extent = f"{spec.x_min:.3f},{spec.y_max:.3f},{spec.n_rows},{spec.n_cols}"
    extent_key = hashlib.sha1(extent.encode(), usedforsecurity=False).hexdigest()[:8]
    path = directory / f"descent_{input_key}_{extent_key}.npz"
    try:
        directory.mkdir(parents=True, exist_ok=True)
        np.savez(
            path,
            spec=np.array([spec.x_min, spec.y_max, spec.cell_m, spec.n_rows, spec.n_cols]),
            band_elevation=descent.band_elevation,
            valid=descent.valid,
            receiver=descent.receiver,
            step_length=descent.step_length,
            azimuth=descent.azimuth,
            levels=np.array(descent.levels, dtype=np.float64),
        )
        olds = sorted(directory.glob("descent_*.npz"), key=lambda item: item.stat().st_mtime)
        for old in olds[:-_DESCENT_KEEP]:
            old.unlink(missing_ok=True)
    except OSError:
        logger.warning("배수유역: 흐름 방향 저장본을 남기지 못했습니다 (%s).", path)


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)