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

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

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

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

하강 방향장은 처음에 **넉넉한 큰 격자**로 한 번 짓고, 회차마다 그 격자에서 잘라 쓴다
(`DESCENT_SPARE_M` · 342-3). 격자가 커질 때마다 다시 짓던 것(8회 · 30 s)을 한 번으로 줄인다.
저장소(`descent_store` = 폴더 · 도엽 등고선 지문)를 주면 지은 것을 남기고, 덮는 저장본이 있으면
읽기만 한다(347-1).
"""

from __future__ import annotations

import logging
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
from shapely.geometry import LineString

from B04_PreProcess.B04_PreProcess_Engine_Watershed_Descent import (
    ContourDescent,
    build_contour_descent,
    descent_input_key,
    iter_saved_descents,
    save_descent,
)
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Flow import (
    FlowClassification,
    RoadRaster,
    burn_stream_flow,
    classify_flow,
    outermost_cells,
    rasterize_road,
)
from B04_PreProcess.B04_PreProcess_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__)

# 하강 방향장을 지을 때 해석 격자 둘레에 더 두는 여유(m). 확장이 이 안에서 끝나면 방향장을
# 한 번만 짓는다 — 검증_sub3 실측 확장은 첫 격자에서 남 350 m · 서 200 m 까지 넓어졌다.
# 넘으면 그때 넓힌 격자 둘레로 다시 짓는다(결과는 같고 시간만 듦).
DESCENT_SPARE_M = 400.0


class AnalysisStopped(Exception):
    """새 차례가 와서 계산을 그만둠(「지금 노선으로 셈」 · 347-2) — 단계 사이에서만 확인."""


@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`)을 받으면 잘라 쓴다.
    해석 영역은 마지막에 마스크로만 씌운다.
    """
    from B03_FileInput.B03_FileInput_Service_Chain import _log_steps

    marks = [("시작", time.perf_counter())]
    fitted = crop_descent(descent, spec) if descent is not None else None
    if fitted is None:
        descent = build_contour_descent(spec, contour_features, None, elevation_floor_m)
    else:
        descent = fitted
    marks.append(
        (
            "하강 방향장" + ("(잘라 씀)" if fitted is not None else "(새로 만듦)"),
            time.perf_counter(),
        )
    )
    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,
    )
    marks.append(("지형 격자 조립", time.perf_counter()))
    road = rasterize_road(spec, route_line, DRAINAGE_ROAD_WIDTH_M)
    marks.append(("도로 굽기(rasterize_road)", time.perf_counter()))
    terrain, burned = burn_stream_flow(terrain, road, upstream_streams)
    marks.append(("세류 새김(burn_stream_flow)", time.perf_counter()))
    flow = classify_flow(terrain, road, burned, azimuth=descent.azimuth)
    marks.append(("도로 도달 판정(classify_flow)", time.perf_counter()))
    _log_steps(f"격자 해석 1회({spec.n_rows}×{spec.n_cols})", marks)
    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,
    descent_store: tuple[Path, str] | None = None,
    should_stop: Callable[[], bool] | None = None,
) -> RedExpansion | None:
    """최외곽 적색 셀 주변으로 넓히며, 새로 추가한 셀에 적색이 없을 때까지 반복한다."""
    _expand_started = time.perf_counter()
    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())
    spare_cells = max(0, int(round(DESCENT_SPARE_M / spec.cell_m)))
    store = None
    if descent_store is not None:
        folder, contour_key = descent_store
        store = (
            folder,
            descent_input_key(contour_key, contour_features, elevation_floor_m, spec.cell_m),
        )

    def source_for(grid: GridSpec, mask: np.ndarray) -> ContourDescent:
        """`grid` 를 덮는 큰 방향장 — 저장본이 덮으면 읽기만 · 아니면 둘레 여유째 짓고 남김."""
        if store is not None:
            for saved in iter_saved_descents(*store):
                if crop_descent(saved, grid) is not None:
                    logger.info(
                        "배수유역: 흐름 방향 저장본을 읽어 씁니다(%d×%d).",
                        saved.spec.n_rows,
                        saved.spec.n_cols,
                    )
                    return saved
        built = build_contour_descent(
            _pad_spec(grid, mask, spare_cells)[0], contour_features, None, elevation_floor_m
        )
        if store is not None:
            save_descent(*store, built)
        return built

    source = source_for(spec, domain)
    analysis = analyze_domain(
        spec, domain, contour_features, route_line, upstream_streams, elevation_floor_m, source
    )
    if analysis is None:
        return None

    logger.info(
        "[계측] 확장 0회차(첫 해석) %.1fs · 격자 %d×%d · 셀 %d",
        time.perf_counter() - _expand_started,
        analysis.spec.n_rows,
        analysis.spec.n_cols,
        started_cells,
    )
    rounds = 0
    closed = False
    for attempt in range(max_rounds):
        if should_stop is not None and should_stop():
            raise AnalysisStopped
        _round_started = time.perf_counter()
        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
        if crop_descent(source, grown_spec) is None:
            # 넉넉히 둔 방향장 밖으로 넓어짐 — 넓힌 격자 둘레로 다시 짓는다.
            source = source_for(grown_spec, grown_domain)

        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=source,
        )
        if widened_analysis is None:
            break
        analysis = widened_analysis
        rounds += 1
        logger.info(
            "[계측] 확장 %d회차 %.1fs · 격자 %d×%d · 셀 %d (+%d)",
            attempt + 1,
            time.perf_counter() - _round_started,
            grown_spec.n_rows,
            grown_spec.n_cols,
            int(widened.sum()),
            int(added_mask.sum()),
        )

        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 crop_descent(source: ContourDescent, spec: GridSpec) -> ContourDescent | None:
    """큰 방향장에서 `spec` 격자 몫을 잘라 낸다. 덮지 못하거나 셀이 안 맞으면 None.

    잘린 가장자리 밖을 가리키던 셀은 제자리(싱크)로 둔다 — 해석 영역 밖으로 빠지는 물과
    같은 판정(미도달)이다. 화살표(방위)는 원래 방향 그대로 둔다.
    """
    big = source.spec
    if big == spec:
        return source
    if big.cell_m != spec.cell_m:
        return None
    row_shift = (big.y_max - spec.y_max) / spec.cell_m
    col_shift = (spec.x_min - big.x_min) / spec.cell_m
    row0, col0 = int(round(row_shift)), int(round(col_shift))
    if (
        abs(row_shift - row0) > 1e-6
        or abs(col_shift - col0) > 1e-6
        or row0 < 0
        or col0 < 0
        or row0 + spec.n_rows > big.n_rows
        or col0 + spec.n_cols > big.n_cols
    ):
        return None
    rows = slice(row0, row0 + spec.n_rows)
    cols = slice(col0, col0 + spec.n_cols)

    def window(flat: np.ndarray) -> np.ndarray:
        return flat.reshape(big.n_rows, big.n_cols)[rows, cols]

    target = window(source.receiver).astype(np.int64)
    target_row = target // big.n_cols - row0
    target_col = target % big.n_cols - col0
    inside = (
        (target_row >= 0)
        & (target_row < spec.n_rows)
        & (target_col >= 0)
        & (target_col < spec.n_cols)
    )
    local = np.arange(spec.size, dtype=np.int64).reshape(spec.n_rows, spec.n_cols)
    receiver = np.where(inside, target_row * spec.n_cols + target_col, local)
    return ContourDescent(
        spec=spec,
        band_elevation=source.band_elevation[rows, cols].copy(),
        valid=source.valid[rows, cols].copy(),
        receiver=receiver.reshape(-1).astype(np.int32),
        step_length=np.where(inside, window(source.step_length), 0.0)
        .reshape(-1)
        .astype(np.float32),
        azimuth=window(source.azimuth).reshape(-1).copy(),
        levels=source.levels,
    )


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에 닿았으면 그 방향으로 격자를 넓히고 마스크를 옮겨 담는다.

    격자는 셀 정수배로만 넓히므로 좌표 정수 m 기준 격자점이 그대로 유지된다.
    """
    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