"""E03 엄밀해 — Morgenstern-Price(GLE) · 절편간 힘 X = λ·f(x)·E · f = 반정현(기본) 또는 상수.

식 = LEM 과 같은 절편 평형(`E03_StabSection_Engine_LEM` 머리글) — 절편간 힘 함수만 경계마다 다름:
  절편 i 연직 평형 → N·(m_α − λ·f_{i+1}·A) = W − c'·sinα/F − X_i + λ·f_{i+1}·(E_i + B0)
  수평 평형 → E_{i+1} = E_i + B0 + N·A · X_{i+1} = λ·f_{i+1}·E_{i+1} · 머리 E_0 = 균열 수압 · X_0 = 0
  f ≡ 1 이면 Spencer 와 같은 식(시험으로 확인). F_m · F_f 는 LEM 과 같은 전체 평형식.
풀이 = λ 훑기(F_m(λ) · F_f(λ) 곡선 — 화면 그래프 자료) → 교점 → (F, λ) 2 차 뉴턴(수치 야코비안).
계산은 임계면 몇 개만(탐색 고리 밖) — 면을 λ 갯수만큼 복제해 한 번에 셈. 외부 코드 없음(NumPy 만).
"""

import numpy as np

from E03_StabSection import E03_StabSection_Engine_LEM as lem

FUNCTIONS = ("half_sine", "constant")
_TINY = 1e-9
CURVE_POINTS = 21
#: M-P λ 범위 — 반정현 평균 2/π 라 평균 X/E 가 Spencer 범위(LEM.LAMBDA_RANGE)와 같게 π/2 배
LAMBDA_RANGE = tuple(v * np.pi / 2 for v in lem.LAMBDA_RANGE)
#: 위 범위에 해가 없을 때만 다시 찾는 넓은 범위(같은 π/2 배 · 거친 훑기 → 뉴턴) — 그래도 없으면 「해 없음」
LAMBDA_WIDE = tuple(v * np.pi / 2 for v in lem.LAMBDA_WIDE)


def interslice_function(xe: np.ndarray, kind: str = "half_sine") -> np.ndarray:
    """절편 경계마다 f(x) — 반정현 sin(π·(x − x_진입)/(x_진출 − x_진입)) · 상수 1."""
    if kind == "constant":
        return np.ones_like(xe)
    if kind != "half_sine":
        raise ValueError(f"절편간 힘 함수 모름: {kind}")
    x0, x1 = xe[:, :1], xe[:, -1:]
    return np.sin(np.pi * (xe - x0) / np.maximum(x1 - x0, _TINY))


def forces(sl: dict, F, lam, f):
    """절편마다 N · 분모 · 경계 E · X (꼴 (S, n) · (S, n) · (S, n+1) · (S, n+1)) — 머리부터 차례로."""
    p = lem._pre(sl)
    F = np.asarray(F, float)
    lam = np.asarray(lam, float)
    g = 1.0 / F[:, None]
    m = sl["ca"] + g * p["sat"]
    A = sl["sa"] - g * p["tca"]
    B0 = p["khws"] - g * p["cuca"]
    rhs = sl["W"] - g * p["cusa"]
    S, n = rhs.shape
    N = np.empty((S, n))
    den = np.empty((S, n))
    E = np.empty((S, n + 1))
    X = np.zeros((S, n + 1))
    E[:, 0] = sl["Pw"]
    for i in range(n):
        lf = lam * f[:, i + 1]
        d = m[:, i] - lf * A[:, i]
        den[:, i] = d
        N[:, i] = (rhs[:, i] - X[:, i] + lf * (E[:, i] + B0[:, i])) / np.where(
            np.abs(d) < _TINY, _TINY, d
        )
        E[:, i + 1] = E[:, i] + B0[:, i] + N[:, i] * A[:, i]
        X[:, i + 1] = lf * E[:, i + 1]
    return N, den, E, X


def _fm_ff(sl, center, F, lam, f):
    N, den, *_ = forces(sl, F, lam, f)
    at, gn, d0, ma = lem._mom(sl, center)
    p = lem._pre(sl)
    fm = lem._ratio(ma + (at * N).sum(1), d0 + (gn * N).sum(1))
    ff = lem._ratio(p["rc0"] + (p["tca"] * N).sum(1), p["fd0"] + (sl["sa"] * N).sum(1), neg=False)
    return fm, ff, den


