"""E04 3D 활동면 탐색 — Scoops3D 방식 구 격자 + 타원체 · 셀별 최소 안전율 · 임계 체적.

근거: Reid 외(2015) USGS TM 14-A1 3.2(탐색 격자 · 반지름 증분 · 부피/면적 제약 ·
거친→고운 탐색) — 방식만 · 코드 자체.
다른 점(정하고 보고):
- 노드 높이 = 노드 아래 지표 + h — 고저차 큰 산지에서 땅속 · 너무 높은 노드를 덜어 냄.
- 타원체 = 활동 단면(d−z)은 원 · 회전축 방향 반축만 ratio 배(Hungr 1989 꼴)
  — 수직력이 축을 지나 Bishop 그대로. 방향 = 노드 지표 낙하 방향(5×5 평활).
  구는 활동체 지표 평균 낙하 방향(Scoops3D 2.1).
- 한 구가 지표를 여러 덩이로 자르면 덩이마다 셈(ndimage.label · 8 이웃).
- 격자 가장자리 · DEM 빈칸에 닿은 덩이는 버림(잘린 질량 막음).

다스리는 안전율(governing):
- 못 뚫는 층(기본 발파암 = 연암)에 잘린 바닥 면적 비 ≥ 10 % → Janbu 보정으로 전환
  · 단 Bishop 이 더 낮으면 Bishop(작은 쪽). Hungr 1989 = 약층 복합면에서 Bishop 도
  쓸 만함 · Zhang 1988 약층 경우 Janbu 보정 1.70 > Bishop 1.62 라 전환만은 덜 보수적.
- 10 % 문턱 = 기둥 몇 개 닿아 방법이 바뀌는 지도 이음매를 막음(Zhang 약층 잘림 비 33 %).
- 그 밖 → Hungr 3D Bishop · Bishop 이 수렴 못 하면 그 후보는 버림(Scoops3D 2.3.2).
"""

from __future__ import annotations

import math
import time
from collections.abc import Callable
from dataclasses import dataclass, field

import numpy as np
from scipy import ndimage

from E04_Stab3D.E04_Stab3D_Engine_Column import Columns, Strata, solve

TRUNCATION_JANBU = 0.10
_STRUCT = np.zeros((3, 3, 3), bool)
_STRUCT[1] = True  # 같은 후보 안에서만 8 이웃 · 후보끼리는 안 이음
BATCH_CELLS = 3_000_000  # 한 번에 만드는 (후보 × 창) 칸 수 — 메모리 약 0.5 GB 아래


@dataclass
class SearchParams:
    """탐색 조건 — 길이 m · 부피 m³ · 면적 m²."""

    box: tuple[float, float, float, float] | None = None  # (x0, x1, y0, y1) · None = 격자 전체
    node_mask: np.ndarray | None = (
        None  # (ny, nx) — 노드 아래 셀이 False 면 노드 뺌(노선 띠만 등 · Scoops3D 파일 탐색)
    )
    spacing: float = 2.0  # 가장 고운 노드 수평 간격
    h_min: float = 2.0  # 노드 높이(지표 위) 범위 · 간격
    h_max: float = 30.0
    h_step: float = 2.0
    coarse: int = 4  # 거친 격자 배수(Scoops3D multres)
    dr: float = 1.0  # 반지름 증분
    max_depth: float = 15.0  # 반지름 상한 = h + max_depth
    v_min: float = 10.0
    v_max: float = 10_000.0
    a_min: float = 0.0
    a_max: float = math.inf
    ratios: tuple[float, ...] = (1.0, 2.0)  # 구(늘 봄) + 타원체 축 방향 배수(구 임계 노드에서)
    sub: int = 2  # 셀 안 표본 sub×sub(부분 기둥)
    refine_rounds: int = 4  # 가장 고운 간격에서 더 내려가 보는 횟수 상한
    # 고운 탐색 · 타원체는 F < max(이 값, 전역 최소 × 1.2) 로 고친 노드만(높은 곳 지도는 거친 값)
    refine_below: float = 2.0
    groundwater: dict | None = None
    seismic_kh: float = 0.0
    # 화면 고리(533-3) — 결과 불변 · on_progress(percent, message) · should_stop() 참이면 InterruptedError
    on_progress: Callable[[int, str], None] | None = None
    should_stop: Callable[[], bool] | None = None


