"""도엽등고선 → 표고 격자 보간 방식 모음 (비교용).

문헌(Hutchinson 1988/89 ANUDEM; Chaplot 2006; Arun 2013 등)은 지형 복잡도·자료 밀도에
따라 우열이 갈리며 단일 최적해가 없다고 본다. 그래서 방식을 하나로 고르지 않고 여기
모아 두고 B04 화면에서 바꿔 가며 보게 한다(2026-08-30 사용자 지시).

각 builder는 `(spec, burned, features, cell_m) -> (R, C) float32` 격자를 돌려준다.
`burned`는 등고 라인이 구워진 격자(라인 셀 = 표고, 그 외 NaN)다. 폐합 링 안쪽 처리와
프리뷰·저장은 호출측(`_SheetSurface`)이 방식과 무관하게 똑같이 해 준다.
"""

import logging
import warnings
from typing import Any, Callable

import numpy as np

logger = logging.getLogger(__name__)

# 화면 버튼에 쓰는 이름 — 키는 surface_models.generation_params.source_filter 접미사다.
SHEET_METHOD_LABELS: dict[str, str] = {
    "distance": "거리비례",
    "biharmonic": "TPS(박판)",
    "anudem": "ANUDEM형",
    "multires": "다중해상도",
    "laplace": "라플라스",
    "tin": "TIN 선형",
    "clough": "TIN 곡면",
    "idw": "IDW",
}


def _laplacian(values: np.ndarray) -> np.ndarray:
    padded = np.pad(values, 1, mode="edge")
    return (
        padded[:-2, 1:-1] + padded[2:, 1:-1] + padded[1:-1, :-2] + padded[1:-1, 2:] - 4.0 * values
    )


def _contour_vertices(burned: np.ndarray, spec: Any) -> tuple[np.ndarray, np.ndarray]:
    """등고 라인 셀을 (N,2) 세계좌표와 표고로 바꾼다."""
    rows, cols = np.nonzero(np.isfinite(burned))
    xs = spec.cell_centers_x()[cols]
    ys = spec.cell_centers_y()[rows]
    return np.column_stack([xs, ys]), burned[rows, cols].astype(np.float64)


def _grid_points(spec: Any) -> tuple[np.ndarray, np.ndarray]:
    grid_x, grid_y = np.meshgrid(spec.cell_centers_x(), spec.cell_centers_y())
    return grid_x, grid_y


