"""E03 2D 확률 해석 · 민감도 (요구 351 · PLAN 537-1) — Slide2 · SLOPE/W Probabilistic 과 같은 생각.

확률 변수 = 물성(c · φ · γ) · 지하수위(수위선 위아래 이동 m) · r_u — 정규 · 대수정규 · 절단(평균 아래 · 위
거리 = Slide2 Relative Min · Max · 물리 한계 c ≥ 0 · 0 ≤ φ < 89 · γ > 0 · 0 ≤ r_u < 1 도 절단) · 변수 사이 상관
(Slide2 c–φ Correlation · 표준정규에서 촐레스키) → 몬테카를로 · 라틴 하이퍼큐브 N 회
→ 안전율 표본 → 평균 · 표준편차 · Pf = P(F < 1) · β 정규 (μ − 1)/σ · β 대수정규
ln(μ/√(1+V²))/√ln(1+V²) · 히스토그램 · 누적분포 · P(F < 기준안전율).
갈래: fixed = 결정론 임계면 고정(절편 기하 한 번 · 물성만 바꿔 묶음 평가 · 빠름 · Slide2 Global Minimum) ·
search = 표본마다 임계면 다시 찾음(`Engine_Search.analyze` · 느림 · 프로세스 병렬 · Slide2 Overall Slope).
민감도 = 변수 하나씩 범위(절단 끝 · 없으면 ±3σ) 50 등분 훑기(나머지 평균 · 임계면 고정 · Slide2 Sensitivity 와 같음)
· 토네이도 = 끝값 F 폭 순.
γ 변수 = γt 표본 · γsat 는 같은 차만큼 옮김(γsat − γt 그대로).
search 일꾼 = `python -m E03_StabSection.E03_StabSection_Engine_Prob <job.json>`(서버 프로세스에서 풀을 열지 않음 ·
`Worker` 머리글) · 한 줄마다 {"done": n} · 끝에 job["out"] 파일.
NumPy · SciPy(stats) 만 · 계산 자리 = ③ 서버 단독.
"""

from __future__ import annotations

import copy
import json
import math
import sys
import time
from pathlib import Path

import numpy as np
from scipy.stats import norm

from E03_StabSection import E03_StabSection_Engine_LEM as lem
from E03_StabSection.E03_StabSection_Engine_Search import analyze
from E03_StabSection.E03_StabSection_Engine_Station import MIN_DEPTH_M
from E03_StabSection.E03_StabSection_Engine_Surface import (
    circle_slices,
    evaluate,
    poly_slices,
    three_point_center,
)

PROPS = ("c", "phi", "gamma", "water", "ru")
#: 물리 한계(절단) — 지하수 이동은 한계 없음
BOUNDS = {"c": (0.0, math.inf), "phi": (0.0, 89.0), "gamma": (1e-3, math.inf), "ru": (0.0, 0.99)}
#: 절단 없을 때 민감도 훑기 범위(±σ 배)
SWEEP_SIGMA = 3.0
SWEEP_POINTS = 51
HIST_BINS = 30
CDF_POINTS = 101
N_SLICES = 50


# ---- 분포 · 표본 --------------------------------------------------------------
def _bounds(v: dict) -> tuple[float, float]:
    """절단 구간 [a, b] — 평균 ∓ 아래 · 위 거리 ∩ 물리 한계."""
    lo_phys, hi_phys = BOUNDS.get(v["prop"], (-math.inf, math.inf))
    mean = float(v["mean"])
    a = mean - float(v["lo"]) if v.get("lo") is not None else -math.inf
    b = mean + float(v["hi"]) if v.get("hi") is not None else math.inf
    return max(a, lo_phys), min(b, hi_phys)


def _lognormal(v: dict) -> tuple[float, float]:
    """대수정규 (μ_ln, σ_ln) — 평균 · 표준편차 그대로 맞춤."""
    mean, std = float(v["mean"]), float(v["std"])
    s2 = math.log1p((std / mean) ** 2)
    return math.log(mean) - 0.5 * s2, math.sqrt(s2)