@dataclass
class SearchResult:
    f_min: np.ndarray  # 셀별 최소 안전율(다스리는 값) · 닿지 않은 셀 = nan
    volume: np.ndarray  # 그 임계 활동체 부피
    method: np.ndarray  # 0 = Bishop · 1 = Janbu 보정
    best: dict  # 전체 최소 활동체(중심 · 반지름 · 비 · 방향 · 방법별 F · 부피)
    trials: int = 0
    evaluated: int = 0
    seconds: float = 0.0
    stats: dict = field(default_factory=dict)


def trial_columns(
    strata: Strata,
    centers: np.ndarray,
    r: float,
    ratio: float = 1.0,
    directions: np.ndarray | None = None,
    v_min: float = 0.0,
    v_max: float = math.inf,
    a_min: float = 0.0,
    a_max: float = math.inf,
    sub: int = 2,
) -> tuple[Columns, np.ndarray, np.ndarray]:
    """후보 T 개(같은 r · ratio)의 기둥 — (Columns, 덩이의 후보 번호, 후보별 가장 큰 덩이 부피).

    부분 기둥 — 셀마다 sub×sub 점에서 활동면 · 지표(쌍선형)를 보고 걸친 비율 · 평균 깊이로
    기둥을 만듦(Scoops3D 2.1 부분 기둥과 같은 뜻 · 가장자리 적분 오차 줄임). sub=1 = 셀 중심.
    """
    centers = np.atleast_2d(np.asarray(centers, float))
    n_t = len(centers)
    dx = strata.dx
    ny, nx = strata.z_top.shape
    w = int(math.ceil(r * max(ratio, 1.0) / dx)) + 1
    off = np.arange(-w, w + 1)
    ix0 = np.rint((centers[:, 0] - strata.x[0]) / dx).astype(int)
    iy0 = np.rint((centers[:, 1] - strata.y[0]) / dx).astype(int)
    ixg = ix0[:, None, None] + off[None, None, :]
    iyg = iy0[:, None, None] + off[None, :, None]
    inb = (ixg >= 0) & (ixg < nx) & (iyg >= 0) & (iyg < ny)
    ixc, iyc, inb = np.broadcast_arrays(np.clip(ixg, 0, nx - 1), np.clip(iyg, 0, ny - 1), inb)
    shape = ixc.shape
    d = None if ratio == 1.0 else np.asarray(directions, float)
    k2 = 1.0 / (ratio * ratio)

    def surface(t, xs, ys):
        """후보 t 의 활동면 — (rem, z_s, ∂z/∂x, ∂z/∂y). 원소별(납작 · 창 모두)."""
        dxp = xs - centers[t, 0]
        dyp = ys - centers[t, 1]
        if d is None:
            rem = r * r - dxp * dxp - dyp * dyp
            root = np.sqrt(np.maximum(rem, 1e-12))
            return rem, centers[t, 2] - root, dxp / root, dyp / root
        dxd, dyd = d[t, 0], d[t, 1]
        pp = dxp * dxd + dyp * dyd
        qq = -dxp * dyd + dyp * dxd  # 축 방향 a = (−d_y, d_x)
        rem = r * r - pp * pp - k2 * qq * qq
        root = np.sqrt(np.maximum(rem, 1e-12))
        gx = (pp * dxd - k2 * qq * dyd) / root
        gy = (pp * dyd + k2 * qq * dxd) / root
        return rem, centers[t, 2] - root, gx, gy

    # 1) 셀 중심으로 거르기 — 걸친 셀과 그 이웃만 셀 안 표본을 봄(창 대부분은 허공 · 땅속 밖)
    t_grid = np.broadcast_to(np.arange(n_t)[:, None, None], shape)
    xs0 = strata.x[0] + ixg * dx
    ys0 = strata.y[0] + iyg * dx
    rem0, zs0, _, _ = surface(t_grid, xs0, ys0)
    in0 = np.flatnonzero(inb & (rem0 > 0))
    yy0, xx0 = iyc.ravel()[in0], ixc.ravel()[in0]
    zb0 = np.maximum(zs0.ravel()[in0], strata.hard_top[yy0, xx0])
    act0 = np.zeros(shape, bool)
    act0.ravel()[in0[zb0 < strata.z_top[yy0, xx0]]] = True
    cand = act0 if sub == 1 else ndimage.binary_dilation(act0, structure=_STRUCT) & inb
    ci = np.flatnonzero(cand)
    big = np.zeros(n_t)
    e = np.zeros(0, int)
    empty = (Columns(e, e, e, np.zeros(0), np.zeros(0), np.zeros(0), np.zeros((0, 3))), e, big)
    if not len(ci):
        return empty
    tt = t_grid.ravel()[ci]
    cy = iyc.ravel()[ci]
    cx = ixc.ravel()[ci]
    hard = strata.hard_top[cy, cx]

    # 2) 부분 기둥 — sub×sub 점
    n_act = np.zeros(len(ci))
    acc = {k: np.zeros(len(ci)) for k in ("depth", "zs", "gx", "gy", "px", "py")}
    fr = ((np.arange(sub) + 0.5) / sub - 0.5) if sub > 1 else np.zeros(1)
    for fy in fr:
        for fx in fr:
            xs = strata.x[0] + (cx + fx) * dx
            ys = strata.y[0] + (cy + fy) * dx
            rem, z_s, gx, gy = surface(tt, xs, ys)
            zg = strata.ground_at(cy, cx, fy, fx)
            z_b = np.maximum(z_s, hard)
            act = (rem > 0) & (z_b < zg)
            n_act += act
            acc["depth"] += np.where(act, zg - z_b, 0.0)
            for key, val in (("zs", z_s), ("gx", gx), ("gy", gy), ("px", xs), ("py", ys)):
                acc[key] += np.where(act, val, 0.0)
    keep = n_act > 0
    ci, tt, cy, cx, n_act = ci[keep], tt[keep], cy[keep], cx[keep], n_act[keep]
    acc = {k: v[keep] for k, v in acc.items()}
    active = np.zeros(shape, bool)
    active.ravel()[ci] = True
    labels, n_lab = ndimage.label(active, structure=_STRUCT)
    if n_lab == 0:
        return empty
    lab = labels.ravel()[ci]
    a_h = dx * dx
    depth_cell = acc["depth"] / (sub * sub)
    vol = np.bincount(lab, weights=depth_cell, minlength=n_lab + 1) * a_h
    area = np.bincount(lab, weights=n_act / (sub * sub), minlength=n_lab + 1) * a_h
    edge = np.bincount(lab, weights=strata.edge[cy, cx].astype(float), minlength=n_lab + 1)
    trial_of = np.zeros(n_lab + 1, int)
    trial_of[lab] = tt
    np.maximum.at(big, trial_of[1:], vol[1:])
    ok = (vol >= v_min) & (vol <= v_max) & (area >= a_min) & (area <= a_max) & (edge == 0)
    ok[0] = False
    keep_lab = np.flatnonzero(ok)
    remap = np.full(n_lab + 1, -1)
    remap[keep_lab] = np.arange(len(keep_lab))
    sel = remap[lab] >= 0
    na_s = n_act[sel]
    mean = {k: acc[k][sel] / na_s for k in ("zs", "gx", "gy", "px", "py")}
    frac = na_s / (sub * sub)
    trial = trial_of[keep_lab]
    cols = Columns(
        iy=cy[sel],
        ix=cx[sel],
        group=remap[lab][sel],
        z_slip=mean["zs"],
        gx=mean["gx"],
        gy=mean["gy"],
        center=centers[trial],
        direction=None if d is None else d[trial],
        px=mean["px"],
        py=mean["py"],
        frac=frac,
        depth=depth_cell[sel] / frac,
    )
    return cols, trial, big


