Aislo/B04_wf1_Surface/B04_wf1_Surface_Engine_Watershed_Expand.py

"""해석 영역 확장 — 최외곽의 **적색 셀** 주변으로 넓히며 다시 분석한다.

적색 셀이 해석 영역 최외곽에 있다는 것은 그 바깥에서 물이 더 흘러 들어온다는 뜻이다.
거기서 멈추면 유역이 잘린다. 반대로 최외곽이 전부 파랑이면 그 바깥 물은 도로로 오지
않으므로 더 볼 필요가 없다.

  ① 현재 해석 영역의 최외곽 셀 중 **적색**인 것을 찾는다
  ② 그 주변으로 한 겹(설정 폭) 넓힌다
  ③ 넓힌 영역으로 흐름 방향·색을 다시 분석한다
  ④ **새로 추가된 셀에 적색이 하나도 없으면 종료** (2026-07-31 사용자 지시)

격자 bbox에 닿으면 격자 자체도 셀 정수배로 넓힌다 — 도로 시작점 기준 격자점은 유지된다.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

import numpy as np
from shapely.geometry import LineString

from B04_wf1_Surface.B04_wf1_Surface_Engine_Watershed_Descent import (
    ContourDescent,
    build_contour_descent,
)
from B04_wf1_Surface.B04_wf1_Surface_Engine_Watershed_Flow import (
    FlowClassification,
    RoadRaster,
    burn_stream_flow,
    classify_flow,
    outermost_cells,
    rasterize_road,
)
from B04_wf1_Surface.B04_wf1_Surface_Engine_Watershed_Grid import GridSpec, TerrainGrid
from config.config_system import (
    DRAINAGE_RED_EXPAND_BAND_M,
    DRAINAGE_RED_EXPAND_MAX_ROUNDS,
    DRAINAGE_ROAD_WIDTH_M,
)

logger = logging.getLogger(__name__)


@dataclass
class GridAnalysis:
    """한 회차 분석 결과 — 격자·해석 영역·방향·색까지 한 묶음."""

    spec: GridSpec
    domain: np.ndarray  # (R, C) bool — 해석 대상 셀
    descent: ContourDescent
    terrain: TerrainGrid
    road: RoadRaster
    flow: FlowClassification


@dataclass
class RedExpansion:
    """확장 루프 결과."""

    analysis: GridAnalysis
    rounds: int  # 실제로 넓힌 횟수 (0 = 처음부터 최외곽에 적색이 없었음)
    closed: bool  # 새로 추가한 셀에 적색이 없어 스스로 멈췄는가
    added_cells: int  # 확장으로 늘어난 셀 수


def analyze_domain(
    spec: GridSpec,
    domain: np.ndarray,
    contour_features: list[dict[str, Any]],
    route_line: LineString,
    upstream_streams: list[LineString],
    elevation_floor_m: float | None = None,
    descent: ContourDescent | None = None,
) -> GridAnalysis | None:
    """주어진 격자·해석 영역에 대해 등고선 하강 방향과 도로 도달 색을 한 번 계산한다.

    **하강 방향장은 해석 영역과 무관하다** — 등고선 기하만으로 정해진다. 그래서 확장
    회차마다 다시 계산하지 않고, 격자가 커졌을 때만 새로 만들어 넘겨받는다(`descent`).
    해석 영역은 마지막에 마스크로만 씌운다.
    """
    if descent is None or descent.spec != spec:
        descent = build_contour_descent(spec, contour_features, None, elevation_floor_m)
    valid = descent.valid & domain
    if not valid.any():
        return None
    terrain = TerrainGrid(
        spec=spec,
        elevation=np.where(valid, descent.band_elevation, np.nan).astype(np.float32),
        valid=valid,
        receiver=descent.receiver,
        step_length=descent.step_length,
    )
    road = rasterize_road(spec, route_line, DRAINAGE_ROAD_WIDTH_M)
    terrain, burned = burn_stream_flow(terrain, road, upstream_streams)
    flow = classify_flow(terrain, road, burned, azimuth=descent.azimuth)
    return GridAnalysis(
        spec=spec, domain=domain, descent=descent, terrain=terrain, road=road, flow=flow
    )


def expand_by_red_boundary(
    spec: GridSpec,
    domain: np.ndarray,
    contour_features: list[dict[str, Any]],
    route_line: LineString,
    upstream_streams: list[LineString],
    elevation_floor_m: float | None = None,
    band_m: float = DRAINAGE_RED_EXPAND_BAND_M,
    max_rounds: int = DRAINAGE_RED_EXPAND_MAX_ROUNDS,
) -> RedExpansion | None:
    """최외곽 적색 셀 주변으로 넓히며, 새로 추가한 셀에 적색이 없을 때까지 반복한다."""
    band_cells = max(1, int(round(band_m / spec.cell_m)))
    # 1차 영역의 bbox는 영역에 딱 붙어 있어 첫 회차부터 격자를 넓혀야 한다. 미리 여유를
    # 두면 방향장을 다시 만들지 않고 해석 영역만 넓히며 몇 회차를 돌 수 있다.
    spec, domain = _pad_spec(spec, domain, band_cells * 2)
    started_cells = int(domain.sum())
    analysis = analyze_domain(
        spec, domain, contour_features, route_line, upstream_streams, elevation_floor_m
    )
    if analysis is None:
        return None

    rounds = 0
    closed = False
    for attempt in range(max_rounds):
        current = analysis.spec
        reaches = analysis.flow.reaches_road.reshape(current.n_rows, current.n_cols)
        rim_red = outermost_cells(analysis.domain) & reaches
        if not rim_red.any():
            closed = True  # 최외곽이 전부 파랑 — 바깥 물은 도로로 오지 않는다
            break

        grown_spec, grown_domain, grown_rim = _grow_for_rim(
            current, analysis.domain, rim_red, band_cells
        )
        widened = grown_domain | _dilate_by(grown_rim, band_cells)
        added_mask = widened & ~grown_domain
        if not added_mask.any():
            closed = True
            break

        logger.info(
            "배수유역: %d회차 확장 — 최외곽 적색 %d셀 주변 %.0fm, 셀 %d개 추가",
            attempt + 1,
            int(rim_red.sum()),
            band_m,
            int(added_mask.sum()),
        )
        widened_analysis = analyze_domain(
            grown_spec,
            widened,
            contour_features,
            route_line,
            upstream_streams,
            elevation_floor_m,
            # 격자가 그대로면 방향장을 재사용한다 — 등고선 기하가 안 바뀌었으므로 결과는 같다.
            descent=analysis.descent if grown_spec == current else None,
        )
        if widened_analysis is None:
            break
        analysis = widened_analysis
        rounds += 1

        added_reaches = widened_analysis.flow.reaches_road.reshape(
            grown_spec.n_rows, grown_spec.n_cols
        )
        if not (added_mask & added_reaches).any():
            closed = True  # 새로 추가한 셀에 적색이 없다 — 여기까지가 유역이다
            logger.info("배수유역: 새로 추가한 셀에 적색이 없어 확장을 멈춥니다.")
            break
    else:
        logger.warning("배수유역: 확장 상한(%d회)에 도달했습니다.", max_rounds)

    added = int(analysis.domain.sum()) - started_cells
    logger.info(
        "배수유역: 확장 %d회, 셀 %d → %d (+%d), %s",
        rounds,
        started_cells,
        int(analysis.domain.sum()),
        added,
        "닫힘" if closed else "미닫힘",
    )
    return RedExpansion(analysis=analysis, rounds=rounds, closed=closed, added_cells=added)


def _dilate_by(mask: np.ndarray, steps: int) -> np.ndarray:
    """8이웃 팽창을 `steps`회 반복한다(정사각 커널이라 반경 = steps 셀)."""
    rows, cols = mask.shape
    result = mask
    for _ in range(steps):
        padded = np.zeros((rows + 2, cols + 2), dtype=bool)
        padded[1:-1, 1:-1] = result
        grown = np.zeros_like(result)
        for row_shift in (0, 1, 2):
            for col_shift in (0, 1, 2):
                grown |= padded[row_shift : row_shift + rows, col_shift : col_shift + cols]
        result = grown
    return result


def _pad_spec(spec: GridSpec, domain: np.ndarray, cells: int) -> tuple[GridSpec, np.ndarray]:
    """격자에 사방 여유를 두고 해석 영역 마스크를 그 안으로 옮겨 담는다."""
    if cells <= 0:
        return spec, domain
    padded_spec = GridSpec(
        x_min=spec.x_min - cells * spec.cell_m,
        y_max=spec.y_max + cells * spec.cell_m,
        cell_m=spec.cell_m,
        n_rows=spec.n_rows + 2 * cells,
        n_cols=spec.n_cols + 2 * cells,
    )
    padded = np.zeros((padded_spec.n_rows, padded_spec.n_cols), dtype=bool)
    padded[cells : cells + spec.n_rows, cells : cells + spec.n_cols] = domain
    return padded_spec, padded


def _grow_for_rim(
    spec: GridSpec, domain: np.ndarray, rim: np.ndarray, band_cells: int
) -> tuple[GridSpec, np.ndarray, np.ndarray]:
    """적색 최외곽이 격자 bbox에 닿았으면 그 방향으로 격자를 넓히고 마스크를 옮겨 담는다.

    격자는 셀 정수배로만 넓히므로 도로 시작점 기준 격자점이 그대로 유지된다.
    """
    north = band_cells if rim[0, :].any() else 0
    south = band_cells if rim[-1, :].any() else 0
    west = band_cells if rim[:, 0].any() else 0
    east = band_cells if rim[:, -1].any() else 0
    if not (north or south or west or east):
        return spec, domain, rim

    grown = GridSpec(
        x_min=spec.x_min - west * spec.cell_m,
        y_max=spec.y_max + north * spec.cell_m,
        cell_m=spec.cell_m,
        n_rows=spec.n_rows + north + south,
        n_cols=spec.n_cols + west + east,
    )
    new_domain = np.zeros((grown.n_rows, grown.n_cols), dtype=bool)
    new_rim = np.zeros_like(new_domain)
    new_domain[north : north + spec.n_rows, west : west + spec.n_cols] = domain
    new_rim[north : north + spec.n_rows, west : west + spec.n_cols] = rim
    logger.info(
        "배수유역: 격자 확대 %d×%d → %d×%d (북%d 남%d 서%d 동%d 셀)",
        spec.n_rows,
        spec.n_cols,
        grown.n_rows,
        grown.n_cols,
        north,
        south,
        west,
        east,
    )
    return grown, new_domain, new_rim