def ppf(v: dict, u: np.ndarray) -> np.ndarray:
    """균등 u(0~1) → 변수 값(절단 분포 역함수)."""
    u = np.asarray(u, float)
    mean, std = float(v["mean"]), float(v.get("std") or 0.0)
    if std <= 0:
        return np.full(u.shape, mean)
    a, b = _bounds(v)
    if v.get("dist") == "lognormal" and v["prop"] != "water":
        if mean <= 0:
            raise ValueError(f"대수정규는 평균 > 0: {v}")
        mu, s = _lognormal(v)
        za = (math.log(a) - mu) / s if a > 0 else -math.inf
        zb = (math.log(b) - mu) / s if b < math.inf else math.inf
        pa, pb = norm.cdf(za), norm.cdf(zb)
        return np.exp(mu + s * norm.ppf(pa + u * (pb - pa)))
    pa, pb = norm.cdf((a - mean) / std), norm.cdf((b - mean) / std)
    return mean + std * norm.ppf(pa + u * (pb - pa))


def sample(
    variables: list[dict],
    n: int,
    *,
    sampling: str = "lhs",
    seed: int = 0,
    correlations: list | None = None,
) -> np.ndarray:
    """표본 (n, 변수 수) — lhs = 라틴 하이퍼큐브(변수마다 n 칸 하나씩 · 칸 차례 섞음) · mc = 단순 무작위.

    correlations = [(i, j, ρ), …] — 표준정규 칸에서 촐레스키로 섞은 뒤 분포로(LHS 층은 섞인 변수만 흐려짐)."""
    rng = np.random.default_rng(int(seed))
    k = len(variables)
    if sampling == "lhs":
        order = np.argsort(rng.random((n, k)), axis=0)
        u = (order + rng.random((n, k))) / n
    else:
        u = rng.random((n, k))
    u = np.clip(u, 1e-12, 1 - 1e-12)
    if correlations:
        corr = np.eye(k)
        for i, j, rho in correlations:
            corr[i, j] = corr[j, i] = float(rho)
        u = norm.cdf(norm.ppf(u) @ np.linalg.cholesky(corr).T)
        u = np.clip(u, 1e-12, 1 - 1e-12)
    return np.column_stack([ppf(v, u[:, j]) for j, v in enumerate(variables)]) if k else u


def sweep_range(v: dict) -> tuple[float, float]:
    """민감도 훑기 범위 — 절단 끝 · 없으면 평균 ±3σ(물리 한계 안)."""
    mean, std = float(v["mean"]), float(v.get("std") or 0.0)
    a, b = _bounds(v)
    lo = a if math.isfinite(a) and v.get("lo") is not None else mean - SWEEP_SIGMA * std
    hi = b if math.isfinite(b) and v.get("hi") is not None else mean + SWEEP_SIGMA * std
    lo_phys, hi_phys = BOUNDS.get(v["prop"], (-math.inf, math.inf))
    return max(lo, lo_phys), min(hi, hi_phys)