# ── ① 거리 비례 ──────────────────────────────────────────────────────────────
def build_distance(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """가장 가까운 서로 다른 표고 두 라인 사이를 거리 비례로 나눈다.

        z = (L1·d2 + L2·d1) / (d1 + d2)

    지도 제작의 고전적 손보간을 그대로 옮긴 것이다. 원뿔·능선(z=r) 형상을 정확히
    재현하고 계단이 생기지 않는다. 표고별 거리장을 돌며 가장 작은 두 값을 추적하므로
    L1≠L2가 보장된다.
    """
    from scipy.ndimage import distance_transform_edt

    levels = np.unique(burned[np.isfinite(burned)])
    if len(levels) < 2:
        return np.full(burned.shape, levels[0] if len(levels) else np.nan, dtype=np.float32)

    infinity = np.float32(np.inf)
    first_d = np.full(burned.shape, infinity, dtype=np.float32)
    first_z = np.zeros(burned.shape, dtype=np.float32)
    second_d = np.full(burned.shape, infinity, dtype=np.float32)
    second_z = np.zeros(burned.shape, dtype=np.float32)
    for level in levels:
        distance = distance_transform_edt(burned != level, sampling=cell_m).astype(np.float32)
        beats_first = distance < first_d
        second_d = np.where(beats_first, first_d, second_d)
        second_z = np.where(beats_first, first_z, second_z)
        first_d = np.where(beats_first, distance, first_d)
        first_z = np.where(beats_first, np.float32(level), first_z)
        beats_second = ~beats_first & (distance < second_d)
        second_d = np.where(beats_second, distance, second_d)
        second_z = np.where(beats_second, np.float32(level), second_z)

    total = first_d + second_d
    usable = np.isfinite(second_d) & (total > 1e-9)
    surface = first_z.astype(np.float32)
    surface[usable] = (
        (
            first_z[usable].astype(np.float64) * second_d[usable].astype(np.float64)
            + second_z[usable].astype(np.float64) * first_d[usable].astype(np.float64)
        )
        / total[usable].astype(np.float64)
    ).astype(np.float32)
    return surface


def relax_laplace(surface: np.ndarray, fixed: np.ndarray, iterations: int) -> None:
    """등고 라인을 고정한 채 이웃 평균으로 다듬는다 (in-place, red-black 순서)."""
    if iterations <= 0:
        return
    free = ~fixed & np.isfinite(surface)
    if not free.any():
        return
    rows, cols = np.indices(surface.shape)
    red = free & (((rows + cols) & 1) == 0)
    black = free & ~red
    padded = np.zeros((surface.shape[0] + 2, surface.shape[1] + 2), dtype=np.float32)
    for _ in range(iterations):
        for colour in (red, black):
            padded[1:-1, 1:-1] = surface
            padded[0, 1:-1] = surface[0]
            padded[-1, 1:-1] = surface[-1]
            padded[1:-1, 0] = surface[:, 0]
            padded[1:-1, -1] = surface[:, -1]
            neighbours = (
                padded[:-2, 1:-1] + padded[2:, 1:-1] + padded[1:-1, :-2] + padded[1:-1, 2:]
            ) * np.float32(0.25)
            surface[colour] = neighbours[colour]


# ── ② 라플라스(조화) ─────────────────────────────────────────────────────────
def build_laplace(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """등고선을 경계값으로 두고 Δz=0을 푼다.

    면이 매끈해지지만 z=r(원뿔·능선)은 조화함수가 아니라 마루가 눌린다. 비교 기준으로
    남겨 둔다 — ANUDEM이 라플라스 대신 박판 스플라인을 쓰는 이유를 눈으로 보기 위함.
    """
    surface = build_distance(spec, burned, features, cell_m)
    relax_laplace(surface, np.isfinite(burned), 400)
    return surface


# ── ③ 박판 스플라인(중조화) ──────────────────────────────────────────────────
def solve_min_curvature(constrained: np.ndarray, guess: np.ndarray) -> np.ndarray:
    """제약 셀을 고정하고 Δ²z=0(박판 스플라인)을 최소곡률 최소제곱으로 푼다.

    ANUDEM/Topo to Raster가 쓰는 박판 스플라인과 같은 연산자다. 라플라스와 달리 z=r을
    그대로 통과시켜 능선·마루가 눌리지 않고, 경사가 등고선 너머로 자연스럽게 이어진다.

    `constrained`: 값이 고정된 셀(등고 라인, 필요하면 구조선 앵커) — 그 외는 NaN.
    `guess`: 시작값이자 감쇠 기준(보통 거리 보간 결과).
    """
    from scipy.sparse.linalg import LinearOperator, lsmr

    guess = guess.astype(np.float64)
    burned = constrained
    # 제약 셀 + **격자 테두리**를 고정한다. 최외곽 등고선 바깥이 통째로 자유면
    # 1차함수가 Δ²의 영공간에 남아 해가 하나로 정해지지 않고 켤레기울기가 발산한다
    # (2026-08-30 실측: |Δz| 2092m). 테두리는 거리 보간값으로 묶는다.
    fixed = np.isfinite(burned)
    fixed[0, :] = fixed[-1, :] = True
    fixed[:, 0] = fixed[:, -1] = True
    free = ~fixed & np.isfinite(guess)
    if not free.any():
        return guess.astype(np.float32)
    index = np.flatnonzero(free.ravel())
    values = np.where(np.isfinite(burned), burned, guess)
    base = np.where(fixed, np.nan_to_num(values), 0.0)

    # Δ²z=0을 정규방정식(CG)으로 풀면 조건수가 격자변 4제곱이라 발산한다(실측).
    # 대신 **최소곡률** 최소제곱으로 세운다 — 자유 셀에 대해 ‖Δz‖를 최소화하며,
    # 그 정상해가 곧 Δ²z=0이다. 조건수가 제곱으로 줄어 LSMR이 안정적으로 푼다
    # (Briggs 1974의 최소곡률 격자화와 같은 목적함수).
    # 최소곡률만으로는 제약(등고선)에서 먼 영역이 정해지지 않아 해가 폭주한다.
    # ANUDEM의 거칠기 벌점과 같은 취지로 감쇠항을 붙여 거리 보간값에 묶어 둔다:
    #     minimize ‖Δz‖² + λ‖z − 거리보간‖²
    # λ가 작을수록 더 매끈하고 클수록 거리 보간에 가깝다.
    total_cells = base.size
    free_count = len(index)
    damping = np.float64(np.sqrt(0.02))
    anchor = guess.ravel()[index]

    def forward(vector: np.ndarray) -> np.ndarray:
        # 반드시 **선형**이어야 한다 — 고정 셀 기여(base)를 여기서 더하면 아핀이 되어
        # LSMR의 전제가 깨지고 해가 폭주한다. base 몫은 우변으로만 넘긴다.
        scattered = np.zeros_like(base)
        scattered.ravel()[index] = vector
        return np.concatenate([_laplacian(scattered).ravel(), damping * vector])

    def adjoint(vector: np.ndarray) -> np.ndarray:
        curvature = _laplacian(vector[:total_cells].reshape(base.shape)).ravel()[index]
        return curvature + damping * vector[total_cells:]

    rhs = np.concatenate([-_laplacian(base).ravel(), damping * anchor])
    linear = LinearOperator(
        (total_cells + free_count, free_count),
        matvec=forward,
        rmatvec=adjoint,
        dtype=np.float64,
    )
    result = lsmr(linear, rhs, x0=anchor, maxiter=400, atol=1e-8, btol=1e-8)
    solution, info = result[0], result[1]
    surface = base.copy()
    surface.ravel()[index] = solution
    surface[fixed] = values[fixed]
    # 안전장치 — 발산하면 조용히 틀린 지형을 넘기지 말고 거리 보간으로 되돌린다.
    drift = float(np.nanmax(np.abs(surface - guess)))
    span = float(np.nanmax(guess) - np.nanmin(guess))
    if not np.isfinite(drift) or drift > max(span, 1.0):
        logger.warning(
            "도엽 서피스(TPS): 해가 발산해(최대 %.1fm) 거리 보간으로 되돌립니다 (info=%s).",
            drift,
            info,
        )
        return guess.astype(np.float32)
    logger.info("도엽 서피스(TPS): 최소제곱 info=%s, 최대 변화 %.2fm", info, drift)
    return surface.astype(np.float32)


def build_biharmonic(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """등고선만 제약으로 둔 박판 스플라인."""
    return solve_min_curvature(burned, build_distance(spec, burned, features, cell_m))


# ── ④ TIN 선형 / ⑤ TIN 곡면 ─────────────────────────────────────────────────
def _triangulated(
    spec: Any, burned: np.ndarray, cell_m: float, smooth: bool, max_points: int = 120_000
) -> np.ndarray:
    from scipy.interpolate import CloughTocher2DInterpolator, LinearNDInterpolator

    points, values = _contour_vertices(burned, spec)
    if len(points) < 3:
        return np.full(burned.shape, np.nan, dtype=np.float32)
    if len(points) > max_points:  # 삼각망 비용은 정점 수에 비례한다
        step = int(np.ceil(len(points) / max_points))
        points, values = points[::step], values[::step]
    factory = CloughTocher2DInterpolator if smooth else LinearNDInterpolator
    interpolator = factory(points, values)
    grid_x, grid_y = _grid_points(spec)
    surface = np.empty(burned.shape, dtype=np.float32)
    chunk = max(1, int(4_000_000 // max(spec.n_cols, 1)))
    for start in range(0, spec.n_rows, chunk):
        stop = min(start + chunk, spec.n_rows)
        surface[start:stop] = interpolator(grid_x[start:stop], grid_y[start:stop]).astype(
            np.float32
        )
    return surface


def build_tin(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """등고선 정점 Delaunay 삼각망 선형 보간 — 가장 흔한 고전 방식.

    같은 표고 정점 3개로 이루어진 평탄 삼각형이 굴곡부·마루에 계단을 만든다. 비교
    기준으로 남긴다.
    """
    return _triangulated(spec, burned, cell_m, smooth=False)


def build_clough(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """같은 삼각망을 C1 곡면(Clough–Tocher)으로 채운다 — TIN 계단을 곡면으로 눌러 준다."""
    return _triangulated(spec, burned, cell_m, smooth=True)


# ── ⑥ IDW ───────────────────────────────────────────────────────────────────
def build_idw(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """등고 라인 셀을 표본으로 한 역거리가중(k=12, 거듭제곱 2).

    구현이 단순해 널리 쓰이지만 표본이 곧 등고선이라 라인 주변에 표고가 뭉치는
    '황소눈' 결함이 잘 드러난다. 비교 기준.
    """
    from scipy.spatial import cKDTree

    points, values = _contour_vertices(burned, spec)
    if len(points) < 2:
        return np.full(burned.shape, np.nan, dtype=np.float32)
    tree = cKDTree(points)
    grid_x, grid_y = _grid_points(spec)
    query = np.column_stack([grid_x.ravel(), grid_y.ravel()])
    neighbours = min(12, len(points))
    distance, index = tree.query(query, k=neighbours, workers=-1)
    if neighbours == 1:
        distance = distance[:, None]
        index = index[:, None]
    weight = 1.0 / np.maximum(distance, 1e-6) ** 2
    surface = (weight * values[index]).sum(axis=1) / weight.sum(axis=1)
    exact = distance[:, 0] < 1e-6
    surface[exact] = values[index[exact, 0]]
    return surface.reshape(burned.shape).astype(np.float32)


# ── ⑦ ANUDEM형 (구조선 + 배수 강제) ─────────────────────────────────────────
def _contour_corner_anchors(
    spec: Any, burned: np.ndarray, guess: np.ndarray, cell_m: float
) -> tuple[np.ndarray, np.ndarray] | None:
    """등고선의 국소 최대 곡률점(코너)에서 능선·계곡 구조선 앵커를 만든다.

    ANUDEM은 등고선 자체의 곡률에서 능선·계곡망을 먼저 뽑아 흐름 구조를 세운다
    (Hutchinson 1988/89). 여기서도 같은 순서를 따른다.

      ① 라인마다 정점 곡률을 재 국소 최대점(V자 꼭짓점)을 고른다
      ② 굽은 안쪽이 더 높으면 **계곡**(등고선 V가 상류를 가리킴), 낮으면 **능선**
      ③ 같은 종류의 코너를 이웃 표고끼리 이어 그 사이를 선형 보간해 앵커로 심는다

    앵커는 박판 해의 제약으로 들어가 계곡 바닥이 이어져 내려가고 능선 마루가 선다.
    """
    from scipy.ndimage import label
    from scipy.spatial import cKDTree

    levels = np.unique(burned[np.isfinite(burned)])
    if len(levels) < 2:
        return None
    interval = float(np.diff(levels).min())
    xs = spec.cell_centers_x()
    ys = spec.cell_centers_y()

    corners: list[tuple[float, float, float, int]] = []  # x, y, level, +1 계곡 / -1 능선
    span = 6  # 곡률을 재는 정점 간격(px) — 짧으면 노이즈, 길면 꼭짓점을 놓친다
    for level in levels:
        labelled, count = label(burned == level)
        for component_id in range(1, count + 1):
            line_rows, line_cols = np.nonzero(labelled == component_id)
            if len(line_rows) < 3 * span:
                continue
            # 라인 셀을 한 줄로 세운다 — 좌표 정렬로 근사한다(정밀 추적은 과하다).
            order = np.argsort(line_cols + line_rows * 1e-3)
            path = np.column_stack([line_cols[order], line_rows[order]]).astype(np.float64)
            before = np.roll(path, span, axis=0)
            after = np.roll(path, -span, axis=0)
            first = path - before
            second = after - path
            first_len = np.hypot(first[:, 0], first[:, 1])
            second_len = np.hypot(second[:, 0], second[:, 1])
            valid = (first_len > 1e-6) & (second_len > 1e-6)
            cosine = np.ones(len(path))
            cosine[valid] = (first[valid] * second[valid]).sum(axis=1) / (
                first_len[valid] * second_len[valid]
            )
            sharp = np.flatnonzero(valid & (cosine < 0.3))  # 70도 이상 꺾인 자리
            for i in sharp[:: max(1, span)]:
                # 굽은 안쪽 방향 = 두 변 단위벡터 합의 반대
                inward = -(first[i] / first_len[i] + second[i] / second_len[i])
                norm = float(np.hypot(inward[0], inward[1]))
                if norm < 1e-6:
                    continue
                probe = path[i] + inward / norm * 6.0
                probe_col = int(round(probe[0]))
                probe_row = int(round(probe[1]))
                if not (0 <= probe_row < guess.shape[0] and 0 <= probe_col < guess.shape[1]):
                    continue
                inside = float(guess[probe_row, probe_col])
                if not np.isfinite(inside) or abs(inside - level) < interval * 0.15:
                    continue
                corners.append(
                    (
                        float(xs[int(path[i, 0])]),
                        float(ys[int(path[i, 1])]),
                        float(level),
                        1 if inside > level else -1,
                    )
                )
    if len(corners) < 4:
        logger.info("도엽 서피스(ANUDEM형): 등고선 코너가 부족해 구조선을 건너뜁니다.")
        return None

    array = np.asarray(corners, dtype=np.float64)
    anchor_xy: list[np.ndarray] = []
    anchor_z: list[np.ndarray] = []
    reach = interval * 20.0  # 이보다 먼 코너는 같은 구조선으로 보지 않는다
    for level in levels[:-1]:
        upper = level + interval
        lower_set = array[np.abs(array[:, 2] - level) < 1e-6]
        upper_set = array[np.abs(array[:, 2] - upper) < 1e-6]
        if not len(lower_set) or not len(upper_set):
            continue
        tree = cKDTree(upper_set[:, :2])
        distance, index = tree.query(lower_set[:, :2], k=1)
        for i in range(len(lower_set)):
            j = int(index[i])
            if distance[i] > reach or lower_set[i, 3] != upper_set[j, 3]:
                continue
            start = lower_set[i, :2]
            end = upper_set[j, :2]
            steps = max(2, int(distance[i] / max(cell_m, 1e-6) / 4))
            fraction = np.linspace(0.0, 1.0, steps + 1)[1:-1]
            if not len(fraction):
                continue
            anchor_xy.append(start + (end - start) * fraction[:, None])
            anchor_z.append(level + interval * fraction)
    if not anchor_xy:
        return None
    return np.vstack(anchor_xy), np.concatenate(anchor_z)


def _enforce_drainage(surface: np.ndarray, epsilon: float = 0.01) -> int:
    """가짜 웅덩이를 메운다 — ANUDEM의 배수 강제와 같은 목적.

    등고선만으로 만든 면에는 흐름이 끊기는 웅덩이가 남는다. 형태학적 재구성(erosion)
    으로 채우되 완전 평탄해지지 않게 아주 작은 값을 얹는다. 채운 셀 수를 반환한다.
    """
    from skimage.morphology import reconstruction

    if not np.isfinite(surface).all():
        return 0
    seed = np.full(surface.shape, float(surface.max()), dtype=np.float64)
    seed[0, :] = surface[0, :]
    seed[-1, :] = surface[-1, :]
    seed[:, 0] = surface[:, 0]
    seed[:, -1] = surface[:, -1]
    filled = reconstruction(seed, surface.astype(np.float64), method="erosion")
    raised = filled > surface + 1e-6
    if not raised.any():
        return 0
    # 그냥 채우면 웅덩이가 통째로 평탄해져 흐름 방향이 없어진다. 채운 영역 안쪽으로
    # 갈수록 아주 조금 높아지게 해서 물이 가장자리(넘침점)로 빠져나가게 둔다
    # (Garbrecht·Martz의 평탄면 해소를 간단히 옮긴 것 — 표고 변화는 cm 단위다).
    from scipy.ndimage import distance_transform_edt

    inner = distance_transform_edt(raised)
    surface[raised] = (filled[raised] + epsilon * inner[raised]).astype(surface.dtype)
    return int(raised.sum())


def build_anudem(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """ANUDEM형 — 등고선 곡률에서 능선·계곡 구조선을 뽑아 제약에 더하고, 박판으로 풀고,
    가짜 웅덩이를 메운다. Topo to Raster가 밟는 세 단계를 그대로 옮긴 것이다."""
    guess = build_distance(spec, burned, features, cell_m).astype(np.float64)
    constrained = burned.astype(np.float64).copy()
    anchors = _contour_corner_anchors(spec, burned, guess, cell_m)
    if anchors is not None:
        xy, z = anchors
        row, col = spec.world_to_rc(xy[:, 0].copy(), xy[:, 1].copy())
        inside = (row >= 0) & (col >= 0)
        row, col, z = row[inside], col[inside], z[inside]
        free = ~np.isfinite(constrained[row, col])
        constrained[row[free], col[free]] = z[free]
        logger.info("도엽 서피스(ANUDEM형): 구조선 앵커 %d셀", int(free.sum()))
    surface = solve_min_curvature(constrained, guess).astype(np.float32)
    logger.info("도엽 서피스(ANUDEM형): 가짜 웅덩이 %d셀 메움", _enforce_drainage(surface))
    return surface


# ── ⑧ 다중해상도 (coarse → fine) ────────────────────────────────────────────
def build_multires(spec: Any, burned: np.ndarray, features: Any, cell_m: float) -> np.ndarray:
    """성긴 격자에서 풀고 점차 세밀화한다 — ANUDEM의 다중해상도 전략.

    전체 형상은 성긴 격자에서 싸게 잡고, 세밀한 격자에서는 등고선 근처만 다듬는다.
    한 해상도에서만 풀 때보다 넓은 밴드가 고르게 퍼지고 값싸게 수렴한다.
    """
    from scipy.ndimage import zoom

    surface: np.ndarray | None = None
    for factor in (8, 4, 2, 1):
        if factor == 1:
            coarse = burned
        else:
            # 성긴 격자의 제약 — 블록 안 등고 라인의 평균 표고를 대표로 쓴다.
            rows = burned.shape[0] // factor * factor
            cols = burned.shape[1] // factor * factor
            blocks = burned[:rows, :cols].reshape(rows // factor, factor, cols // factor, factor)
            # 라인이 하나도 없는 블록은 NaN이 정상이라 경고를 삼킨다.
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", RuntimeWarning)
                coarse = np.nanmean(blocks, axis=(1, 3)).astype(np.float32)
        level = build_distance(spec, coarse, features, cell_m * factor)
        if surface is not None:
            # 앞 단계 해를 지금 해상도로 올려 절반씩 섞는다 — 성긴 단계의 넓은 추세를
            # 이어받되 이번 해상도의 등고선 정보를 덮지 않는다.
            scale = (level.shape[0] / surface.shape[0], level.shape[1] / surface.shape[1])
            upscaled = zoom(surface, scale, order=1)
            free = ~np.isfinite(coarse)
            level[free] = (level[free] + upscaled[free]) * 0.5
        relax_laplace(level, np.isfinite(coarse), 8)
        surface = level
    assert surface is not None
    if surface.shape != burned.shape:  # 블록 자르기로 남은 가장자리 보정
        scale = (burned.shape[0] / surface.shape[0], burned.shape[1] / surface.shape[1])
        surface = zoom(surface, scale, order=1)
    line = np.isfinite(burned)
    surface[line] = burned[line]
    return surface.astype(np.float32)


SHEET_METHOD_BUILDERS: dict[str, Callable[[Any, np.ndarray, Any, float], np.ndarray]] = {
    "distance": build_distance,
    "biharmonic": build_biharmonic,
    "anudem": build_anudem,
    "multires": build_multires,
    "laplace": build_laplace,
    "tin": build_tin,
    "clough": build_clough,
    "idw": build_idw,
}