Aislo/B04_PreProcess/B04_PreProcess_Engine_ModelBuild.py

"""B04 지표면 5종 표현 빌더 (TIN/DTM/NURBS/implicit/meshfree).

TerrainContext에서 파생 격자·샘플을 받아 각 표현의 모델(npz)과 프리뷰
(GLB/PLY)를 저장하고 메타데이터를 반환한다.
"""

from pathlib import Path
from typing import Any

import numpy as np
from scipy.interpolate import RBFInterpolator
from scipy.spatial import Delaunay

from B04_PreProcess.B04_PreProcess_Engine_ModelContext import (
    ProgressCallback,
    TerrainContext,
    artifact_size,
    atomic_npz,
    clip_and_compact_mesh,
    evaluate_nurbs_spline,
    fit_nurbs_spline,
    grid_faces,
    grid_vertices,
    with_footprint,
    write_binary_ply,
    write_glb,
)


def build_tin(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    progress(5)
    points = context.sample(int(context.config["tin_max_input_points"]))
    unique_xy, unique_indices = np.unique(points[:, :2], axis=0, return_index=True)
    points = points[unique_indices]
    if len(points) < 3:
        raise ValueError("TIN 생성에 필요한 포인트가 부족합니다.")
    progress(25)
    faces = np.asarray(Delaunay(unique_xy).simplices, dtype=np.uint32)
    if len(faces):
        triangle_xy = points[faces, :2]
        edges = np.stack(
            [
                np.linalg.norm(triangle_xy[:, 0] - triangle_xy[:, 1], axis=1),
                np.linalg.norm(triangle_xy[:, 1] - triangle_xy[:, 2], axis=1),
                np.linalg.norm(triangle_xy[:, 2] - triangle_xy[:, 0], axis=1),
            ],
            axis=1,
        )
        faces = faces[np.max(edges, axis=1) <= float(context.config["tile_size_meters"]) * 2]
    valid_vertices = context.contains_xy(points[:, 0], points[:, 1])
    points, faces = clip_and_compact_mesh(points, faces, valid_vertices)
    if not len(faces):
        raise ValueError("외곽 안쪽 기준 적용 후 TIN 면이 남지 않았습니다.")
    progress(65)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(model_path, vertices=points, faces=faces)
    write_glb(preview_path, points, faces, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "triangular_mesh",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "vertex_count": int(len(points)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


def build_dtm(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    progress(10)
    resolution = float(context.config["dtm_grid_resolution_meters"])
    x_coords, y_coords, z_grid = context.grid(resolution)
    progress(55)
    preview_x, preview_y, preview_z = context.preview_grid(resolution)
    vertices = grid_vertices(preview_x, preview_y, preview_z)
    faces = grid_faces(len(preview_y), len(preview_x))
    valid_grid = context.contains_xy(*np.meshgrid(x_coords, y_coords)).reshape(
        len(y_coords), len(x_coords)
    )
    preview_valid = context.contains_xy(vertices[:, 0], vertices[:, 1])
    vertices, faces = clip_and_compact_mesh(vertices, faces, preview_valid)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(
        model_path,
        x=x_coords,
        y=y_coords,
        z=z_grid,
        valid_mask=valid_grid,
        resolution=np.array([resolution], np.float32),
    )
    progress(75)
    write_glb(preview_path, vertices, faces, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "regular_grid",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "grid_rows": int(len(y_coords)),
            "grid_columns": int(len(x_coords)),
            "grid_resolution_meters": resolution,
            "vertex_count": int(len(vertices)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


def build_nurbs(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    degree = max(1, min(5, int(context.config["nurbs_degree"])))
    patch_size = float(context.config["nurbs_patch_size_meters"])
    controls = max(degree + 1, int(context.config["nurbs_control_points_per_axis"]))
    control_resolution = max(patch_size / max(controls - 1, 1), 0.25)
    x_control, y_control, z_control = context.grid(control_resolution)
    progress(30)
    spline, z_range = fit_nurbs_spline(x_control, y_control, z_control, degree)
    x_preview, y_preview, _ = context.preview_grid(
        float(context.config["dtm_grid_resolution_meters"])
    )
    z_preview = evaluate_nurbs_spline(spline, y_preview, x_preview, z_range, stem).astype(
        np.float32
    )
    progress(65)
    vertices = grid_vertices(x_preview, y_preview, z_preview)
    faces = grid_faces(len(y_preview), len(x_preview))
    valid_preview = context.contains_xy(vertices[:, 0], vertices[:, 1])
    vertices, faces = clip_and_compact_mesh(vertices, faces, valid_preview)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(
        model_path,
        control_x=x_control,
        control_y=y_control,
        control_z=z_control,
        degree=np.array([degree], np.int16),
        patch_size_meters=np.array([patch_size], np.float32),
    )
    write_glb(preview_path, vertices, faces, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "bspline_surface",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "degree": degree,
            "control_rows": int(len(y_control)),
            "control_columns": int(len(x_control)),
            "vertex_count": int(len(vertices)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


def build_implicit(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    maximum = max(100, int(context.config["implicit_max_points_per_tile"]))
    points = context.sample(maximum)
    unique_xy, unique_indices = np.unique(points[:, :2], axis=0, return_index=True)
    points = points[unique_indices]
    if len(points) < 4:
        raise ValueError("Implicit 생성에 필요한 포인트가 부족합니다.")
    progress(20)
    interpolator = RBFInterpolator(
        unique_xy.astype(np.float64),
        points[:, 2].astype(np.float64),
        neighbors=min(64, len(points)),
        smoothing=float(context.config["implicit_smoothing"]),
        kernel="thin_plate_spline",
    )
    x_preview, y_preview, _ = context.preview_grid(
        float(context.config["dtm_grid_resolution_meters"])
    )
    xx, yy = np.meshgrid(x_preview, y_preview)
    query = np.column_stack([xx.ravel(), yy.ravel()])
    z_values = np.empty(len(query), dtype=np.float32)
    for start in range(0, len(query), 50_000):
        end = min(start + 50_000, len(query))
        z_values[start:end] = interpolator(query[start:end]).astype(np.float32)
        progress(25 + int(45 * end / len(query)))
    z_grid = z_values.reshape(len(y_preview), len(x_preview))
    vertices = grid_vertices(x_preview, y_preview, z_grid)
    faces = grid_faces(len(y_preview), len(x_preview))
    valid_preview = context.contains_xy(vertices[:, 0], vertices[:, 1])
    vertices, faces = clip_and_compact_mesh(vertices, faces, valid_preview)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(
        model_path,
        centers_xy=unique_xy.astype(np.float32),
        center_z=points[:, 2].astype(np.float32),
        smoothing=np.array([float(context.config["implicit_smoothing"])], np.float32),
    )
    write_glb(preview_path, vertices, faces, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "local_rbf_height_field",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "center_count": int(len(points)),
            "vertex_count": int(len(vertices)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


def build_meshfree(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    points = context.sample(int(context.config["meshfree_max_model_points"]))
    points = points[context.contains_xy(points[:, 0], points[:, 1])]
    if not len(points):
        raise ValueError("외곽 안쪽 기준 적용 후 Meshfree 포인트가 남지 않았습니다.")
    resolution = float(context.config["dtm_grid_resolution_meters"])
    x_grid, y_grid, z_grid = context.grid(resolution)
    dz_dy, dz_dx = np.gradient(z_grid, resolution, resolution)
    gx = np.clip(np.searchsorted(x_grid, points[:, 0]), 0, len(x_grid) - 1)
    gy = np.clip(np.searchsorted(y_grid, points[:, 1]), 0, len(y_grid) - 1)
    normals = np.column_stack([-dz_dx[gy, gx], -dz_dy[gy, gx], np.ones(len(points), np.float32)])
    normals /= np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1e-9)
    progress(55)
    preview_max = int(context.config["max_preview_vertices"])
    if len(points) > preview_max:
        selection = np.linspace(0, len(points) - 1, preview_max, dtype=np.int64)
        preview_points, preview_normals = points[selection], normals[selection]
    else:
        preview_points, preview_normals = points, normals
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.ply"
    radius = float(context.config["meshfree_point_radius_meters"])
    atomic_npz(
        model_path,
        points=points,
        normals=normals.astype(np.float32),
        radius=np.array([radius], np.float32),
    )
    write_binary_ply(preview_path, preview_points, preview_normals, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "meshfree_surfels",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "application/octet-stream",
            "point_count": int(len(points)),
            "preview_point_count": int(len(preview_points)),
            "point_radius_meters": radius,
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


BUILDERS = {
    "tin": build_tin,
    "dtm": build_dtm,
    "nurbs": build_nurbs,
    "implicit": build_implicit,
    "meshfree": build_meshfree,
}