# ---- 통계 -------------------------------------------------------------------
def statistics(F: np.ndarray, required: float | None = None, bins: int = HIST_BINS) -> dict:
    """안전율 표본 → 평균 · 표준편차 · Pf · β 정규 · 대수정규 · 히스토그램 · 누적분포(해 없는 표본 = 셈 밖)."""
    F = np.asarray(F, float)
    ok = F[np.isfinite(F)]
    out = {"n": int(len(F)), "n_valid": int(len(ok))}
    if len(ok) < 2:
        return {**out, "mean": None, "std": None, "pf": None, "beta_normal": None}
    mean, std = float(ok.mean()), float(ok.std(ddof=1))
    V = std / mean if mean > 0 else math.nan
    beta_ln = (
        math.log(mean / math.sqrt(1 + V * V)) / math.sqrt(math.log1p(V * V))
        if mean > 0 and V > 0
        else None
    )
    counts, edges = np.histogram(ok, bins=bins)
    p = np.linspace(0.0, 1.0, CDF_POINTS)
    out.update(
        mean=mean,
        std=std,
        cov=V if math.isfinite(V) else None,
        min=float(ok.min()),
        max=float(ok.max()),
        pf=float(np.mean(ok < 1.0)),
        n_fail=int(np.sum(ok < 1.0)),
        beta_normal=(mean - 1.0) / std if std > 0 else None,
        beta_lognormal=beta_ln,
        required=required,
        p_below_required=float(np.mean(ok < required)) if required else None,
        histogram={"edges": edges.tolist(), "counts": counts.tolist()},
        cdf={"p": p.tolist(), "F": np.quantile(ok, p).tolist()},
    )
    return out


# ---- 변수 → 모형 ---------------------------------------------------------------
def perturb(model: dict, variables: list[dict], row) -> dict:
    """모형 사본 — 표본 한 줄을 물성 · 지하수위에 넣음(search 갈래)."""
    out = dict(model)
    mats = {m["id"]: dict(m) for m in model["materials"]}
    dz = 0.0
    for v, x in zip(variables, row):
        if v["prop"] == "water":
            dz += float(x)
            continue
        if v["prop"] == "ru":
            gw = model.get("groundwater") or {}
            if gw.get("mode") == "ru":
                out["groundwater"] = {**gw, "ru": float(x)}
            continue
        m = mats[v["material"]]
        if v["prop"] == "gamma":
            gt = float(m["gamma_t"])
            m["gamma_sat"] = float(m.get("gamma_sat", gt)) + float(x) - gt
            m["gamma_t"] = float(x)
        else:
            m[v["prop"]] = float(x)
    out["materials"] = list(mats.values())
    gw = out.get("groundwater") or {}
    if dz and gw.get("mode") == "line":
        out["groundwater"] = {**gw, "line": [[x, z + dz] for x, z in gw["line"]]}
    return out


class Fixed:
    """임계면 고정 — 모형 · 면 하나를 한 번 준비해 표본마다 물성 칸만 갈아 끼우고 절편 묶음으로 평가."""

    def __init__(self, model: dict, surface: dict, direction: int, n_slices: int = N_SLICES):
        self.prep = lem.Prepared(model, direction)
        self.n = n_slices
        s = self.prep.direction
        self.index = {m["id"]: i for i, m in enumerate(self.prep.materials)}
        if surface["kind"] == "circle":
            xs = sorted((s * surface["entry"][0], s * surface["exit"][0]))
            self.geom = (
                "circle",
                s * surface["center"][0],
                surface["center"][1],
                surface["radius"],
            )
            self.geom += tuple(xs)
        else:
            pts = sorted((s * x, z) for x, z in surface["points"])
            self.geom = ("polyline", np.array([p[0] for p in pts]), np.array([p[1] for p in pts]))

    def _prep(self, variables: list[dict], row) -> lem.Prepared:
        base = self.prep
        p = copy.copy(base)
        gt, gs = base.gamma_t.copy(), base.gamma_sat.copy()
        c, tp = base.c.copy(), base.tanphi.copy()
        dz = 0.0
        for v, x in zip(variables, row):
            if v["prop"] == "water":
                dz += float(x)
                continue
            if v["prop"] == "ru":
                if base.gw_mode == "ru":
                    p.ru = float(x)
                continue
            i = self.index[v["material"]]
            if v["prop"] == "c":
                c[i] = x
            elif v["prop"] == "phi":
                tp[i] = math.tan(math.radians(x))
            else:
                gs[i] += x - gt[i]
                gt[i] = x
        p.gamma_t, p.gamma_sat, p.c, p.tanphi = gt, gs, c, tp
        p.table = base.table.copy()
        p.table[:, 4, :] = gt[base.mat]
        p.table[:, 5, :] = gs[base.mat] - gt[base.mat]
        if dz and base.gw_mode == "line":
            p.water_z = base.water_z + dz
        return p

    def _slices(self, prep: lem.Prepared) -> tuple[dict, np.ndarray]:
        g = self.geom
        if g[0] == "circle":
            _, xc, zc, R, xE, xX = g
            sl = circle_slices(prep, [xc], [zc], [R], [xE], [xX], self.n)
            return sl, np.array([[xc, zc]])
        vx, vz = g[1][None, :], g[2][None, :]
        return poly_slices(prep, vx, vz, self.n), three_point_center(vx, vz)

    def F(self, variables: list[dict], X: np.ndarray, method: str) -> np.ndarray:
        """표본 줄마다 안전율(해 없음 = nan)."""
        X = np.atleast_2d(np.asarray(X, float))
        parts = [self._slices(self._prep(variables, row)) for row in X]
        sl = _stack([p[0] for p in parts])
        center = np.vstack([p[1] for p in parts])
        mp_seed = None
        if method == "morgenstern_price":  # describe 와 같은 출발 참고값(같은 면 Spencer)
            mp_seed = evaluate(sl, center, ("spencer",))["spencer"]["F"]
            sl.pop("_pre", None)
        F = evaluate(sl, center, (method,), mp_seed=mp_seed)[method]["F"]
        return np.where(sl["valid"], F, np.nan)


