"""E04 3D 기둥 솔버 — Hovland · Hungr 3D Bishop 간편 · Hungr 3D Janbu 간편 + 보정.

근거(수식만 · 코드는 자체):
- Reid 외(2015) USGS TM 14-A1 Scoops3D 2.1~2.5 — 기둥 기하 · 모멘트 평형 · 간극수압 · 특수 경우.
- Hungr(1987) Géotechnique 37(1) · Hungr·Salgado·Byrne(1989) CGJ 26(4) — 3D Bishop · 3D Janbu.
- Hovland(1977) ASCE GT9 — 기둥 사이 힘 0 · N = W cosψ.
- Janbu(1973) 보정 계수 f0 = 1 + b1 (d/L − 1.4 (d/L)²).
- 지식DB `original/안정성평가/05_3D_사면안정해석_논문_기술정보.md` 1장.

꼴: 기둥은 납작한 배열 + `group`(활동체 번호) — 여러 활동체를 한 번에 셈(np.bincount).
좌표: 활동 방향 d(수평 단위) · 회전축 = 중심을 지나는 수평선 ⊥ d.
모멘트(내리미는 쪽 +) — 자중 −p·W · 지진 e·kh·W · 수직력 N·m_n · 저항 전단 −S·arm_s.
  p = (P − C)·d · e = 축 높이 − 기둥 무게중심 높이 · g = 바닥면의 d 방향 기울기
  arm_s = (p·g − z_b')/√(1+g²) · m_n = cosψ (p + z_b'·g)
  구 · 원형 단면 타원체는 수직력이 축을 지나 m_n = 0 → Scoops3D 식 2.22 그대로.
Bishop 수직력(식 2.31): N = (W+Q − (cA − uA tanφ) sinα / F) / mα · mα = cosψ + sinα tanφ / F
  → 저항 S_res = [c A_h + (W+Q − u A_h) tanφ] / mα (식 2.33 · A_h = 수평 면적).
Janbu 힘 평형(d 방향): F0 = Σ S_res / cosα ÷ Σ((W+Q) tanα + kh W) · 보정 F = f0 F0.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
from scipy import ndimage

from E03_StabSection.E03_StabSection_Types import GAMMA_W, Groundwater, Material, groundwater

MA_MIN = 0.2  # mα 하한 경고(Duncan & Wright) — 밑이면 결과에 표시
F_CAP = 100.0  # Scoops3D 2.5 — 너무 큰 F 는 100
TOL = 1e-4
MAX_ITER = 50


@dataclass
class Strata:
    """지층 격자 — 셀 중심 축 x(nx) · y(ny) · 지표 z_top(ny, nx) · 층 바닥(K−1, ny, nx).

    층 k 는 위(층 k−1 바닥 · 0 층은 지표)부터 bottoms[k] 까지 · 맨 아래 층은 끝없음.
    hard[k] = 활동면이 못 뚫는 층(암반) — 그 윗면에서 잘림(Search 머리 주석 규칙).
    z_water = 지하수위면(mode "line" 의 3D 꼴) · q = 상재하중(kPa) 격자.
    """

    x: np.ndarray
    y: np.ndarray
    z_top: np.ndarray
    bottoms: np.ndarray
    materials: list[Material]
    hard: list[bool] = field(default_factory=list)
    z_water: np.ndarray | None = None
    q: np.ndarray | None = None

    def __post_init__(self) -> None:
        self.x = np.asarray(self.x, float)
        self.y = np.asarray(self.y, float)
        self.z_top = np.asarray(self.z_top, float)
        self.bottoms = np.asarray(self.bottoms, float).reshape(-1, *self.z_top.shape)
        if len(self.materials) != len(self.bottoms) + 1:
            raise ValueError("물성 수 = 층 바닥 수 + 1 이어야 함")
        self.hard = list(self.hard) or [False] * len(self.materials)
        self.dx = float(self.x[1] - self.x[0])
        if abs(float(self.y[1] - self.y[0]) - self.dx) > 1e-4 * self.dx:
            raise ValueError("격자 간격은 x · y 같아야 함")
        # 층 바닥은 지표 아래 · 위에서 아래로 내려감
        tops = np.concatenate([self.z_top[None], self.bottoms])
        self.bottoms = np.minimum.accumulate(tops, axis=0)[1:]
        self.gamma_t = np.array([m["gamma_t"] for m in self.materials])
        self.gamma_sat = np.array([m["gamma_sat"] for m in self.materials])
        self.c = np.array([m["c"] for m in self.materials])
        self.tan_phi = np.tan(np.radians([m["phi"] for m in self.materials]))
        # 못 뚫는 층 윗면(없으면 −inf) · 그 기울기
        tops = np.concatenate([self.z_top[None], self.bottoms])
        first = next((k for k, h in enumerate(self.hard) if h), None)
        hard_top = np.full(self.z_top.shape, -np.inf) if first is None else tops[first].copy()
        self.hard_top = hard_top
        finite = np.where(np.isfinite(hard_top), hard_top, 0.0)
        self.hard_gy, self.hard_gx = np.gradient(finite, self.dx)
        self.gy, self.gx = np.gradient(self.z_top, self.dx)
        # 가장자리 · 빈칸 이웃 — 여기 닿은 활동체는 질량이 잘려 버림(Search)
        edge = ndimage.binary_dilation(~np.isfinite(self.z_top))
        edge[[0, -1], :] = True
        edge[:, [0, -1]] = True
        self.edge = edge
        self._z_pad = np.pad(self.z_top, 1, mode="edge")
        self._shifted: dict = {}

    def ground_at(self, iy: np.ndarray, ix: np.ndarray, fy: float, fx: float) -> np.ndarray:
        """셀 (iy, ix) 중심에서 (fy, fx) 칸만큼 옮긴 점의 지표(쌍선형 · |f| ≤ 0.5)."""
        if fx == 0 and fy == 0:
            return self.z_top[iy, ix]
        key = (fy, fx)
        if key not in self._shifted:  # 옮긴 점 지표 격자를 한 번 만들어 둠
            sx, sy = int(np.sign(fx)), int(np.sign(fy))
            ax, ay = abs(fx), abs(fy)
            zp = self._z_pad
            ny, nx = self.z_top.shape
            z00 = zp[1 : ny + 1, 1 : nx + 1]
            z01 = zp[1 : ny + 1, 1 + sx : nx + 1 + sx]
            z10 = zp[1 + sy : ny + 1 + sy, 1 : nx + 1]
            z11 = zp[1 + sy : ny + 1 + sy, 1 + sx : nx + 1 + sx]
            self._shifted[key] = (1 - ay) * ((1 - ax) * z00 + ax * z01) + ay * (
                (1 - ax) * z10 + ax * z11
            )
        return self._shifted[key][iy, ix]


@dataclass
class Columns:
    """활동체 묶음의 기둥 — 납작한 배열(길이 n) · group 은 0..G−1."""

    iy: np.ndarray
    ix: np.ndarray
    group: np.ndarray
    z_slip: np.ndarray  # 시험 활동면 높이(잘리기 전)
    gx: np.ndarray  # 활동면 기울기 ∂z/∂x
    gy: np.ndarray  # ∂z/∂y
    center: np.ndarray  # (G, 3) 회전축이 지나는 점
    direction: np.ndarray | None = None  # (G, 2) 활동 방향 · None 이면 지표 낙하 방향(Scoops3D)
    # 부분 기둥(셀 안 표본) — 대표 위치 · 걸친 비율 · 걸친 부분 평균 깊이
    # None = 셀 중심 · 한 칸 · 지표 − 바닥
    px: np.ndarray | None = None
    py: np.ndarray | None = None
    frac: np.ndarray | None = None
    depth: np.ndarray | None = None


def _gsum(group: np.ndarray, values: np.ndarray, n: int) -> np.ndarray:
    return np.bincount(group, weights=values, minlength=n)


def fall_direction(strata: Strata, cols: Columns, n: int) -> np.ndarray:
    """활동체 지표의 평균 낙하 방향(Scoops3D 2.1) — (n, 2) 단위 벡터."""
    sx = _gsum(cols.group, -strata.gx[cols.iy, cols.ix], n)
    sy = _gsum(cols.group, -strata.gy[cols.iy, cols.ix], n)
    norm = np.hypot(sx, sy)
    norm[norm == 0] = 1.0
    return np.column_stack([sx / norm, sy / norm])


def _column_state(strata: Strata, cols: Columns, gw: Groundwater, kh: float) -> dict:
    """기둥마다 자중 · 바닥 물성 · 간극수압 · 기하 — 활동면을 못 뚫는 층에서 자름."""
    iy, ix, grp = cols.iy, cols.ix, cols.group
    hard_top = strata.hard_top[iy, ix]
    truncated = cols.z_slip < hard_top
    z_b = np.where(truncated, hard_top, cols.z_slip)
    gx = np.where(truncated, strata.hard_gx[iy, ix], cols.gx)
    gy = np.where(truncated, strata.hard_gy[iy, ix], cols.gy)
    z_top = strata.z_top[iy, ix] if cols.depth is None else z_b + cols.depth
    depth = np.maximum(z_top - z_b, 0.0)
    x = strata.x[ix] if cols.px is None else cols.px
    y = strata.y[iy] if cols.py is None else cols.py

    # 층별 두께 → 자중(지하수위 아래는 γsat) · 부분 기둥은 걸친 비율만큼 면적
    a_h = strata.dx**2 * (np.ones(len(iy)) if cols.frac is None else cols.frac)
    k_n = len(strata.materials)
    tops = np.concatenate([z_top[None], strata.bottoms[:, iy, ix]])
    bots = np.concatenate([strata.bottoms[:, iy, ix], np.full((1, len(iy)), -np.inf)])
    upper = np.minimum(tops, z_top)
    lower = np.maximum(bots, z_b)
    thick = np.clip(upper - lower, 0.0, None)
    if gw.get("mode") == "line" and strata.z_water is not None:
        zw = strata.z_water[iy, ix]
        wet = np.clip(np.minimum(upper, zw) - lower, 0.0, None)
        weight = a_h * (strata.gamma_t[:, None] * (thick - wet) + strata.gamma_sat[:, None] * wet)
    else:
        weight = a_h * strata.gamma_t[:, None] * thick
    w = weight.sum(0)

    # 바닥이 놓인 층 — 경계 위면 위층(잘린 면 = 암반 위 흙의 강도)
    layer = (strata.bottoms[:, iy, ix] >= z_b + 1e-9).sum(0) if k_n > 1 else np.zeros(len(iy), int)
    c = strata.c[layer]
    tan_phi = strata.tan_phi[layer]

    q = strata.q[iy, ix] * a_h if strata.q is not None else np.zeros_like(w)
    mode = gw.get("mode", "none")
    if mode == "ru":
        u_ah = gw["ru"] * w  # Scoops3D 식 2.37 — u A_h = ru W
    elif mode == "line" and strata.z_water is not None:
        u_ah = GAMMA_W * np.maximum(strata.z_water[iy, ix] - z_b, 0.0) * a_h
    else:
        u_ah = np.zeros_like(w)

    n_g = len(cols.center)
    d = cols.direction if cols.direction is not None else fall_direction(strata, cols, n_g)
    dxg, dyg = d[grp, 0], d[grp, 1]
    xc, yc, zc = cols.center[grp, 0], cols.center[grp, 1], cols.center[grp, 2]
    p = (x - xc) * dxg + (y - yc) * dyg
    g = gx * dxg + gy * dyg
    sec_g = np.sqrt(1.0 + g * g)
    cos_psi = 1.0 / np.sqrt(1.0 + gx * gx + gy * gy)
    zr = z_b - zc
    return {
        "w": w,
        "q": q,
        "c": c,
        "tan_phi": tan_phi,
        "u_ah": u_ah,
        "a_h": a_h,
        "cos_psi": cos_psi,
        "sin_a": -g / sec_g,
        "cos_a": 1.0 / sec_g,
        "tan_a": -g,
        "p": p,
        "arm_s": (p * g - zr) / sec_g,
        "m_n": cos_psi * (p + zr * g),
        "e": zc - 0.5 * (z_top + z_b),
        "kh": kh,
        "depth": depth,
        "z_b": z_b,
        "z_top": z_top,
        "truncated": truncated,
        "direction": d,
        "group": grp,
        "n": n_g,
        "layer": layer,
    }


def _s_res(st: dict, f: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Bishop 수직력에서 바닥 저항력 S_res(F) · mα (Scoops3D 식 2.33 · 2.43)."""
    fg = np.maximum(f[st["group"]], 1e-3)
    m_a = st["cos_psi"] + st["sin_a"] * st["tan_phi"] / fg
    fric = np.maximum(st["w"] + st["q"] - st["u_ah"], 0.0) * st["tan_phi"]
    return (st["c"] * st["a_h"] + fric) / m_a, m_a