def governing(res: dict) -> tuple[np.ndarray, np.ndarray]:
    """다스리는 안전율 · 방법(0 Bishop · 1 Janbu 보정) — 머리 주석 규칙."""
    jc = np.where(np.isfinite(res["janbu_corrected"]), res["janbu_corrected"], np.inf)
    use_j = (res["truncated_ratio"] >= TRUNCATION_JANBU) & (jc < res["bishop"])
    f = np.where(use_j, jc, res["bishop"])
    f = np.where(np.isfinite(res["bishop"]), f, np.nan)
    return f, use_j.astype(np.int8)


def evaluate_surface(
    strata: Strata,
    center,
    r: float,
    ratio: float = 1.0,
    direction=None,
    groundwater: dict | None = None,
    seismic_kh: float = 0.0,
    sub: int = 2,
) -> dict:
    """활동면 하나(구 · 타원체) — 가장 큰 덩이의 방법별 안전율 · 부피."""
    dirs = None if direction is None else np.asarray(direction, float)[None]
    cols, _, _ = trial_columns(strata, np.asarray(center, float)[None], r, ratio, dirs, sub=sub)
    if not len(cols.group):
        raise ValueError("활동면이 지표를 자르지 않음")
    res = solve(strata, cols, groundwater, seismic_kh)
    k = int(np.argmax(res["volume"]))
    out = {key: (v[k] if isinstance(v, np.ndarray) else v) for key, v in res.items()}
    out["governing"], out["method"] = (float(x[k]) for x in governing(res))
    return out