def _repeat(sl: dict, k: int) -> dict:
    """면마다 k 번 복제(λ 훑기 · 야코비안을 한 번에) — 미리 묶은 칸은 버림."""
    S = len(sl["W"])
    return {
        key: (np.repeat(v, k, axis=0) if isinstance(v, np.ndarray) and v.shape[:1] == (S,) else v)
        for key, v in sl.items()
        if key not in ("_pre", "_mom")
    }


def _solve_f(fun, F0, iters=60):
    """F = fun(F) — 수치 미분 뉴턴 · 안 되면 반쯤 섞은 고정점(LEM `_iterate` 와 같은 안전장치)."""
    F = np.asarray(F0, float).copy()
    done = np.zeros(len(F), bool)
    for _ in range(iters):
        f1 = fun(F)
        h1 = f1 - F
        dlt = 1e-6 * np.maximum(np.abs(F), 0.1)
        dh = ((fun(F + dlt) - (F + dlt)) - h1) / dlt
        newton = F - h1 / np.where(np.abs(dh) < 1e-12, -1.0, dh)
        use = np.isfinite(newton) & (dh < -1e-9) & (newton > 0.5 * F) & (newton < 2.0 * F)
        Fn = np.where(use, newton, 0.5 * (F + f1))
        Fn = np.where(np.isfinite(Fn) & (Fn > 0), Fn, np.nan)
        conv = np.abs(h1) < lem.TOL * np.maximum(1.0, np.abs(F))
        F = np.where(done | conv, F, Fn)
        done |= conv | ~np.isfinite(F)
        if done.all():
            break
    return np.where(done, F, np.nan)


def lambda_curves(sl, center, f, grid=None, F0=None):
    """λ 마다 F_m(λ) · F_f(λ) — 꼴 (G,) · (S, G) · (S, G)."""
    grid = np.linspace(*LAMBDA_RANGE, CURVE_POINTS) if grid is None else np.asarray(grid)
    S, G = len(sl["W"]), len(grid)
    rep = _repeat(sl, G)
    rc = np.repeat(center, G, axis=0)
    rf = np.repeat(f, G, axis=0)
    lam = np.tile(grid, S)
    start = np.full(S * G, 1.5) if F0 is None else np.repeat(F0, G)
    fm = _solve_f(lambda F: _fm_ff(rep, rc, F, lam, rf)[0], start)
    ff = _solve_f(lambda F: _fm_ff(rep, rc, F, lam, rf)[1], start)
    return grid, fm.reshape(S, G), ff.reshape(S, G)


def _crossing(grid, fm, ff):
    """F_m − F_f 부호가 바뀌는 곳 중 λ = 0 가까운 쪽 → (F, λ) · 없으면 nan."""
    S = len(fm)
    F, L = np.full(S, np.nan), np.full(S, np.nan)
    best = np.full(S, np.inf)
    d = fm - ff
    for k in range(len(grid) - 1):
        d0, d1 = d[:, k], d[:, k + 1]
        cross = np.isfinite(d0) & np.isfinite(d1) & (np.sign(d0) != np.sign(d1))
        t = np.where(cross, d0 / np.where(d0 == d1, 1, d0 - d1), 0)
        lc = grid[k] + t * (grid[k + 1] - grid[k])
        fc = fm[:, k] + t * (fm[:, k + 1] - fm[:, k])
        better = cross & (np.abs(lc) < best)
        best = np.where(better, np.abs(lc), best)
        F, L = np.where(better, fc, F), np.where(better, lc, L)
    return F, L


