"""D 노선분석 지도 0차곡 겹 — 지형 추정 골짜기 머리(PLAN 584-1 · 608-2 · 요구 443 · 486).

0차곡 = 뚜렷한 유로가 없는 곡두의 집수지형(`노선_배치.md:31` · 교본 `3. 임도망의 배치 절차.md:156`)
→ 사업 DEM(평활 안 한 `models_plan/dtm_sheet_<방법>.npz`)을 2 m 로 솎아 D8 흐름누적 ·
위 집수 넓이가 `min_m2` 이상 · `max_m2` 미만인 칸 = 0차곡(그보다 크면 1차 계류) →
세선화 · 선 뽑기(B05 `_mask_to_polylines`) · 경위도 선 + 가닥 · 길이 합.
임계값 = 지식DB 수치 없음 → 정함(브레인) 0.1 ha · 1 ha · 바꿈 = 요청 `min_m2` · `max_m2`.
파생 저장본(열쇠 = 격자 지문 · 임계 · 칸 · 쓴 도엽 누적 지문).
흐름누적 = 공용 지형 창고 도엽마다(도엽 + 여유 200 m · 창고에 있는 이웃 도엽만) 한 번 셈 →
`storage/_shared/terrain/<도엽>/zero_acc_<epsg>_<칸>m.npz`(열쇠 = 쓴 도엽 지형 열쇠) ·
임계 바꿈 = 저장 누적에서 칸 고르기 · 선 뽑기만 · 사업 DEM 이 창고 격자가 아니면(옛 받기) 통째 셈.

계산 자리 = ③ 서버 단독.
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any

import numpy as np

from common_util import common_util_terrain_store as terrain_store
from common_util.common_util_atomic import atomic_write_npz
from common_util.common_util_derived_store import (
    derived_key,
    file_stamp,
    read_derived,
    write_derived,
)
from common_util.common_util_json import atomic_write_json
from config.config_system_terrain import SHEET_SURFACE_DEFAULT_METHOD

ZERO_ORDER_KIND = "plan_zero_order"
ZERO_ORDER_VERSION = 2
ZERO_ORDER_MIN_M2 = 1_000.0  # 집수 시작(0.1 ha)
ZERO_ORDER_MAX_M2 = 10_000.0  # 이 넓이부터 1차 계류(1 ha)
ZERO_ORDER_CELL_M = 2.0  # 흐름누적 칸(1 m 격자는 솎음)
ZERO_ORDER_MIN_LEN_M = 10.0  # 세선화 갈림목 토막 거름
ZERO_ORDER_MARGIN_M = 200.0  # 도엽 누적 여유(통째 셈과 0차곡 칸 차 0 · 분석 486 3-1)
ACC_CAP = 65535  # 저장 누적 상한(칸 수 · 2 m 면 26 ha · 임계 1 ha 보다 넉넉)
SIMPLIFY_M = 1.0

_D8 = ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1))


def _shift(grid: np.ndarray, dr: int, dc: int, fill: Any) -> np.ndarray:
    """칸마다 (dr, dc) 이웃 값 — 격자 밖 = `fill`."""
    rows, cols = grid.shape
    pad = np.pad(grid, 1, constant_values=fill)
    return pad[1 + dr : 1 + dr + rows, 1 + dc : 1 + dc + cols]


def d8_accumulation(z: np.ndarray, valid: np.ndarray) -> np.ndarray:
    """D8 흐름누적(칸 수 · 자기 칸 포함 · 빈 칸 0).

    웅덩이 메움(형태 재구성 · 가장자리 · 빈 칸 곁 = 출구) → 가장 가파른 아래 이웃 하나로 흐름 →
    평지(등고선 칸 · 메운 웅덩이)는 이미 흐르는 같은 높이 이웃으로 이어 감 → 위에서 아래로 더함.
    """
    from skimage.morphology import reconstruction

    rows, cols = z.shape
    height = np.where(valid, np.asarray(z, float), np.nan)
    floor = np.nanmin(height) if valid.any() else 0.0
    edge = valid & ~np.all([_shift(valid, dr, dc, False) for dr, dc in _D8], axis=0)
    ground = np.where(valid, height, floor)
    seed = np.where(edge | ~valid, ground, ground.max())
    filled = np.where(valid, reconstruction(seed, ground, method="erosion"), np.nan)
    index = np.arange(rows * cols).reshape(rows, cols)
    best = np.zeros(z.shape)
    receiver = np.full(z.shape, -1, np.int64)
    with np.errstate(invalid="ignore"):
        for dr, dc in _D8:
            drop = (filled - _shift(filled, dr, dc, np.nan)) / math.hypot(dr, dc)
            better = drop > best
            best = np.where(better, drop, best)
            receiver = np.where(better, index + dr * cols + dc, receiver)
        done = (receiver >= 0) | edge | ~valid
        while True:
            grew = False
            for dr, dc in _D8:
                take = ~done & _shift(done & valid, dr, dc, False)
                take &= _shift(filled, dr, dc, np.nan) == filled
                if take.any():
                    receiver[take] = index[take] + dr * cols + dc
                    done |= take
                    grew = True
            if not grew:
                break
    down = receiver.ravel().tolist()
    acc = valid.astype(float).ravel().tolist()
    upstream = np.bincount(receiver[receiver >= 0], minlength=rows * cols).tolist()
    stack = [i for i, n in enumerate(upstream) if n == 0]
    while stack:
        i = stack.pop()
        j = down[i]
        if j >= 0:
            acc[j] += acc[i]
            upstream[j] -= 1
            if upstream[j] == 0:
                stack.append(j)
    return np.asarray(acc).reshape(rows, cols)


def zero_order_mask(
    acc: np.ndarray, valid: np.ndarray, cell_m: float, min_m2: float, max_m2: float
) -> np.ndarray:
    """위 집수 넓이(칸 수 × 칸 넓이)가 [min_m2, max_m2) 인 칸."""
    area = acc * cell_m * cell_m
    return valid & (area >= min_m2) & (area < max_m2)


def mask_lines(
    acc: np.ndarray,
    z: np.ndarray,
    valid: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    cell_m: float,
    min_m2: float = ZERO_ORDER_MIN_M2,
    max_m2: float = ZERO_ORDER_MAX_M2,
) -> list[list[tuple[float, float]]]:
    """흐름누적 → 0차곡 칸 → 선(작업 좌표계 m 점 목록 · `ZERO_ORDER_MIN_LEN_M` 미만 거름)."""
    from B05_Profile.B05_Profile_Engine_Skeleton import _mask_to_polylines

    mask = zero_order_mask(acc, valid, cell_m, min_m2, max_m2)
    lines = [[(p[0], p[1]) for p in line["polyline"]] for line in _mask_to_polylines(mask, x, y, z)]
    return [
        line
        for line in lines
        if sum(math.dist(a, b) for a, b in zip(line, line[1:])) >= ZERO_ORDER_MIN_LEN_M
    ]


def zero_order_lines(
    z: np.ndarray,
    valid: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    cell_m: float,
    min_m2: float = ZERO_ORDER_MIN_M2,
    max_m2: float = ZERO_ORDER_MAX_M2,
) -> list[list[tuple[float, float]]]:
    """격자(행 = 남 → 북) 통째 흐름누적 → 0차곡 선."""
    return mask_lines(d8_accumulation(z, valid), z, valid, x, y, cell_m, min_m2, max_m2)


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


def _acc_names(sheet_no: str, epsg: int, cell_m: float) -> tuple[Path, Path, Path]:
    folder = terrain_store.TERRAIN_ROOT / str(sheet_no)
    stem = f"zero_acc_{int(epsg)}_{cell_m:g}m"
    return folder / f"{stem}.npz", folder / f"{stem}.key.json", folder / f".{stem}.lock"


def _dem_key(sheet_no: str, epsg: int, cell_m: float) -> dict[str, Any] | None:
    """창고 지형 열쇠(셈 시간 뺌) — 안 셌거나 원천이 바뀌었으면 None."""
    stored = terrain_store._stored_ok(sheet_no, epsg, cell_m)
    return None if stored is None else {k: v for k, v in stored.items() if k != "seconds"}


def sheet_accumulation(sheet_no: str, epsg: int, cell_m: float = ZERO_ORDER_CELL_M) -> Path | None:
    """도엽 하나의 흐름누적(창고) — 도엽 + 여유 200 m 지형(창고에 있는 도엽만 · 없는 이웃 =
    출구)으로 셈 · 도엽 칸만 저장(uint16 · `ACC_CAP` 로 자름) · 열쇠 = 쓴 도엽 지형 열쇠 ·
    지형이 없으면 None."""
    own = _dem_key(sheet_no, epsg, cell_m)
    dem_npz = terrain_store._names(sheet_no, epsg, cell_m)[0]
    if own is None or own.get("empty") or not dem_npz.is_file():
        return None
    with np.load(dem_npz) as data:
        ix0, iy0 = int(data["ix0"]), int(data["iy0"])
        rows, cols = data["z"].shape
    pad = math.ceil(ZERO_ORDER_MARGIN_M / cell_m)
    box = (
        (ix0 - pad) * cell_m,
        (iy0 - pad) * cell_m,
        (ix0 + cols + pad) * cell_m,
        (iy0 + rows + pad) * cell_m,
    )
    used = {}
    for no in terrain_store.sheets_for_box(box, epsg):
        dem = _dem_key(no, epsg, cell_m)
        if dem is not None and not dem.get("empty"):
            used[no] = dem
    key = {"version": ZERO_ORDER_VERSION, "margin_m": ZERO_ORDER_MARGIN_M, "dem": used}
    npz, key_path, lock_path = _acc_names(sheet_no, epsg, cell_m)
    if npz.is_file() and _read_json(key_path) == key:
        return npz
    with terrain_store._Lock(lock_path):
        if npz.is_file() and _read_json(key_path) == key:  # 기다리는 동안 남이 셌을 수 있음
            return npz
        grid = terrain_store.load_terrain(epsg, sheets=sorted(used), box=box, cell_m=cell_m)
        acc = d8_accumulation(grid["z"], grid["valid_mask"])[pad : pad + rows, pad : pad + cols]
        key_path.unlink(missing_ok=True)  # 쓰는 도중 끊기면 옛 열쇠가 새 파일을 덮지 않게
        atomic_write_npz(
            npz,
            acc=np.minimum(acc, ACC_CAP).astype(np.uint16),
            ix0=np.int64(ix0),
            iy0=np.int64(iy0),
            cell_m=np.float64(cell_m),
        )
        atomic_write_json(key_path, key)
    return npz


def store_accumulation_paths(
    x: np.ndarray, y: np.ndarray, epsg: int, cell_m: float
) -> list[Path] | None:
    """사업 DEM 격자(창고 격자 위)를 덮는 도엽 누적 — 창고 격자 아님 · 지형 안 센 도엽 = None."""
    fx, fy = x[0] / cell_m - 0.5, y[0] / cell_m - 0.5
    if cell_m != ZERO_ORDER_CELL_M or abs(fx - round(fx)) > 1e-6 or abs(fy - round(fy)) > 1e-6:
        return None
    half = cell_m / 2
    box = (x[0] - half, y[0] - half, x[-1] + half, y[-1] + half)
    paths: list[Path] = []
    for no in terrain_store.sheets_for_box(box, epsg):
        dem = _dem_key(no, epsg, cell_m)
        if dem is None:
            return None
        path = None if dem.get("empty") else sheet_accumulation(no, epsg, cell_m)
        if path is not None:
            paths.append(path)
    return paths


def paste_accumulation(
    paths: list[Path], x: np.ndarray, y: np.ndarray, cell_m: float
) -> np.ndarray:
    """도엽 누적을 사업 격자(y 오름차순)에 이어 붙임 — 없는 칸 0."""
    ix0, iy0 = round(x[0] / cell_m - 0.5), round(y[0] / cell_m - 0.5)
    acc = np.zeros((len(y), len(x)))
    for path in paths:
        with np.load(path) as data:
            px, py, part = int(data["ix0"]), int(data["iy0"]), data["acc"]
        r0, c0 = max(py, iy0), max(px, ix0)
        r1, c1 = min(py + part.shape[0], iy0 + len(y)), min(px + part.shape[1], ix0 + len(x))
        if r0 < r1 and c0 < c1:
            src = part[r0 - py : r1 - py, c0 - px : c1 - px]
            np.copyto(acc[r0 - iy0 : r1 - iy0, c0 - ix0 : c1 - ix0], src, where=src > 0)
    return acc


def zero_order_layer(
    project_root: Path,
    models_dir: Path,
    epsg: int,
    min_m2: float = ZERO_ORDER_MIN_M2,
    max_m2: float = ZERO_ORDER_MAX_M2,
) -> dict[str, Any]:
    """0차곡 겹 — 지문(격자 파일 · 임계 · 칸 · 도엽 누적) 같으면 저장본 · 경위도 선 + 가닥 ·
    길이 합 m."""
    from shapely.geometry import LineString, mapping
    from shapely.ops import transform

    from D03_PlanAnalysis.D03_PlanAnalysis_Layers import _transforms

    npz = models_dir / f"dtm_sheet_{SHEET_SURFACE_DEFAULT_METHOD}.npz"
    stamp = file_stamp(npz)
    if stamp is None:
        raise FileNotFoundError("사업 지형면이 없습니다 — 지형 받기 뒤 보입니다.")
    with np.load(npz) as dem:
        x, y = np.asarray(dem["x"], float), np.asarray(dem["y"], float)
        cell = float(dem["resolution"][0]) if "resolution" in dem.files else float(x[1] - x[0])
    step = max(1, round(ZERO_ORDER_CELL_M / cell))
    x, y, cell_m = x[::step], y[::step], cell * step
    paths = store_accumulation_paths(x, y, epsg, cell_m)
    key = derived_key(
        ZERO_ORDER_KIND,
        ZERO_ORDER_VERSION,
        npz.name,
        stamp,
        int(epsg),
        float(min_m2),
        float(max_m2),
        ZERO_ORDER_CELL_M,
        ZERO_ORDER_MIN_LEN_M,
        SIMPLIFY_M,
        None if paths is None else [f"{p.parent.name}:{file_stamp(p)}" for p in paths],
    )
    saved = read_derived(project_root, ZERO_ORDER_KIND, key)
    if saved is not None:
        return saved["value"]
    with np.load(npz) as dem:
        z = np.asarray(dem["z"], float)[::step, ::step]
        valid = np.asarray(dem["valid_mask"], bool)[::step, ::step]
    if paths is None:
        lines = zero_order_lines(z, valid, x, y, cell_m, min_m2, max_m2)
    else:
        acc = paste_accumulation(paths, x, y, cell_m)
        lines = mask_lines(acc, z, valid, x, y, cell_m, min_m2, max_m2)
    _, to_ll = _transforms(epsg)
    shapes = [LineString(line).simplify(SIMPLIFY_M) for line in lines]
    value = {
        "min_m2": float(min_m2),
        "max_m2": float(max_m2),
        "cell_m": cell_m,
        "count": len(shapes),
        "length_m": round(sum(s.length for s in shapes), 1),
        "collection": {
            "type": "FeatureCollection",
            "features": [
                {"type": "Feature", "geometry": mapping(transform(to_ll, s)), "properties": {}}
                for s in shapes
            ],
        },
    }
    write_derived(project_root, ZERO_ORDER_KIND, key, value)
    return value