Aislo/E07_StabArea/E07_StabArea_Sheet.py

"""E07 광역 — 도엽마다 경사 · 집수면적 · 기본 물성 안정 지수를 공용 창고에 미리
(PLAN 610-1 · 요구 487 · 488).

지형 원천 = 공용 지형 창고(`common_util_terrain_store` · 전국 등고선 2 m) — LAS DTM 복사본 안 씀.
자리 = 창고 도엽 폴더 `storage/_shared/terrain/<도엽>/`
  e07_<epsg>_<격자>m.npz     경사 β(rad) · 단위 집수면적 a(m) — 도엽 지형 + 여유 400 m 로 셈 →
                            도엽 안쪽 칸만 저장 → 이어 붙여도 위쪽 유역 안 잘림(분석 486 3-1)
  e07si_<epsg>_<격자>m.npz   기본 물성(마스터 E02 붕적토) 안정 지수 — 물성이 다르면 부르는 쪽이
                            저장된 β · a 로 칸 계산만(흐름 다시 안 셈)
열쇠 = 판 + 여유 + 쓴 도엽 지형 열쇠(전국 파일 지문 · 방식 판) · SI 는 + 물성.
격자 = 창고와 같음(칸 중심 (k + ½)·cell · 행 0 = 남쪽).
"""

from __future__ import annotations

import json
import math
from dataclasses import asdict
from pathlib import Path
from typing import Any, Callable

import numpy as np

from common_util import common_util_terrain_store as store
from common_util.common_util_atomic import atomic_write_npz
from common_util.common_util_json import atomic_write_json
from E07_StabArea import E07_StabArea_Engine as engine

SHEET_VER = 1
MARGIN_M = 400.0

Report = Callable[[int, str], None]
Box = tuple[float, float, float, float]


def _paths(sheet_no: str, epsg: int, cell_m: float) -> dict[str, Path]:
    folder = store.TERRAIN_ROOT / str(sheet_no)
    stem = f"{int(epsg)}_{cell_m:g}m"
    return {
        "flow": folder / f"e07_{stem}.npz",
        "key": folder / f"e07_{stem}.key.json",
        "si": folder / f"e07si_{stem}.npz",
        "lock": folder / f".e07_{stem}.lock",
    }


def _dem_stamp(sheet_no: str, epsg: int, cell_m: float) -> Any:
    """도엽 지형 열쇠(셈 시간 뺌) — 지형이 다시 셈되면 달라짐."""
    path = store.TERRAIN_ROOT / str(sheet_no) / f"dem_{int(epsg)}_{cell_m:g}m.key.json"
    try:
        key = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    key.pop("seconds", None)
    return key


def _read_key(path: Path) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None


