"""등고 격자 → 라플라스 면 — 공용 지형 창고가 쓰는 빠른 판(PLAN 617-2 · 요구 504 · 506).

B04 짝: `B04_PreProcess_Engine_SheetMethods.build_laplace`(거리 비례 + `relax_laplace` 400 회)와
비트 동일 · 짝 시험 `test_617_2_laplace_mirror.py` · B 파일은 안 고침.
빠른 까닭 둘:
  ① 라플라스 — 격자를 홀짝 네 덩이(`S[p::2, q::2]` · 연속 메모리)로 나눠 그 색 칸만 셈 ·
     덧셈 차례(위 + 아래 + 왼 + 오른) × 0.25 는 원본과 같음.
  ② 거리 비례 — `np.where` 새 배열 대신 `np.copyto(where=)`.
"""

from __future__ import annotations

import numpy as np


def distance_surface(burned: np.ndarray, cell_m: float) -> np.ndarray:
    """B04 `build_distance` 와 같은 값 — 가장 가까운 서로 다른 두 수위 사이 거리 비례."""
    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)
    beats_second = np.empty(burned.shape, dtype=bool)
    for level in levels:
        distance = distance_transform_edt(burned != level, sampling=cell_m).astype(np.float32)
        beats_first = distance < first_d
        np.copyto(second_d, first_d, where=beats_first)
        np.copyto(second_z, first_z, where=beats_first)
        np.copyto(first_d, distance, where=beats_first)
        np.copyto(first_z, np.float32(level), where=beats_first)
        np.less(distance, second_d, out=beats_second)
        beats_second &= ~beats_first
        np.copyto(second_d, distance, where=beats_second)
        np.copyto(second_z, np.float32(level), where=beats_second)

    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_checker(surface: np.ndarray, fixed: np.ndarray, iterations: int) -> None:
    """B04 `relax_laplace` 와 같은 값(in-place · 빨강 → 검정) — 홀짝 네 덩이로 셈.

    덩이 (p, q) 마다 가장자리 한 칸 덧댄 버퍼.
    (p, q) 칸의 위 · 아래 = (1-p, q) 버퍼의 연이은 두 행 ·
    왼 · 오른 = (p, 1-q) 버퍼의 연이은 두 열 · 격자 끝은 자기 값(원본 edge 덧대기) → 셈 직전
    그 덧댄 줄만 자기 첫 · 끝 줄로 채움.
    """
    if iterations <= 0:
        return
    free = ~fixed & np.isfinite(surface)
    if not free.any():
        return
    quads = {(p, q): (p + q) & 1 for p in (0, 1) for q in (0, 1)}  # 0 = 빨강
    buffers: dict[tuple[int, int], np.ndarray] = {}
    for p, q in quads:
        block = surface[p::2, q::2]
        buffer = np.zeros((block.shape[0] + 2, block.shape[1] + 2), dtype=np.float32)
        buffer[1:-1, 1:-1] = block
        buffers[p, q] = buffer
    frees = {key: np.ascontiguousarray(free[key[0] :: 2, key[1] :: 2]) for key in quads}
    temps = {key: np.empty(buffers[key][1:-1, 1:-1].shape, dtype=np.float32) for key in quads}
    quarter = np.float32(0.25)
    for _ in range(iterations):
        for colour in (0, 1):
            for (p, q), quad_colour in quads.items():
                if quad_colour != colour:
                    continue
                own = buffers[p, q][1:-1, 1:-1]
                rows, cols = own.shape
                if own.size == 0:
                    continue
                vertical = buffers[1 - p, q]
                horizontal = buffers[p, 1 - q]
                if p == 0:
                    vertical[0, 1:-1] = own[0]
                if p + rows == vertical.shape[0] - 1:
                    vertical[-1, 1:-1] = own[-1]
                if q == 0:
                    horizontal[1:-1, 0] = own[:, 0]
                if q + cols == horizontal.shape[1] - 1:
                    horizontal[1:-1, -1] = own[:, -1]
                t = temps[p, q]
                np.add(vertical[p : p + rows, 1:-1], vertical[p + 1 : p + 1 + rows, 1:-1], out=t)
                t += horizontal[1:-1, q : q + cols]
                t += horizontal[1:-1, q + 1 : q + 1 + cols]
                t *= quarter
                np.copyto(own, t, where=frees[p, q])
    for p, q in quads:
        surface[p::2, q::2] = buffers[p, q][1:-1, 1:-1]


def laplace_surface(burned: np.ndarray, cell_m: float, iterations: int = 400) -> np.ndarray:
    """B04 `build_laplace` 와 같은 값 — 거리 비례 초기값 + 라플라스 다듬기."""
    surface = distance_surface(burned, cell_m)
    relax_checker(surface, np.isfinite(burned), iterations)
    return surface