Aislo/B04_PreProcess/B04_PreProcess_Engine_ModelContext.py

"""B04 지표면 모델 공통 컨텍스트 및 메시 유틸리티.

지면 마스크가 적용된 포인트에서 footprint(외곽), 격자, 프리뷰 격자를 만들고,
GLB/PLY 프리뷰 및 npz 모델을 원자적으로 저장하는 공통 기능을 제공한다.
5개 표현(TIN/DTM/NURBS/implicit/meshfree) 빌더가 이 컨텍스트를 공유한다.
"""

import hashlib
import json
import logging
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

import numpy as np
import trimesh
from scipy import ndimage
from scipy.interpolate import RectBivariateSpline

from common_util.common_util_atomic import atomic_write_bytes, atomic_write_npz

logger = logging.getLogger(__name__)

MODEL_VERSION = 1
MODEL_METHODS = ("tin", "dtm", "nurbs", "implicit", "meshfree")
SOURCE_FILTER_LABELS = {"grid_min_z": "Grid Min-Z", "csf": "CSF", "pmf": "PMF"}
ProgressCallback = Callable[[int], None]

# 대용량 포인트 배치 처리 크기
_BATCH_SIZE = 500_000

# NURBS 평활 계수 — 제어점당 허용 잔차제곱(m²).
# 0.01(RMS 0.1m)은 요구가 너무 빡빡해 FITPACK이 제어점 수(122x113)보다 많은
# knot(126x117)을 밀어넣고 계수가 발산했다 — csf 표고 1e105, grid_min_z 1e10
# (2026-08-17 실측). 1.0(RMS 1m)이면 knot이 17~35로 떨어지고 세 지면 필터 모두
# 데이터 표고 범위 안에 머문다. NURBS는 평활 곡면 표현이라 이 정도 완화가 맞다.
NURBS_RESIDUAL_PER_CONTROL_POINT = 1.0


def fit_nurbs_spline(
    x_control: np.ndarray, y_control: np.ndarray, z_control: np.ndarray, degree: int
) -> tuple[RectBivariateSpline, tuple[float, float]]:
    """제어 격자에 B-spline 곡면을 맞추고, 발산 판정용 표고 허용 범위를 함께 준다.

    모델 빌더와 등고선 엔진이 같은 곡면을 각각 만들므로 여기서 한 번에 정의한다.
    """
    spline = RectBivariateSpline(
        y_control,
        x_control,
        z_control,
        kx=min(degree, len(y_control) - 1),
        ky=min(degree, len(x_control) - 1),
        s=float(len(x_control) * len(y_control)) * NURBS_RESIDUAL_PER_CONTROL_POINT,
    )
    z_min = float(np.min(z_control))
    z_max = float(np.max(z_control))
    margin = max((z_max - z_min) * 0.5, 1.0)
    return spline, (z_min - margin, z_max + margin)


def evaluate_nurbs_spline(
    spline: RectBivariateSpline,
    y_coords: np.ndarray,
    x_coords: np.ndarray,
    z_range: tuple[float, float],
    label: str = "",
) -> np.ndarray:
    """스플라인을 평가하고 표고 허용 범위를 벗어난 값을 잘라낸다.

    평활 스플라인은 제어점이 성긴 구석에서 발산할 수 있다. 그대로 두면 float32
    캐스팅에서 inf가 되고 프리뷰 색상(nan)과 등고선 레벨 산출(`np.arange`가
    "Maximum allowed size exceeded"로 실패)까지 번진다.
    """
    z_values = np.asarray(spline(y_coords, x_coords), dtype=np.float64)
    low, high = z_range
    outside = np.count_nonzero(~np.isfinite(z_values) | (z_values < low) | (z_values > high))
    if outside:
        logger.warning(
            "NURBS 곡면이 표고 범위를 벗어나 잘라냈습니다: %s%d개 셀 (허용 %.1f~%.1fm)",
            f"{label} " if label else "",
            int(outside),
            low,
            high,
        )
    z_values = np.nan_to_num(z_values, nan=low, posinf=high, neginf=low)
    return np.clip(z_values, low, high)


def config_signature(config: dict[str, Any]) -> str:
    """지오메트리 원본과 무관한 등고선·스무딩 설정을 제외한 캐시 서명."""
    sig_config = {
        k: v
        for k, v in config.items()
        if not k.startswith("contour_") and not k.startswith("smoothing_")
    }
    encoded = json.dumps(sig_config, sort_keys=True, default=list).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()[:16]