def _stack(parts: list[dict]) -> dict:
    """한 줄짜리 절편 묶음들 → (S, …) 하나(배열 칸만 잇고 스칼라는 그대로)."""
    out = {}
    for key, value in parts[0].items():
        if key == "_pre":
            continue
        if isinstance(value, np.ndarray) and value.ndim >= 1 and value.shape[0] == 1:
            out[key] = np.concatenate([p[key] for p in parts])
        else:
            out[key] = value
    return out


def _chunks(n: int, size: int):
    for i in range(0, n, size):
        yield slice(i, min(i + size, n))


def sensitivity(
    fixed: Fixed, variables: list[dict], method: str, points: int = SWEEP_POINTS
) -> dict:
    """변수 하나씩 범위 훑기(나머지 평균) · 토네이도 = 끝값 F 폭 큰 순."""
    means = np.array([float(v["mean"]) for v in variables])
    rows, spans = [], []
    for j, v in enumerate(variables):
        lo, hi = sweep_range(v)
        xs = np.linspace(lo, hi, points) if hi > lo else np.full(points, means[j])
        block = np.tile(means, (points, 1))
        block[:, j] = xs
        rows.append(block)
        spans.append(xs)
    F0 = float(fixed.F(variables, means[None, :], method)[0])
    if not rows:
        return {"F_mean": F0, "curves": [], "tornado": []}
    F = fixed.F(variables, np.vstack(rows), method).reshape(len(variables), points)
    num = lambda a: [float(x) if np.isfinite(x) else None for x in a]  # noqa: E731
    curves = []
    for v, xs, f in zip(variables, spans, F):
        ends = [f[0], f[-1]]
        swing = abs(ends[1] - ends[0]) if all(np.isfinite(ends)) else None
        curves.append(
            {
                "name": v["name"],
                "x": num(xs),
                "F": num(f),
                "F_lo": num([ends[0]])[0],
                "F_hi": num([ends[1]])[0],
                "swing": swing,
            }
        )
    tornado = sorted(
        (c for c in curves if c["swing"] is not None), key=lambda c: c["swing"], reverse=True
    )
    return {"F_mean": F0 if math.isfinite(F0) else None, "curves": curves, "tornado": [
        {"name": c["name"], "F_lo": c["F_lo"], "F_hi": c["F_hi"], "swing": c["swing"]}
        for c in tornado
    ]}  # fmt: skip


# ---- 매회 탐색(search) ------------------------------------------------------------
_CTX: dict = {}