def _local_fall(strata: Strata) -> tuple[np.ndarray, np.ndarray]:
    zs = ndimage.uniform_filter(np.nan_to_num(strata.z_top, nan=float(np.nanmean(strata.z_top))), 5)
    gy, gx = np.gradient(zs, strata.dx)
    norm = np.hypot(gx, gy)
    norm[norm == 0] = 1.0
    return -gx / norm, -gy / norm


def search(strata: Strata, params: SearchParams | None = None) -> SearchResult:
    """거친 격자 → 고운 격자(Scoops3D 3.2.2) 로 셀별 최소 안전율 · 임계 체적."""
    p = params or SearchParams()
    t0 = time.perf_counter()
    ny, nx = strata.z_top.shape
    x0, x1, y0, y1 = p.box or (strata.x[0], strata.x[-1], strata.y[0], strata.y[-1])
    nxs = int((x1 - x0) // p.spacing) + 1
    nys = int((y1 - y0) // p.spacing) + 1
    nhs = int(round((p.h_max - p.h_min) / p.h_step)) + 1
    fall_x, fall_y = _local_fall(strata)

    f_min = np.full((ny, nx), np.inf)
    volume = np.full((ny, nx), np.nan)
    method = np.full((ny, nx), -1, np.int8)
    node_map = np.full((ny, nx), -1, np.int64)
    best = {"governing": np.inf}
    seen: set[int] = set()
    counters = {"trials": 0, "evaluated": 0}
    stage = [0, 40, "거친 탐색"]  # 진행률 — 시작 % · 폭 % · 이름

    def tick(done: float) -> None:
        if p.should_stop and p.should_stop():
            raise InterruptedError
        if p.on_progress:
            p.on_progress(int(stage[0] + stage[1] * done), str(stage[2]))

    def node_xyz(ids: np.ndarray) -> np.ndarray:
        ix = ids % nxs
        iy = (ids // nxs) % nys
        ih = ids // (nxs * nys)
        x = x0 + ix * p.spacing
        y = y0 + iy * p.spacing
        cx = np.clip(np.rint((x - strata.x[0]) / strata.dx).astype(int), 0, nx - 1)
        cy = np.clip(np.rint((y - strata.y[0]) / strata.dx).astype(int), 0, ny - 1)
        z = strata.z_top[cy, cx] + p.h_min + ih * p.h_step
        return np.column_stack([x, y, z]), cy, cx, p.h_min + ih * p.h_step

    def run(ids: np.ndarray, ratios: tuple[float, ...] = (1.0,)) -> set[int]:
        """노드들의 모든 반지름 · 비 — 셀 최소를 고친 노드 번호를 돌려줌."""
        improved: set[int] = set()
        xyz0, cy0, cx0, _ = node_xyz(ids)
        keep = np.isfinite(xyz0[:, 2])
        if p.node_mask is not None:
            keep &= p.node_mask[cy0, cx0]
        ids = ids[keep]
        if not len(ids):
            return improved
        xyz, cy, cx, h = node_xyz(ids)
        dirs = np.column_stack([fall_x[cy, cx], fall_y[cy, cx]])
        r_lo = np.maximum(p.dr, np.floor(0.5 * h / p.dr) * p.dr)
        r_hi = h + p.max_depth
        for k_ratio, ratio in enumerate(ratios):
            done = np.zeros(len(ids), bool)
            radii = np.arange(r_lo.min(), r_hi.max() + 1e-9, p.dr)
            for k_r, r in enumerate(radii):
                tick((k_ratio + k_r / len(radii)) / len(ratios))
                live = np.flatnonzero(~done & (r_lo <= r + 1e-9) & (r <= r_hi + 1e-9))
                if not len(live):
                    continue
                w_now = (2 * math.ceil(r * max(ratio, 1) / strata.dx) + 3) ** 2
                chunk = max(1, BATCH_CELLS // max(w_now, 1))
                for s in range(0, len(live), chunk):
                    part = live[s : s + chunk]
                    d_part = dirs[part] if ratio != 1.0 else None
                    cols, trial, big = trial_columns(
                        strata,
                        xyz[part],
                        r,
                        ratio,
                        d_part,
                        p.v_min,
                        p.v_max,
                        p.a_min,
                        p.a_max,
                        p.sub,
                    )
                    done[part[big > p.v_max]] = True
                    counters["trials"] += len(part)
                    if not len(trial):
                        continue
                    res = solve(strata, cols, p.groundwater, p.seismic_kh)
                    counters["evaluated"] += len(trial)
                    f, meth = governing(res)
                    improved |= _update(
                        f,
                        meth,
                        res,
                        cols,
                        trial,
                        ids[part],
                        f_min,
                        volume,
                        method,
                        node_map,
                        max(p.refine_below, 1.2 * best["governing"]),
                    )
                    k = (
                        int(np.nanargmin(np.where(np.isfinite(f), f, np.inf)))
                        if np.isfinite(f).any()
                        else -1
                    )
                    if k >= 0 and f[k] < best["governing"]:
                        best.clear()
                        best.update(
                            {
                                key: (v[k].tolist() if isinstance(v, np.ndarray) else v)
                                for key, v in res.items()
                            }
                        )
                        best.update(
                            governing=float(f[k]),
                            method=int(meth[k]),
                            center=xyz[part][trial[k]].tolist(),
                            radius=float(r),
                            ratio=float(ratio),
                        )
        return improved

    step = max(1, int(p.coarse))
    grid = np.stack(
        np.meshgrid(
            np.arange(0, nhs, step), np.arange(0, nys, step), np.arange(0, nxs, step), indexing="ij"
        ),
        -1,
    ).reshape(-1, 3)
    ids = (grid[:, 0] * nys + grid[:, 1]) * nxs + grid[:, 2]
    seen.update(ids.tolist())
    seeds = run(ids)
    rounds = passes = 0
    while seeds and (step > 1 or rounds < p.refine_rounds):
        if step == 1:
            rounds += 1
        step = max(1, step // 2)
        passes += 1
        stage[:] = [min(40 + 8 * (passes - 1), 77), 8, f"고운 탐색 {passes}"]
        cand = set()
        for nid in seeds:
            ih, rem = divmod(nid, nys * nxs)
            iy, ix = divmod(rem, nxs)
            for dh in (-step, 0, step):
                for dy in (-step, 0, step):
                    for dxs in (-step, 0, step):
                        h2, y2, x2 = ih + dh, iy + dy, ix + dxs
                        if 0 <= h2 < nhs and 0 <= y2 < nys and 0 <= x2 < nxs:
                            cid = (h2 * nys + y2) * nxs + x2
                            if cid not in seen:
                                cand.add(cid)
        if not cand:
            if step == 1:
                break
            continue
        seen.update(cand)
        seeds = run(np.fromiter(cand, np.int64))

    # 타원체 — 구 탐색이 찾은 임계 노드(F < refine_below)에서만(시간 1/3 · 노선 1 km 실측)
    ell = tuple(r for r in p.ratios if r != 1.0)
    limit = max(p.refine_below, 1.2 * best["governing"])
    crit = np.unique(node_map[(node_map >= 0) & (f_min < limit)])
    if ell and len(crit):
        stage[:] = [85, 14, "타원체"]
        run(crit, ell)

    f_out = np.where(np.isfinite(f_min), f_min, np.nan)
    return SearchResult(
        f_min=f_out,
        volume=volume,
        method=method,
        best=best,
        trials=counters["trials"],
        evaluated=counters["evaluated"],
        seconds=time.perf_counter() - t0,
        stats={"nodes": len(seen), "cells": int(np.isfinite(f_out).sum())},
    )


def _update(
    f, meth, res, cols, trial, node_ids, f_min, volume, method, node_map, refine_below=np.inf
) -> set[int]:
    """셀마다 더 낮은 F 면 바꿈 — 같은 셀 여럿이면 가장 낮은 것(내림차순으로 써 마지막이 이김)."""
    fc = f[cols.group]
    ok = np.isfinite(fc)
    iy, ix, grp, fc = cols.iy[ok], cols.ix[ok], cols.group[ok], fc[ok]
    order = np.argsort(-fc, kind="stable")
    iy, ix, grp, fc = iy[order], ix[order], grp[order], fc[order]
    better = fc < f_min[iy, ix] * (1 - 1e-6)
    iy, ix, grp, fc = iy[better], ix[better], grp[better], fc[better]
    f_min[iy, ix] = fc
    volume[iy, ix] = res["volume"][grp]
    method[iy, ix] = meth[grp]
    nid = node_ids[trial[grp]]
    node_map[iy, ix] = nid
    return set(np.unique(nid[fc < refine_below]).tolist())