def ensure_sheet(
    sheet_no: str,
    epsg: int,
    cell_m: float,
    defaults: engine.AreaParams,
    report: Report | None = None,
) -> Path | None:
    """도엽 하나의 E07 흐름(β · a) + 기본 물성 SI — 있으면 그대로 · 없거나 지형이 바뀌었으면 셈.

    빈 도엽(바다) = None · 지형 받기 실패 = RuntimeError(저장 안 함 · 다음에 다시)."""

    def say(pct: int, stage: str) -> None:
        if report is not None:
            report(pct, stage)

    sheet_no = str(sheet_no)
    paths = _paths(sheet_no, epsg, cell_m)
    say(0, f"지형 · {sheet_no}")
    dem = store.ensure_sheet_terrain(sheet_no, epsg, cell_m, lambda p, s: say(p * 30 // 100, s))
    if dem is None:
        return None
    with np.load(dem) as data:
        ix0, iy0 = int(data["ix0"]), int(data["iy0"])
        core_ok = np.isfinite(data["z"])
    ny, nx = core_ok.shape
    pad = math.ceil(MARGIN_M / cell_m)
    box = (
        (ix0 - pad) * cell_m,
        (iy0 - pad) * cell_m,
        (ix0 + nx + pad) * cell_m,
        (iy0 + ny + pad) * cell_m,
    )
    around = store.sheets_for_box(box, epsg)
    stamp: dict[str, Any] = {}
    for i, no in enumerate(around):  # 이웃 지형(처음이면 셈 · 있으면 열쇠만 봄)
        say(30 + 40 * i // len(around), f"이웃 지형 {i + 1}/{len(around)} · {no}")
        store.ensure_sheet_terrain(no, epsg, cell_m)
        stamp[no] = _dem_stamp(no, epsg, cell_m)
    key = {"ver": SHEET_VER, "margin_m": MARGIN_M, "dem": stamp}
    if not (_read_key(paths["key"]) == key and paths["flow"].is_file()):
        with store._Lock(paths["lock"]):
            if not (_read_key(paths["key"]) == key and paths["flow"].is_file()):
                say(70, f"경사 · 집수면적 · {sheet_no}")
                region = store.load_terrain(epsg, sheets=around, box=box, cell_m=cell_m)
                if region["failed"]:
                    raise RuntimeError(f"도엽 {sheet_no} 이웃 지형 실패: {region['failed']}")
                z = region["z"]
                inner = (slice(pad, pad + ny), slice(pad, pad + nx))
                beta = engine.slope_rad(z, cell_m)[inner]
                area = engine.catchment_area(z, cell_m)[inner]
                beta[~core_ok] = np.nan
                area[~core_ok] = np.nan
                paths["key"].unlink(missing_ok=True)
                atomic_write_npz(
                    paths["flow"],
                    slope=beta.astype(np.float32),
                    a=area.astype(np.float32),
                    ix0=np.int64(ix0),
                    iy0=np.int64(iy0),
                )
                atomic_write_json(paths["key"], key)
    si_key = json.dumps({"flow": key, "params": asdict(defaults)}, sort_keys=True)
    if _si_key(paths["si"]) != si_key:
        say(90, f"기본 물성 안정 지수 · {sheet_no}")
        with np.load(paths["flow"]) as data:
            si = engine.stability_index(
                data["slope"].astype(float), data["a"].astype(float), defaults
            )
        atomic_write_npz(paths["si"], si=si.astype(np.float32), key=np.array(si_key))
    say(100, f"완료 · {sheet_no}")
    return paths["flow"]


def _si_key(path: Path) -> str | None:
    try:
        with np.load(path) as data:
            return str(data["key"])
    except (OSError, KeyError, ValueError):
        return None


def load_area(
    epsg: int,
    box: Box,
    cell_m: float,
    defaults: engine.AreaParams,
    report: Report | None = None,
) -> dict[str, Any]:
    """TM 상자 → 이어 붙인 {x0 · y0(첫 칸 중심) · slope · a · si(기본 물성)} · 없는 칸 = NaN."""
    sheets = store.sheets_for_box(box, epsg)
    ix0, iy0 = math.ceil(box[0] / cell_m - 0.5), math.ceil(box[1] / cell_m - 0.5)
    ix1, iy1 = math.floor(box[2] / cell_m - 0.5) + 1, math.floor(box[3] / cell_m - 0.5) + 1
    shape = (max(iy1 - iy0, 0), max(ix1 - ix0, 0))
    out = {k: np.full(shape, np.nan) for k in ("slope", "a", "si")}
    total = max(len(sheets), 1)
    for i, no in enumerate(sheets):

        def inner(pct: int, stage: str, i: int = i) -> None:
            if report is not None:
                report((i * 100 + pct) // total, f"{stage} ({i + 1}/{total})")

        path = ensure_sheet(no, epsg, cell_m, defaults, inner)
        if path is None:
            continue
        with np.load(path) as data:
            px, py = int(data["ix0"]), int(data["iy0"])
            pieces = {"slope": data["slope"], "a": data["a"]}
        with np.load(_paths(no, epsg, cell_m)["si"]) as data:
            pieces["si"] = data["si"]
        h, w = pieces["slope"].shape
        r0, c0, r1, c1 = max(py, iy0), max(px, ix0), min(py + h, iy1), min(px + w, ix1)
        if r0 >= r1 or c0 >= c1:
            continue
        for k, v in pieces.items():
            src = v[r0 - py : r1 - py, c0 - px : c1 - px]
            np.copyto(
                out[k][r0 - iy0 : r1 - iy0, c0 - ix0 : c1 - ix0], src, where=~np.isnan(src)
            )  # SI ∞ = 평지
    return {"x0": (ix0 + 0.5) * cell_m, "y0": (iy0 + 0.5) * cell_m, "sheets": sheets, **out}