def _init(ctx: dict) -> None:
    _CTX.clear()
    _CTX.update(ctx)


def _search_one(row) -> float:
    c = _CTX
    m = c["method"]
    methods = tuple(dict.fromkeys(("bishop", m)))
    res = analyze(
        perturb(c["model"], c["variables"], row),
        direction=c["direction"],
        limits=c["limits"],
        methods=methods,
        noncircular=m not in ("ordinary", "bishop"),  # 비원호 Fellenius · Bishop = 참고 칸
        min_depth=MIN_DEPTH_M,
    )
    best = res["methods"][m]["min"]
    return float(best["F"]) if best and best.get("F") is not None else math.nan


def search_F(ctx: dict, X: np.ndarray, workers: int = 1, emit=None) -> np.ndarray:
    """표본마다 임계면 다시 찾음 — workers > 1 = 프로세스 풀(서버 밖 일꾼에서만)."""
    out = np.full(len(X), np.nan)
    if workers <= 1:
        _init(ctx)
        for i, row in enumerate(X):
            out[i] = _search_one(row)
            emit and emit(i + 1)
        return out
    from concurrent.futures import ProcessPoolExecutor, as_completed

    with ProcessPoolExecutor(max_workers=workers, initializer=_init, initargs=(ctx,)) as pool:
        futures = {pool.submit(_search_one, row): i for i, row in enumerate(X)}
        for done, future in enumerate(as_completed(futures), 1):
            out[futures[future]] = future.result()
            emit and emit(done)
    return out


# ---- 한 번 해석 ---------------------------------------------------------------
def run(
    model: dict,
    *,
    surface: dict,
    direction: int,
    limits=None,
    variables: list[dict],
    correlations: list | None = None,
    n: int = 1000,
    sampling: str = "lhs",
    seed: int = 0,
    mode: str = "fixed",
    method: str = "bishop",
    required: float | None = None,
    workers: int = 1,
    emit=None,
) -> dict:
    """확률 해석 + 민감도 한 벌. emit(끝난 표본 수) = 진행률."""
    t0 = time.perf_counter()
    X = sample(variables, n, sampling=sampling, seed=seed, correlations=correlations)
    fixed = Fixed(model, surface, direction)
    t1 = time.perf_counter()
    if mode == "search":
        ctx = {
            "model": model,
            "variables": variables,
            "direction": direction,
            "limits": limits,
            "method": method,
        }
        F = search_F(ctx, X, workers, emit)
    else:
        F = np.full(n, np.nan)
        for part in _chunks(n, 250):
            F[part] = fixed.F(variables, X[part], method)
            emit and emit(part.stop)
    t2 = time.perf_counter()
    sens = sensitivity(fixed, variables, method)
    t3 = time.perf_counter()
    realized = [
        {"name": v["name"], "mean": float(X[:, j].mean()), "std": float(X[:, j].std(ddof=1))}
        for j, v in enumerate(variables)
    ]
    return {
        "mode": mode,
        "method": method,
        "n": n,
        "sampling": sampling,
        "seed": seed,
        "variables": variables,
        "correlations": correlations or [],
        "realized": realized,
        "stats": statistics(F, required),
        "F": [round(float(f), 4) if np.isfinite(f) else None for f in F],
        "sensitivity": sens,
        "elapsed_s": {
            "prepare": round(t1 - t0, 3),
            "samples": round(t2 - t1, 3),
            "sensitivity": round(t3 - t2, 3),
            "total": round(t3 - t0, 3),
        },
    }


def main() -> None:
    """일꾼 — job = run() 인자 + out(결과 파일) · 한 줄마다 진행."""
    path = Path(sys.argv[1])
    job = json.loads(path.read_text(encoding="utf-8"))
    out = Path(job.pop("out"))
    result = run(**job, emit=lambda k: print(json.dumps({"done": k}), flush=True))
    out.write_text(json.dumps(result, ensure_ascii=False), encoding="utf-8")


if __name__ == "__main__":
    main()