Aislo/common_util/common_util_surface_sampler.py

"""확정 지표면 모델 표고 sampler (공용).

종·횡단 생성기가 의존하는 최소 표고 조회 인터페이스와, 확정된 지표면 모델
(B04_wf1_Surface/models)을 일괄 XY 표고 sampler로 여는 팩토리를 제공한다.
DTM valid_mask를 footprint로 결합해 데이터가 없는 영역을 임의 표고로 메우지
않는다.

B05 종횡단 전용이었으나 B04 배수유역 세부 설계도 같은 표고면을 써야 해서
common_util로 옮겼다 — 두 화면의 종단 Z가 갈라지면 세부유역 경계가 달라진다
(2026-08-01 구조 개편).
"""

from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

import numpy as np
from scipy.interpolate import RegularGridInterpolator


class SurfaceElevationSampler(Protocol):
    """종·횡단 생성기가 의존하는 최소 표고 조회 인터페이스."""

    def sample_xy(self, xy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """(N, 2) 모델좌표 XY 배열에 대해 (z, valid)를 반환한다."""


@dataclass
class DtmGridSampler:
    """정규 DTM의 표고와 valid_mask를 보수적으로 조회한다.

    보간점 주변 네 격자 꼭짓점이 모두 유효할 때만 valid=True로 반환한다.
    """

    x: np.ndarray
    y: np.ndarray
    z: np.ndarray
    valid_mask: np.ndarray

    def __post_init__(self) -> None:
        self.x = np.asarray(self.x, dtype=np.float64).reshape(-1)
        self.y = np.asarray(self.y, dtype=np.float64).reshape(-1)
        self.z = np.asarray(self.z, dtype=np.float64)
        self.valid_mask = np.asarray(self.valid_mask, dtype=bool)
        if len(self.x) < 2 or len(self.y) < 2:
            raise ValueError("DTM 표고 조회에는 X/Y 축이 각각 2개 이상 필요합니다.")
        if self.z.shape != (len(self.y), len(self.x)):
            raise ValueError("DTM Z 격자 크기가 X/Y 축과 일치하지 않습니다.")
        if self.valid_mask.shape != self.z.shape:
            raise ValueError("DTM valid_mask 크기가 Z 격자와 일치하지 않습니다.")
        if self.x[0] > self.x[-1]:
            self.x = self.x[::-1]
            self.z = self.z[:, ::-1]
            self.valid_mask = self.valid_mask[:, ::-1]
        if self.y[0] > self.y[-1]:
            self.y = self.y[::-1]
            self.z = self.z[::-1, :]
            self.valid_mask = self.valid_mask[::-1, :]
        self._interpolator = RegularGridInterpolator(
            (self.y, self.x), self.z, method="linear", bounds_error=False, fill_value=np.nan
        )

    @classmethod
    def from_npz(cls, path: Path | str) -> "DtmGridSampler":
        with np.load(Path(path), allow_pickle=False) as data:
            return cls(data["x"], data["y"], data["z"], data["valid_mask"])

    def sample_xy(self, xy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        xy = np.asarray(xy, dtype=np.float64)
        if xy.ndim != 2 or xy.shape[1] != 2:
            raise ValueError("표고 조회 좌표는 (N, 2) XY 배열이어야 합니다.")
        if not len(xy):
            return np.empty(0, dtype=np.float64), np.empty(0, dtype=bool)

        z = np.asarray(self._interpolator(np.column_stack([xy[:, 1], xy[:, 0]])), dtype=np.float64)
        ix = np.searchsorted(self.x, xy[:, 0], side="right") - 1
        iy = np.searchsorted(self.y, xy[:, 1], side="right") - 1
        inside = (ix >= 0) & (iy >= 0) & (ix < len(self.x) - 1) & (iy < len(self.y) - 1)
        valid = np.zeros(len(xy), dtype=bool)
        selected = np.flatnonzero(inside)
        if len(selected):
            sx = ix[selected]
            sy = iy[selected]
            valid[selected] = (
                self.valid_mask[sy, sx]
                & self.valid_mask[sy, sx + 1]
                & self.valid_mask[sy + 1, sx]
                & self.valid_mask[sy + 1, sx + 1]
                & np.isfinite(z[selected])
            )
        z[~valid] = np.nan
        return z, valid


@dataclass
class CallableSurfaceSampler:
    """테스트와 어댑터에 사용할 함수 기반 sampler."""

    function: Callable[[np.ndarray], np.ndarray]

    def sample_xy(self, xy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        values = np.asarray(self.function(np.asarray(xy, dtype=np.float64)), dtype=np.float64)
        if values.shape != (len(xy),):
            raise ValueError("표고 함수는 입력 좌표 수와 같은 길이의 배열을 반환해야 합니다.")
        valid = np.isfinite(values)
        return values, valid


@dataclass
class InterpolatedSurfaceSampler:
    """불규칙/곡면 모델 보간기와 DTM footprint 유효성을 결합한다."""

    interpolator: Callable[[np.ndarray], np.ndarray]
    footprint: SurfaceElevationSampler | None = None

    def sample_xy(self, xy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        xy = np.asarray(xy, dtype=np.float64)
        values = np.asarray(self.interpolator(xy), dtype=np.float64).reshape(-1)
        valid = np.isfinite(values)
        if self.footprint is not None:
            _, footprint_valid = self.footprint.sample_xy(xy)
            valid &= footprint_valid
        values[~valid] = np.nan
        return values, valid


def build_surface_sampler(
    models_dir: Path | str, source_filter: str, method: str, smooth: bool
) -> SurfaceElevationSampler:
    """1단계 확정 모델을 종·횡단용 일괄 XY 표고 sampler로 연다."""
    models_dir = Path(models_dir)
    smooth_suffix = "_smooth" if smooth and method in {"dtm", "tin"} else ""
    dtm_smooth = models_dir / f"dtm_{source_filter}_smooth.npz"
    dtm_original = models_dir / f"dtm_{source_filter}.npz"
    dtm_path = dtm_smooth if smooth and dtm_smooth.exists() else dtm_original
    if not dtm_path.exists():
        raise FileNotFoundError(f"기준 DTM이 없습니다: {dtm_path.name}")
    footprint = DtmGridSampler.from_npz(dtm_path)
    if method == "dtm":
        return footprint

    if method == "tin":
        from scipy.interpolate import LinearNDInterpolator

        path = models_dir / f"tin_{source_filter}{smooth_suffix}.npz"
        if not path.exists() and smooth_suffix:
            path = models_dir / f"tin_{source_filter}.npz"
        with np.load(path, allow_pickle=False) as data:
            vertices = np.asarray(data["vertices"], dtype=np.float64)
        interpolator = LinearNDInterpolator(vertices[:, :2], vertices[:, 2], fill_value=np.nan)
        return InterpolatedSurfaceSampler(lambda xy: interpolator(xy), footprint)

    if method == "nurbs":
        from scipy.interpolate import RectBivariateSpline

        path = models_dir / f"nurbs_{source_filter}.npz"
        with np.load(path, allow_pickle=False) as data:
            control_x = np.asarray(data["control_x"], dtype=np.float64)
            control_y = np.asarray(data["control_y"], dtype=np.float64)
            control_z = np.asarray(data["control_z"], dtype=np.float64)
            degree = int(data["degree"][0]) if "degree" in data else 3
        spline = RectBivariateSpline(
            control_y,
            control_x,
            control_z,
            kx=min(degree, len(control_y) - 1),
            ky=min(degree, len(control_x) - 1),
        )
        return InterpolatedSurfaceSampler(lambda xy: spline.ev(xy[:, 1], xy[:, 0]), footprint)

    if method == "implicit":
        from scipy.interpolate import RBFInterpolator

        path = models_dir / f"implicit_{source_filter}.npz"
        with np.load(path, allow_pickle=False) as data:
            centers = np.asarray(data["centers_xy"], dtype=np.float64)
            center_z = np.asarray(data["center_z"], dtype=np.float64)
            smoothing = float(data["smoothing"][0]) if "smoothing" in data else 0.0
        interpolator = RBFInterpolator(
            centers,
            center_z,
            neighbors=min(64, len(centers)),
            smoothing=smoothing,
            kernel="thin_plate_spline",
        )
        return InterpolatedSurfaceSampler(lambda xy: interpolator(xy), footprint)

    if method == "meshfree":
        from scipy.interpolate import LinearNDInterpolator

        path = models_dir / f"meshfree_{source_filter}.npz"
        with np.load(path, allow_pickle=False) as data:
            points = np.asarray(data["points"], dtype=np.float64)
        interpolator = LinearNDInterpolator(points[:, :2], points[:, 2], fill_value=np.nan)
        return InterpolatedSurfaceSampler(lambda xy: interpolator(xy), footprint)

    raise ValueError(f"지원하지 않는 지표면 모델입니다: {method}")