"""342-3 배수유역 빠른 길 — 도로 닿음 판정 한꺼번에 · 큰 방향장 잘라 씀.

- 판정: 옛 한 칸씩 걷기(최외곽 출발 → 내부 보충 · 색 물려받기)와 칸마다 같음(적색 차이 0).
- 잘라 씀: 같은 격자면 그대로 · 셀 어긋남 · 범위 밖은 None · 안쪽은 제 격자로 지은 것과 같음.
실측(검증_sub3 · 적색 차이 0 · 60 → 약 15 s)은 bench_342_watershed.py.
"""

import math

import numpy as np

from B04_PreProcess.B04_PreProcess_Engine_Watershed_Descent import build_contour_descent
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Expand import _pad_spec, crop_descent
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Flow import (
    _reach_by_doubling,
    outermost_cells,
)
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import GridSpec


def _walk_reference(receiver, valid2d, absorbing, stream):
    """옛 classify_flow 판정(한 칸씩 걷기) 그대로."""
    valid = valid2d.reshape(-1)
    reaches = np.zeros(receiver.size, bool)
    state = np.zeros(receiver.size, np.int8)
    reaches[stream] = True
    state[stream] = 2
    starts = list(np.flatnonzero(outermost_cells(valid2d))) + list(np.flatnonzero(valid))
    for start in starts:
        if state[start] == 2:
            continue
        path, node = [], start
        while True:
            if state[node] == 2:
                verdict = bool(reaches[node])
                break
            if state[node] == 1:
                verdict = False
                break
            state[node] = 1
            path.append(node)
            if absorbing[node]:
                verdict = True
                break
            following = int(receiver[node])
            if following == node:
                verdict = False
                break
            if not valid[following] and state[following] != 2:
                verdict = False
                break
            node = following
        for visited in path:
            reaches[visited] = verdict
            state[visited] = 2
    return reaches


def _random_case(seed, rows=60, cols=70):
    rng = np.random.default_rng(seed)
    size = rows * cols
    potential = rng.random(size)
    index = np.arange(size)
    r, c = index // cols, index % cols
    receiver = index.copy()
    best = potential.copy()
    for dr in (-1, 0, 1):
        for dc in (-1, 0, 1):
            nr, nc = r + dr, c + dc
            ok = (nr >= 0) & (nr < rows) & (nc >= 0) & (nc < cols)
            n = np.where(ok, nr * cols + nc, index)
            better = ok & (potential[n] < best)
            best = np.where(better, potential[n], best)
            receiver = np.where(better, n, receiver)
    valid2d = rng.random((rows, cols)) > 0.15
    road = np.zeros(size, bool)
    road[(r == rows // 2) & (c > 5) & (c < cols - 5)] = True
    stream = (rng.random(size) > 0.97) & ~road  # 표고 없는 세류 셀도 섞임
    # 세류 셀은 물길 다음 셀로 새김 — 순환 하나를 일부러 만든다(걷던 판정 = 미도달).
    receiver = receiver.astype(np.int32)
    receiver[stream] = np.roll(index, 1)[stream]
    receiver[5 * cols + 5], receiver[5 * cols + 6] = 5 * cols + 6, 5 * cols + 5
    absorbing = (road & valid2d.reshape(-1)) | stream
    return receiver, valid2d, absorbing, stream


def test_reach_matches_walk_cell_by_cell():
    for seed in range(8):
        receiver, valid2d, absorbing, stream = _random_case(seed)
        expected = _walk_reference(receiver, valid2d, absorbing, stream)
        got = _reach_by_doubling(receiver, valid2d.reshape(-1), absorbing, stream)
        assert int((got != expected).sum()) == 0, seed
        assert expected.sum() > 0 and (~expected).sum() > 0


def _circle_contours(cx, cy, radii_levels):
    features = []
    for radius, level in radii_levels:
        coords = [
            [cx + radius * math.cos(t), cy + radius * math.sin(t)]
            for t in np.linspace(0, 2 * math.pi, 181)
        ]
        features.append(
            {
                "geometry": {"type": "LineString", "coordinates": coords},
                "properties": {"elevation": level},
            }
        )
    return features


def test_crop_descent_matches_own_grid_inside():
    contours = _circle_contours(100.0, -100.0, [(r, 200.0 - r) for r in range(10, 140, 10)])
    spec = GridSpec(x_min=40.0, y_max=-40.0, cell_m=1.0, n_rows=120, n_cols=120)
    big_spec, _ = _pad_spec(spec, np.ones((120, 120), bool), 30)
    big = build_contour_descent(big_spec, contours)
    own = build_contour_descent(spec, contours)
    cut = crop_descent(big, spec)
    assert cut is not None and cut.spec == spec
    inner = np.zeros((120, 120), bool)
    inner[20:100, 20:100] = True
    flat = inner.reshape(-1)
    assert np.array_equal(cut.valid[inner], own.valid[inner])
    assert np.array_equal(cut.receiver[flat], own.receiver[flat])
    assert np.array_equal(cut.azimuth[flat], own.azimuth[flat])
    # 잘린 몫 밖을 가리키는 셀은 없다(제자리로 둠).
    assert cut.receiver.min() >= 0 and cut.receiver.max() < spec.size
    assert crop_descent(big, big_spec) is big


def test_crop_descent_refuses_misaligned_or_outside():
    contours = _circle_contours(100.0, -100.0, [(r, 200.0 - r) for r in range(10, 140, 10)])
    spec = GridSpec(x_min=40.0, y_max=-40.0, cell_m=1.0, n_rows=50, n_cols=50)
    big = build_contour_descent(spec, contours)
    shifted = GridSpec(x_min=40.5, y_max=-40.0, cell_m=1.0, n_rows=20, n_cols=20)
    outside = GridSpec(x_min=30.0, y_max=-40.0, cell_m=1.0, n_rows=20, n_cols=20)
    assert crop_descent(big, shifted) is None
    assert crop_descent(big, outside) is None