"""등고선 → 지형 격자(라플라스) — 공용 지형 창고가 쓰는 흠 고친 판(PLAN 607-1 · 요구 486~491).

B04 `B04_PreProcess_Engine_SheetSurface.py` 라플라스 길을 본뜸 · B 파일은 안 고침
(B 는 옛 판 그대로).
고친 흠 둘:
  ① 굽기 — rasterio `all_touched=False` 는 선 방향 · 조각에 따라 다른 칸을 칠함(같은 선 뒤집어도
     최대 9 m) → 선분 끝을 정렬한 뒤 큰 축 칸 중심마다 한 칸 + 꼭짓점 칸(방향 무관 · 굵기 1 칸).
  ② 폐합 안쪽 채움 — 봉우리 · 웅덩이 방향 · 기울기를 같은 높이 선 전체가 아니라 그 고리 둘레로만
     → 범위가 바뀌어도 같은 고리는 같은 값(조각 = 통째).
보간 자체(거리 비례 + 라플라스 400 회)는 `common_util_terrain_laplace`
(B04 `build_laplace` 짝 · 비트 동일).
폐합 채움은 수위 칸 상자 · 묶음 창 안에서만 셈(통격자 셈과 같은 값 · PLAN 617-2).
"""

from __future__ import annotations

import logging
from typing import Any

import numpy as np

from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import GridSpec
from common_util.common_util_terrain_laplace import laplace_surface

logger = logging.getLogger(__name__)

# 계산 방식 판 — 바뀌면 올림 · 창고 열쇠가 어긋나 저장본을 다시 셈.
TERRAIN_ALGO = 1

# 채움 기울기를 재는 고리 밖 띠(칸) — B04 와 같은 값.
_BAND_CELLS = 8


def segments_from_lines(lines: list[tuple[np.ndarray, float]]) -> tuple[np.ndarray, ...]:
    """(꼭짓점 (N,2) · 높이) 목록 → 선분 끝 a · b (M,2) · 높이 z (M,)."""
    starts, ends, levels = [], [], []
    for coords, level in lines:
        coords = np.asarray(coords, dtype=np.float64)[:, :2]
        if len(coords) < 2:
            continue
        starts.append(coords[:-1])
        ends.append(coords[1:])
        levels.append(np.full(len(coords) - 1, float(level)))
    if not starts:
        empty = np.zeros((0, 2))
        return empty, empty, np.zeros(0)
    return np.vstack(starts), np.vstack(ends), np.concatenate(levels)


def burn_segments(spec: GridSpec, a: np.ndarray, b: np.ndarray, z: np.ndarray) -> np.ndarray:
    """선분을 격자에 굽는다 — 칸 = 그 위를 지나는 선의 높이 · 그 밖 NaN · 겹치면 낮은 쪽.

    방향 무관: 선분 끝을 (x, y) 사전 차례로 맞춘 뒤 큰 축의 칸 중심마다 한 칸을 칠함 ·
    꼭짓점 칸도 칠해 이음매 빈칸을 막음 · 선분마다 따로라 조각 · 이어 붙임에도 같음.
    """
    burned = np.full((spec.n_rows, spec.n_cols), np.inf, dtype=np.float64)
    if len(z):
        swap = (a[:, 0] > b[:, 0]) | ((a[:, 0] == b[:, 0]) & (a[:, 1] > b[:, 1]))
        p = np.where(swap[:, None], b, a)
        q = np.where(swap[:, None], a, b)
        # 연속 칸 좌표 — 칸 c 의 중심 = c + 0.5
        u0 = (p[:, 0] - spec.x_min) / spec.cell_m
        u1 = (q[:, 0] - spec.x_min) / spec.cell_m
        v0 = (spec.y_max - p[:, 1]) / spec.cell_m
        v1 = (spec.y_max - q[:, 1]) / spec.cell_m
        x_major = np.abs(u1 - u0) >= np.abs(v1 - v0)
        rows, cols, levels = [], [], []
        for major, (s0, s1, t0, t1) in ((True, (u0, u1, v0, v1)), (False, (v0, v1, u0, u1))):
            pick = x_major if major else ~x_major
            s0, s1, t0, t1, zz = s0[pick], s1[pick], t0[pick], t1[pick], z[pick]
            lo = np.ceil(np.minimum(s0, s1) - 0.5).astype(np.int64)
            hi = np.floor(np.maximum(s0, s1) - 0.5).astype(np.int64)
            count = np.clip(hi - lo + 1, 0, None)
            seg = np.repeat(np.arange(len(lo)), count)
            step = np.arange(int(count.sum())) - np.repeat(np.cumsum(count) - count, count)
            major_idx = lo[seg] + step
            span = s1[seg] - s0[seg]
            t = (major_idx + 0.5 - s0[seg]) / np.where(span == 0, 1.0, span)
            minor_idx = np.floor(t0[seg] + t * (t1[seg] - t0[seg])).astype(np.int64)
            r, c = (minor_idx, major_idx) if major else (major_idx, minor_idx)
            rows.append(r)
            cols.append(c)
            levels.append(zz[seg])
        for ends in (a, b):  # 꼭짓점 칸
            rows.append(np.floor((spec.y_max - ends[:, 1]) / spec.cell_m).astype(np.int64))
            cols.append(np.floor((ends[:, 0] - spec.x_min) / spec.cell_m).astype(np.int64))
            levels.append(z)
        r, c, zz = np.concatenate(rows), np.concatenate(cols), np.concatenate(levels)
        inside = (r >= 0) & (r < spec.n_rows) & (c >= 0) & (c < spec.n_cols)
        np.minimum.at(burned, (r[inside], c[inside]), zz[inside])
    burned[np.isinf(burned)] = np.nan
    return burned.astype(np.float32)