def bounds_dict(bounds: np.ndarray) -> dict[str, list[float]]:
    return {
        "x": [float(bounds[0, 0]), float(bounds[0, 1])],
        "y": [float(bounds[1, 0]), float(bounds[1, 1])],
        "z": [float(bounds[2, 0]), float(bounds[2, 1])],
    }


def scene_vertices(vertices: np.ndarray, bounds: np.ndarray) -> np.ndarray:
    """모델 좌표를 뷰어(Y-up) 좌표계로 변환한다."""
    center = bounds.mean(axis=1)
    result = np.empty((len(vertices), 3), dtype=np.float32)
    result[:, 0] = vertices[:, 0] - center[0]
    result[:, 1] = vertices[:, 2] - center[2]
    result[:, 2] = -(vertices[:, 1] - center[1])
    return result


def height_colors(vertices: np.ndarray) -> np.ndarray:
    """표고에 따른 그라디언트 정점 색상(RGBA)을 만든다."""
    if not len(vertices):
        return np.empty((0, 4), dtype=np.uint8)
    z = vertices[:, 2]
    span = max(float(np.max(z) - np.min(z)), 1e-9)
    t = np.clip((z - np.min(z)) / span, 0.0, 1.0)
    colors = np.empty((len(vertices), 4), dtype=np.uint8)
    colors[:, 0] = np.clip(36 + 190 * t, 0, 255).astype(np.uint8)
    colors[:, 1] = np.clip(86 + 95 * np.sin(t * np.pi), 0, 255).astype(np.uint8)
    colors[:, 2] = np.clip(128 - 80 * t, 0, 255).astype(np.uint8)
    colors[:, 3] = 255
    return colors


def write_glb(path: Path, vertices: np.ndarray, faces: np.ndarray, bounds: np.ndarray) -> None:
    mesh = trimesh.Trimesh(
        vertices=scene_vertices(vertices, bounds),
        faces=np.asarray(faces, dtype=np.int64),
        vertex_colors=height_colors(vertices),
        process=False,
    )
    payload = mesh.export(file_type="glb")
    if not isinstance(payload, bytes):
        raise TypeError("GLB exporter did not return bytes")
    atomic_write_bytes(path, payload)


def write_binary_ply(
    path: Path, vertices: np.ndarray, normals: np.ndarray, bounds: np.ndarray
) -> None:
    verts = scene_vertices(vertices, bounds)
    scene_normals = np.empty_like(normals, dtype=np.float32)
    scene_normals[:, 0] = normals[:, 0]
    scene_normals[:, 1] = normals[:, 2]
    scene_normals[:, 2] = -normals[:, 1]
    colors = height_colors(vertices)
    dtype = np.dtype(
        [
            ("x", "<f4"),
            ("y", "<f4"),
            ("z", "<f4"),
            ("nx", "<f4"),
            ("ny", "<f4"),
            ("nz", "<f4"),
            ("red", "u1"),
            ("green", "u1"),
            ("blue", "u1"),
            ("alpha", "u1"),
        ]
    )
    records = np.empty(len(vertices), dtype=dtype)
    records["x"], records["y"], records["z"] = verts.T
    records["nx"], records["ny"], records["nz"] = scene_normals.T
    records["red"], records["green"], records["blue"], records["alpha"] = colors.T
    header = (
        "ply\nformat binary_little_endian 1.0\n"
        f"element vertex {len(vertices)}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property float nx\nproperty float ny\nproperty float nz\n"
        "property uchar red\nproperty uchar green\nproperty uchar blue\nproperty uchar alpha\n"
        "end_header\n"
    ).encode("ascii")
    atomic_write_bytes(path, header + records.tobytes())


def grid_faces(rows: int, cols: int) -> np.ndarray:
    """정규 격자의 삼각형 면 인덱스를 만든다."""
    if rows < 2 or cols < 2:
        return np.empty((0, 3), dtype=np.uint32)
    base = np.arange((rows - 1) * (cols - 1), dtype=np.uint32)
    row = base // (cols - 1)
    col = base % (cols - 1)
    top_left = row * cols + col
    faces = np.empty((len(base) * 2, 3), dtype=np.uint32)
    faces[0::2] = np.stack([top_left, top_left + cols, top_left + 1], axis=1)
    faces[1::2] = np.stack([top_left + 1, top_left + cols, top_left + cols + 1], axis=1)
    return faces