def _hovland(st: dict) -> np.ndarray:
    """N = (W+Q) cosψ · 기둥 사이 힘 0 — 반복 없음(Hovland 1977 · 지식DB 05 식 ①)."""
    grp, n = st["group"], st["n"]
    a = st["a_h"] / st["cos_psi"]
    wq = st["w"] + st["q"]
    nn = wq * st["cos_psi"]
    fric = np.maximum(nn - st["u_ah"] / st["cos_psi"], 0.0) * st["tan_phi"]
    res = _gsum(grp, st["arm_s"] * (st["c"] * a + fric), n)
    drive = _gsum(grp, -st["p"] * wq + st["e"] * st["kh"] * st["w"] + nn * st["m_n"], n)
    return _finish(res, drive)


def _finish(res: np.ndarray, drive: np.ndarray) -> np.ndarray:
    out = np.full(res.shape, F_CAP)
    ok = drive > 0
    out[ok] = np.minimum(np.maximum(res[ok], 0.0) / drive[ok], F_CAP)
    return out


def _fixed_point(st: dict, f: np.ndarray, coef: dict, step) -> tuple[np.ndarray, np.ndarray]:
    """F = step(기둥 묶음, F) 고정점 — 수렴한 활동체의 기둥은 빼며 셈(값은 빼지 않을 때와 같음).

    coef = 기둥별 상수 배열(group 포함) · 반환 (F, 수렴 여부).
    """
    n = st["n"]
    sub = coef
    conv = np.zeros(n, bool)
    for _ in range(MAX_ITER):
        f_new = step(sub, f)
        done = np.abs(f_new - f) < TOL
        f = np.where(conv, f, f_new)
        conv |= done
        if conv.all():
            break
        live = ~conv[sub["group"]]
        if live.mean() < 0.6:  # 남은 기둥만 — 매번 줄이면 줄이는 값이 더 듦
            sub = {k: v[live] for k, v in sub.items()}
    return f, conv