def resolve_enclosed_interiors(
    burned: np.ndarray, present: list[float], surface: np.ndarray, cell_m: float, interval_m: float
) -> int:
    """폐합 등고선 안쪽(봉우리 · 웅덩이)을 그 고리 바깥 사면 경사로 연장한다(in-place).

    B04 `_resolve_enclosed_interiors` 와 같은 규칙 · 다른 점 = 방향(바깥 평균)과 기울기 띠를
    그 칸 묶음을 감싼 고리 칸 바깥으로만 잼(같은 높이 다른 선은 안 섞음). 채운 곳 수를 돌려줌.
    채움 · 묶음은 그 수위 칸 상자 + 1 칸 빈 덧댐 안에서 · 묶음은 `find_objects` 창 —
    통격자와 같은 값 · 차례(수위 → 묶음 번호) 그대로.
    """
    from scipy.ndimage import (
        binary_dilation,
        binary_fill_holes,
        distance_transform_edt,
        find_objects,
        label,
    )

    constrained = np.isfinite(burned)
    limit = max(interval_m - 0.5, 0.5)
    pad = _BAND_CELLS + 2
    eight = np.ones((3, 3), dtype=bool)
    resolved = 0
    rows_all, cols_all = np.nonzero(constrained)
    values = burned[rows_all, cols_all]
    order = np.argsort(values, kind="stable")
    rows_all, cols_all, values = rows_all[order], cols_all[order], values[order]
    for level in present:
        lo = np.searchsorted(values, level, side="left")
        hi = np.searchsorted(values, level, side="right")
        if hi <= lo:
            continue
        r0, r1 = int(rows_all[lo:hi].min()), int(rows_all[lo:hi].max()) + 1
        c0, c1 = int(cols_all[lo:hi].min()), int(cols_all[lo:hi].max()) + 1
        box_mask = np.pad(burned[r0:r1, c0:c1] == level, 1)
        interior = binary_fill_holes(box_mask) & ~box_mask
        if not interior.any():
            continue
        components, count = label(interior)
        for component_id, obj in enumerate(find_objects(components), start=1):
            whole = components[obj] == component_id
            top, left = obj[0].start - 1 + r0, obj[1].start - 1 + c0
            bottom, right = top + whole.shape[0], left + whole.shape[1]
            if (constrained[top:bottom, left:right] & whole).any():
                continue  # 안에 다른 제약이 있으면 마루가 아니다(보통의 감싸는 링)
            win = (
                slice(max(top - pad, 0), bottom + pad),
                slice(max(left - pad, 0), right + pad),
            )
            line = burned[win] == level
            component = np.zeros(line.shape, dtype=bool)
            component[
                top - win[0].start : bottom - win[0].start,
                left - win[1].start : right - win[1].start,
            ] = whole
            local = surface[win]
            # 그 묶음을 감싼 고리 칸 → 고리 바로 바깥
            loop = binary_dilation(component, structure=eight) & line
            ring = binary_dilation(loop, structure=eight) & ~component & ~line
            ring &= np.isfinite(local)
            if not ring.any():
                continue
            direction = 1.0 if float(local[ring].mean()) < level else -1.0
            inner = distance_transform_edt(component, sampling=cell_m)
            outer_distance = distance_transform_edt(~(loop | component), sampling=cell_m)
            band = (outer_distance > 0) & (outer_distance <= _BAND_CELLS * cell_m)
            band &= np.isfinite(local) & ~component & ~line
            if band.any():
                slope = float(
                    np.mean(np.abs(level - local[band].astype(np.float64)) / outer_distance[band])
                )
            else:
                slope = 0.0
            if slope > 1e-3:
                offset = np.minimum(slope * inner[component], limit)
            else:
                peak = float(inner.max())
                offset = (limit / 2.0) * (inner[component] / peak) if peak > 0 else 0.0
            local[component] = level + direction * offset
            resolved += 1
    return resolved


def build_contour_surface(
    spec: GridSpec, a: np.ndarray, b: np.ndarray, z: np.ndarray
) -> np.ndarray | None:
    """선분 → (R, C) float32 표고(북쪽 첫 행 · 없음 NaN). 등고 두 단 미만이면 None."""
    burned = burn_segments(spec, a, b, z)
    present = sorted(np.unique(burned[np.isfinite(burned)]).tolist())
    if len(present) < 2:
        return None
    interval_m = float(np.diff(np.array(present)).min())
    surface = laplace_surface(burned, spec.cell_m)
    resolve_enclosed_interiors(burned, present, surface, spec.cell_m, interval_m)
    return surface


def surface_grid(spec: GridSpec, surface: np.ndarray) -> dict[str, Any]:
    """DtmGridSampler 꼴 — x · y 오름차순 · z (ny, nx) · valid_mask."""
    z = surface[::-1, :].astype(np.float32)
    return {
        "x": spec.cell_centers_x(),
        "y": spec.cell_centers_y()[::-1],
        "z": z,
        "valid_mask": np.isfinite(z),
    }