def clip_and_compact_mesh(
    vertices: np.ndarray, faces: np.ndarray, valid_vertices: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """footprint 내부 정점만 사용하는 면을 남기고 미사용 정점을 제거한다."""
    if not len(faces):
        return np.empty((0, 3), np.float32), np.empty((0, 3), np.uint32)
    kept_faces = faces[np.all(valid_vertices[faces], axis=1)]
    if not len(kept_faces):
        return np.empty((0, 3), np.float32), np.empty((0, 3), np.uint32)
    used = np.unique(kept_faces)
    remap = np.full(len(vertices), -1, dtype=np.int64)
    remap[used] = np.arange(len(used))
    return vertices[used], remap[kept_faces].astype(np.uint32)


def grid_vertices(x_coords: np.ndarray, y_coords: np.ndarray, z_grid: np.ndarray) -> np.ndarray:
    xx, yy = np.meshgrid(x_coords, y_coords)
    return np.column_stack([xx.ravel(), yy.ravel(), z_grid.ravel()]).astype(np.float32)


def artifact_size(*paths: Path) -> int:
    return int(sum(path.stat().st_size for path in paths if path.exists()))


@dataclass
class TerrainContext:
    """지면 마스크가 적용된 포인트 집합에서 파생 격자·footprint를 캐싱한다."""

    xyz: np.ndarray
    mask: np.ndarray
    bounds: np.ndarray
    config: dict[str, Any]
    _indices: np.ndarray | None = None
    _samples: dict[int, np.ndarray] = field(default_factory=dict)
    _grids: dict[float, tuple[np.ndarray, np.ndarray, np.ndarray]] = field(default_factory=dict)
    _footprint: tuple[float, float, float, np.ndarray] | None = None

    @property
    def source_count(self) -> int:
        return int(np.count_nonzero(self.mask))

    def indices(self) -> np.ndarray:
        if self._indices is None:
            self._indices = np.flatnonzero(self.mask)
        return self._indices

    def sample(self, maximum: int) -> np.ndarray:
        maximum = max(3, int(maximum))
        if maximum in self._samples:
            return self._samples[maximum]
        indices = self.indices()
        if len(indices) > maximum:
            positions = np.linspace(0, len(indices) - 1, maximum, dtype=np.int64)
            indices = indices[positions]
        points = np.asarray(self.xyz[indices], dtype=np.float32)
        self._samples[maximum] = points
        return points

    def footprint(self) -> tuple[float, float, float, np.ndarray]:
        if self._footprint is not None:
            return self._footprint
        resolution = max(float(self.config.get("footprint_resolution_meters", 1.0)), 0.1)
        x_min, x_max = self.bounds[0]
        y_min, y_max = self.bounds[1]
        cols = max(2, int(math.ceil((x_max - x_min) / resolution)) + 1)
        rows = max(2, int(math.ceil((y_max - y_min) / resolution)) + 1)
        occupied = np.zeros((rows, cols), dtype=bool)
        indices = self.indices()
        for start in range(0, len(indices), _BATCH_SIZE):
            points = np.asarray(self.xyz[indices[start : start + _BATCH_SIZE]], dtype=np.float32)
            gx = np.clip(((points[:, 0] - x_min) / resolution).astype(np.int32), 0, cols - 1)
            gy = np.clip(((points[:, 1] - y_min) / resolution).astype(np.int32), 0, rows - 1)
            occupied[gy, gx] = True
        if not occupied.any():
            raise ValueError("기준 필터에 footprint를 만들 포인트가 없습니다.")

        close_cells = max(
            0,
            int(math.ceil(float(self.config.get("footprint_gap_close_meters", 1.0)) / resolution)),
        )
        footprint = occupied
        if close_cells:
            padded = np.pad(footprint, close_cells, mode="constant", constant_values=False)
            padded = ndimage.binary_closing(
                padded, structure=np.ones((3, 3), dtype=bool), iterations=close_cells
            )
            footprint = padded[close_cells:-close_cells, close_cells:-close_cells]
        if bool(self.config.get("keep_largest_footprint", True)):
            labels, component_count = ndimage.label(
                footprint, structure=np.ones((3, 3), dtype=bool)
            )
            if component_count:
                sizes = np.bincount(labels.ravel())
                sizes[0] = 0
                footprint = labels == int(np.argmax(sizes))
        footprint = ndimage.binary_fill_holes(footprint)

        inset_cells = max(
            0, int(math.ceil(float(self.config.get("boundary_inset_meters", 1.0)) / resolution))
        )
        if inset_cells:
            footprint = ndimage.binary_erosion(
                footprint,
                structure=np.ones((3, 3), dtype=bool),
                iterations=inset_cells,
                border_value=0,
            )
        if not footprint.any():
            raise ValueError("외곽 안쪽 기준 적용 후 유효한 footprint가 없습니다.")
        self._footprint = (float(x_min), float(y_min), resolution, footprint)
        return self._footprint

    def contains_xy(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        x_min, y_min, resolution, footprint = self.footprint()
        gx = np.floor((np.asarray(x) - x_min) / resolution).astype(np.int64)
        gy = np.floor((np.asarray(y) - y_min) / resolution).astype(np.int64)
        valid = (gx >= 0) & (gx < footprint.shape[1]) & (gy >= 0) & (gy < footprint.shape[0])
        result = np.zeros(np.broadcast(x, y).shape, dtype=bool)
        result[valid] = footprint[gy[valid], gx[valid]]
        return result

    def footprint_metadata(self) -> dict[str, Any]:
        _, _, resolution, footprint = self.footprint()
        return {
            "footprint_area_m2": round(float(footprint.sum()) * resolution * resolution, 3),
            "footprint_resolution_meters": resolution,
            "boundary_inset_meters": float(self.config.get("boundary_inset_meters", 1.0)),
        }

    def grid(self, resolution: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        resolution = max(float(resolution), 0.05)
        cached = self._grids.get(resolution)
        if cached is not None:
            return cached
        x_min, x_max = self.bounds[0]
        y_min, y_max = self.bounds[1]
        cols = max(2, int(math.ceil((x_max - x_min) / resolution)) + 1)
        rows = max(2, int(math.ceil((y_max - y_min) / resolution)) + 1)
        grid = np.full((rows, cols), np.inf, dtype=np.float32)
        indices = self.indices()
        for start in range(0, len(indices), _BATCH_SIZE):
            points = np.asarray(self.xyz[indices[start : start + _BATCH_SIZE]], dtype=np.float32)
            gx = np.clip(((points[:, 0] - x_min) / resolution).astype(np.int32), 0, cols - 1)
            gy = np.clip(((points[:, 1] - y_min) / resolution).astype(np.int32), 0, rows - 1)
            np.minimum.at(grid, (gy, gx), points[:, 2])
        missing = ~np.isfinite(grid)
        if missing.all():
            raise ValueError("기준 필터에 지면 포인트가 없습니다.")
        if missing.any():
            nearest = ndimage.distance_transform_edt(
                missing, return_distances=False, return_indices=True
            )
            grid = grid[tuple(nearest)]
        x_coords = np.linspace(x_min, x_max, cols, dtype=np.float32)
        y_coords = np.linspace(y_min, y_max, rows, dtype=np.float32)
        result = (x_coords, y_coords, grid)
        self._grids[resolution] = result
        return result

    def preview_grid(
        self, preferred_resolution: float
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        x_span = max(float(self.bounds[0, 1] - self.bounds[0, 0]), preferred_resolution)
        y_span = max(float(self.bounds[1, 1] - self.bounds[1, 0]), preferred_resolution)
        maximum = max(4, int(self.config["max_preview_vertices"]))
        predicted = (x_span / preferred_resolution + 1) * (y_span / preferred_resolution + 1)
        if predicted > maximum:
            preferred_resolution *= math.sqrt(predicted / maximum)
        return self.grid(preferred_resolution)

    def clear_caches(self) -> None:
        self._samples.clear()
        self._grids.clear()
        self._indices = None


def with_footprint(context: TerrainContext, metadata: dict[str, Any]) -> dict[str, Any]:
    metadata.update(context.footprint_metadata())
    return metadata


# 모델 빌더가 사용하는 원자적 저장 래퍼 (공통 유틸 재노출)
atomic_npz = atomic_write_npz