def _newton2(sl, center, f, F, lam, iters=40):
    """(F_m − F, F_f − F) = 0 — 세 벌 복제(F · F+h · λ+h)로 수치 야코비안 · 걸음 제한 LEM Spencer 와 같음."""
    S = len(F)
    rep, rc, rf = _repeat(sl, 3), np.repeat(center, 3, axis=0), np.repeat(f, 3, axis=0)
    F, lam = F.copy(), lam.copy()
    done = ~np.isfinite(F)
    it = np.zeros(S, int)
    for _ in range(iters):
        hF = 1e-6 * np.maximum(np.abs(F), 0.1)
        hl = 1e-6
        Fs = np.column_stack([F, F + hF, F]).ravel()
        ls = np.column_stack([lam, lam, lam + hl]).ravel()
        fm, ff, _ = _fm_ff(rep, rc, np.where(np.isfinite(Fs), Fs, 1.0), ls, rf)
        rm = (fm - Fs).reshape(S, 3)
        rf_ = (ff - Fs).reshape(S, 3)
        j11, j12 = (rm[:, 1] - rm[:, 0]) / hF, (rm[:, 2] - rm[:, 0]) / hl
        j21, j22 = (rf_[:, 1] - rf_[:, 0]) / hF, (rf_[:, 2] - rf_[:, 0]) / hl
        det = j11 * j22 - j12 * j21
        det = np.where(np.abs(det) < 1e-14, np.nan, det)
        dF = -(j22 * rm[:, 0] - j12 * rf_[:, 0]) / det
        dl = -(-j21 * rm[:, 0] + j11 * rf_[:, 0]) / det
        scale = np.minimum(
            1.0,
            np.minimum(
                0.5 * np.abs(F) / np.maximum(np.abs(dF), _TINY),
                0.3 / np.maximum(np.abs(dl), _TINY),
            ),
        )
        F = np.where(done, F, F + scale * dF)
        lam = np.where(done, lam, lam + scale * dl)
        it += ~done
        small = (np.abs(dF) < lem.TOL * np.maximum(1, np.abs(F))) & (np.abs(dl) < lem.TOL)
        done |= small | ~np.isfinite(F) | (F < lem.F_MIN)
        if done.all():
            break
    return F, lam, it


def _solve(sl, center, f, grid, fm, ff, F0, rng):
    """곡선 교점 → 뉴턴 → (F, λ, 반복, 수렴, 잔차) — 수렴 = 두 잔차 < 1e-4 · λ 범위 안 · 참고값 0.5~2 배."""
    F, lam = _crossing(grid, fm, ff)
    F, lam, it = _newton2(sl, center, f, F, lam)
    ok = np.isfinite(F)
    safe = np.where(ok & (F > 0), F, 1.0)
    rm, rff, _ = _fm_ff(sl, center, safe, np.where(ok, lam, 0.0), f)
    res = np.where(ok, np.maximum(np.abs(rm - F), np.abs(rff - F)), np.nan)
    conv = ok & (res < 1e-4) & (lam >= rng[0]) & (lam <= rng[1])
    if F0 is not None:
        conv &= (F > 0.5 * F0) & (F < 2.0 * F0)
    return F, lam, it, conv, res


def morgenstern_price(sl: dict, center: np.ndarray, kind="half_sine", F0=None, curve=True):
    """M-P (F, λ) · 수렴 · (curve) F_m · F_f–λ 곡선. F0 = 출발 참고값(Spencer · Bishop) — 엉뚱한 근 거름."""
    f = interslice_function(sl["xe"], kind)
    grid, fm, ff = lambda_curves(sl, center, f, F0=F0)
    F, lam, it, conv, res = _solve(sl, center, f, grid, fm, ff, F0, LAMBDA_RANGE)
    wide = ~conv & sl["valid"]
    if wide.any():
        # 좁은 범위에 해 없음 → 넓은 범위 거친 훑기 → 뉴턴(그 면만)
        sub, cw, fw = lem._subset(sl, wide), center[wide], f[wide]
        g2 = np.linspace(*LAMBDA_WIDE, lem.WIDE_POINTS)
        F0w = None if F0 is None else np.asarray(F0, float)[wide]
        _, fm2, ff2 = lambda_curves(sub, cw, fw, g2, F0w)
        F2, l2, it2, c2, r2 = _solve(sub, cw, fw, g2, fm2, ff2, F0w, LAMBDA_WIDE)
        F[wide], lam[wide], it[wide], conv[wide], res[wide] = F2, l2, it[wide] + it2, c2, r2
    ok = np.isfinite(F)
    safe = np.where(ok & (F > 0), F, 1.0)
    rm, rff, den = _fm_ff(sl, center, safe, np.where(ok, lam, 0.0), f)
    p = lem._pre(sl)
    N = forces(sl, safe, np.where(ok, lam, 0.0), f)[0]
    malpha_ok = np.all(den >= lem.MALPHA_MIN, axis=1)
    R = p["cu"] + N * sl["tanphi"]
    good = sl["valid"] & conv & (F > lem.F_MIN) & malpha_ok & (R.sum(1) > 0)
    out = {
        "F": np.where(good, F, np.nan),
        "lambda": np.where(good, lam, np.nan),
        "iterations": it,
        "converged": conv,
        "malpha_ok": malpha_ok,
        "residual": res,
        "wide": wide,
        "function": kind,
    }
    if curve:
        out["curve"] = (grid, fm, ff)
    return out