def _bishop(st: dict, f0: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    grp, n = st["group"], st["n"]
    wq = st["w"] + st["q"]
    # 거꾸로 선 바닥(α<0)이 있으면 처음 값을 올림 — Scoops3D 식 2.34 · 2.35
    neg = np.where(st["sin_a"] < 0, -st["sin_a"] * st["tan_phi"] / st["cos_psi"], 0.0)
    lift = np.zeros(n)
    np.maximum.at(lift, grp, neg)
    f = np.where(np.isfinite(f0), np.maximum(f0, 0.05), 1.0) + lift
    # S_res = num / mα · N = (W+Q − S sinα)/cosψ → 구동 = d1 − Σ S_res·k_n / F (기둥별 상수 미리)
    num = st["c"] * st["a_h"] + np.maximum(wq - st["u_ah"], 0.0) * st["tan_phi"]
    d1 = _gsum(
        grp, -st["p"] * wq + st["e"] * st["kh"] * st["w"] + wq * st["m_n"] / st["cos_psi"], n
    )
    coef = {
        "group": grp,
        "num": num,
        "cos_psi": st["cos_psi"],
        "k": st["sin_a"] * st["tan_phi"],
        "arm_s": st["arm_s"],
        "k_n": st["sin_a"] * st["m_n"] / st["cos_psi"],
    }

    def step(c: dict, f: np.ndarray) -> np.ndarray:
        g = c["group"]
        fg = np.maximum(f, 1e-3)
        s_res = c["num"] / (c["cos_psi"] + c["k"] / fg[g])
        res = _gsum(g, c["arm_s"] * s_res, n)
        drive = d1 - _gsum(g, s_res * c["k_n"], n) / fg
        return _finish(res, drive)

    f, conv = _fixed_point(st, f, coef, step)
    _, m_a = _s_res(st, f)
    ma_min = _gmin(grp, m_a, n)
    # 못 쓰는 면 — 수렴 안 함 · mα ≤ 0(바닥 수직력이 음) · F ≤ 0 → nan(후보에서 빠짐)
    f[~conv | (ma_min <= 0) | (f <= 0)] = np.nan
    return f, ma_min


def _gmin(grp: np.ndarray, values: np.ndarray, n: int) -> np.ndarray:
    out = np.full(n, np.inf)
    np.minimum.at(out, grp, np.nan_to_num(values, nan=-np.inf))
    return out


def _janbu(st: dict, f0: np.ndarray) -> np.ndarray:
    grp, n = st["group"], st["n"]
    wq = st["w"] + st["q"]
    drive = _gsum(grp, wq * st["tan_a"] + st["kh"] * st["w"], n)
    f = np.where(np.isfinite(f0), np.maximum(f0, 0.05), 1.0)
    num = st["c"] * st["a_h"] + np.maximum(wq - st["u_ah"], 0.0) * st["tan_phi"]
    coef = {
        "group": grp,
        "num": num / st["cos_a"],
        "cos_psi": st["cos_psi"],
        "k": st["sin_a"] * st["tan_phi"],
    }

    def step(c: dict, f: np.ndarray) -> np.ndarray:
        g = c["group"]
        s_res = c["num"] / (c["cos_psi"] + c["k"] / np.maximum(f, 1e-3)[g])
        return _finish(_gsum(g, s_res, n), drive)

    f, conv = _fixed_point(st, f, coef, step)
    f[~conv | (f <= 0)] = np.nan
    return f


def janbu_f0(st: dict) -> np.ndarray:
    """Janbu(1973) 보정 계수 — d/L 은 활동 방향 단면(현 = 윗끝 · 아랫끝 지표를 잇는 선)에서.

    b1 = 0.69(c=0) · 0.31(φ=0) · 0.50(c−φ). L = 활동 방향 길이 · d = 바닥에서 현까지 가장 먼 거리.
    """
    grp, n = st["group"], st["n"]
    s = st["p"]  # 축 기준 d 방향 위치(같은 묶음 안 비교만 씀)
    lo = np.full(n, np.inf)
    hi = np.full(n, -np.inf)
    np.minimum.at(lo, grp, s)
    np.maximum.at(hi, grp, s)
    # 양 끝 지표 높이 — 끝 기둥들 가운데 가장 얕은 바닥 위 지표
    at_lo = s <= lo[grp] + 1e-9
    at_hi = s >= hi[grp] - 1e-9
    z_lo = np.full(n, -np.inf)
    z_hi = np.full(n, -np.inf)
    np.maximum.at(z_lo, grp[at_lo], st["z_top"][at_lo])
    np.maximum.at(z_hi, grp[at_hi], st["z_top"][at_hi])
    length = np.maximum(hi - lo, 1e-9)
    # 현: (lo, z_lo) → (hi, z_hi) · 바닥점과의 수직 거리
    dz = z_hi - z_lo
    norm = np.hypot(length, dz)
    dist = np.abs((s - lo[grp]) * dz[grp] - (st["z_b"] - z_lo[grp]) * length[grp]) / norm[grp]
    depth = np.zeros(n)
    np.maximum.at(depth, grp, dist)
    ratio = depth / norm
    has_c = _gsum(grp, (st["c"] > 0).astype(float), n) > 0
    has_phi = _gsum(grp, (st["tan_phi"] > 0).astype(float), n) > 0
    b1 = np.where(has_c & has_phi, 0.50, np.where(has_phi, 0.69, 0.31))
    return 1.0 + b1 * (ratio - 1.4 * ratio**2)


def solve(
    strata: Strata,
    cols: Columns,
    groundwater_raw: dict | None = None,
    seismic_kh: float = 0.0,
) -> dict:
    """활동체마다 안전율 — hovland · bishop · janbu(간편) · janbu_corrected · 부피 · 잘림 비율.

    group 은 0..G−1 · 활동면이 지표 아래인 기둥만 셈(밖은 버림).
    """
    raw = groundwater_raw or {}
    # 3D 의 "line" = 지하수위면 격자(Strata.z_water) — 2D 선(line 칸)은 없어도 됨
    gw = {"mode": "line"} if raw.get("mode") == "line" and "line" not in raw else groundwater(raw)
    if gw["mode"] == "line" and strata.z_water is None:
        raise ValueError("지하수 mode line — Strata.z_water(지하수위면 격자)가 없음")
    st = _column_state(strata, cols, gw, float(seismic_kh))
    keep = st["depth"] > 0
    n = st["n"]
    st = {
        k: (v[keep] if isinstance(v, np.ndarray) and v.shape == keep.shape else v)
        for k, v in st.items()
    }
    grp = st["group"]
    count = np.bincount(grp, minlength=n)
    f_hov = _hovland(st)
    f_bis, ma_min = _bishop(st, f_hov)
    f_jan = _janbu(st, f_bis)
    f0 = janbu_f0(st)
    area = _gsum(grp, st["a_h"], n)
    trunc_area = _gsum(grp, st["truncated"] * st["a_h"], n)
    empty = count == 0
    out = {
        "hovland": f_hov,
        "bishop": f_bis,
        "janbu": f_jan,
        "janbu_corrected": f_jan * f0,
        "janbu_f0": f0,
        "volume": _gsum(grp, st["depth"] * st["a_h"], n),
        "weight": _gsum(grp, st["w"], n),
        "area": area,
        "truncated_ratio": np.divide(trunc_area, area, out=np.zeros(n), where=area > 0),
        "ma_min": ma_min,
        "columns": count,
        "direction": st["direction"],
    }
    for key in ("hovland", "bishop", "janbu", "janbu_corrected"):
        out[key] = np.where(empty, np.nan, out[key])
    return out