"""D 0차곡 D8 평지 이어 감 고침(PLAN 617-1) — 옛 꼴(통격자 바퀴)과 새 꼴(후보 칸 목록) 비트 동일 ·
합성 격자(넓은 평지 · 계단 · 빈 칸 · 가장자리) + 창고 0차곡 도엽 다시 셈 = 저장본(읽기만).

./venv/Scripts/python.exe -m pytest resources/tester/test_617_1_d8_flat.py -q
"""

import json
import math

import numpy as np
import pytest

from common_util import common_util_terrain_store as store
from D03_PlanAnalysis import D03_PlanAnalysis_ZeroOrder as zo

_D8 = zo._D8
_shift = zo._shift


def _old_d8(z: np.ndarray, valid: np.ndarray) -> np.ndarray:
    """617-1 전 꼴 그대로(대조용)."""
    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 _grids():
    rng = np.random.default_rng(617)
    yy, xx = np.mgrid[0:60, 0:80]
    flat = np.full((60, 80), 5.0)  # 넓은 평지(통째)
    basin = flat.copy()
    basin[10:50, 10:70] = 3.0  # 메운 웅덩이 = 넓은 평지
    steps = np.floor(xx / 9.0) + np.floor(yy / 13.0)  # 계단(층마다 평지)
    noisy = np.round(rng.normal(0, 1, (60, 80)) + yy * 0.05, 0)  # 같은 높이 많음
    slope = 0.3 * xx + np.sin(yy / 5.0)
    out = []
    for z in (flat, basin, steps, noisy, slope):
        out.append((z, np.ones(z.shape, bool)))
        holes = rng.random(z.shape) > 0.1  # 빈 칸
        holes[:, :3] = False  # 빈 띠 가장자리
        out.append((z, holes))
    one = np.zeros((1, 1))
    out.append((one, np.ones((1, 1), bool)))
    out.append((np.zeros((3, 5)), np.zeros((3, 5), bool)))  # 모두 빈 칸
    return out


@pytest.mark.parametrize("z,valid", _grids())
def test_new_equals_old_synthetic(z, valid) -> None:
    assert np.array_equal(zo.d8_accumulation(z, valid), _old_d8(z, valid))


def _stored_tiles(limit: int = 10):
    """열쇠가 지금 지형 파일과 맞는 창고 0차곡 도엽(고루 솎음)."""
    found = sorted(store.TERRAIN_ROOT.glob("*/zero_acc_*_2m.key.json"))
    step = max(len(found) // (limit * 3), 1)
    for key_path in found[::step]:
        key = json.loads(key_path.read_text(encoding="utf-8"))
        if key.get("version") != zo.ZERO_ORDER_VERSION:
            continue
        epsg = int(key_path.name.split("_")[2])
        dems = {}
        for no, dem in key["dem"].items():
            npz, dem_key, _ = store._names(no, epsg, zo.ZERO_ORDER_CELL_M)
            try:
                now = json.loads(dem_key.read_text(encoding="utf-8"))
            except (OSError, ValueError):
                break
            if {k: v for k, v in now.items() if k != "seconds"} != dem or not npz.is_file():
                break
            dems[no] = npz
        else:
            yield key_path.parent.name, epsg, dems
            limit -= 1
            if limit == 0:
                return


def test_store_tiles_recompute_equal() -> None:
    """창고 도엽 누적을 새 꼴로 다시 셈 = 저장본(읽기만 · 창고 안 씀)."""
    tiles = list(_stored_tiles())
    if not tiles:
        pytest.skip("창고 0차곡 도엽 없음")
    cell = zo.ZERO_ORDER_CELL_M
    pad = math.ceil(zo.ZERO_ORDER_MARGIN_M / cell)
    for no, epsg, dems in tiles:
        with np.load(zo._acc_names(no, epsg, cell)[0]) as data:
            want = np.asarray(data["acc"])
            ix0, iy0 = int(data["ix0"]) - pad, int(data["iy0"]) - pad
        rows, cols = want.shape[0] + 2 * pad, want.shape[1] + 2 * pad
        z = np.full((rows, cols), np.nan, np.float32)
        for npz in dems.values():  # load_terrain 이음과 같음(지형 다시 셈 길 없이)
            with np.load(npz) as data:
                px, py, pz = int(data["ix0"]), int(data["iy0"]), data["z"]
            r0, c0 = max(py, iy0), max(px, ix0)
            r1, c1 = min(py + pz.shape[0], iy0 + rows), min(px + pz.shape[1], ix0 + cols)
            if r0 < r1 and c0 < c1:
                src = pz[r0 - py : r1 - py, c0 - px : c1 - px]
                np.copyto(z[r0 - iy0 : r1 - iy0, c0 - ix0 : c1 - ix0], src, where=np.isfinite(src))
        acc = zo.d8_accumulation(z, np.isfinite(z))[pad:-pad, pad:-pad]
        got = np.minimum(acc, zo.ACC_CAP).astype(np.uint16)
        assert np.array_equal(got, want), no