"""E04 3D 입력 격자 — 합성 DEM · 지층 경계 격자 · 노면 하중 격자.

합성 DEM(원지형 + B06 깎기 · 쌓기):
- 측점 i · i+1 횡단선을 t 로 잇는 선직면에 셀을 앉힘
  — P = C(t) + o·L(t) 를 t 의 2차식으로 풂 · 측점 위(t=0 · 1)는 횡단선 그대로.
- δ(o) = 설계선 z − 원지형 z(같은 offset) · 측점 사이는 t 선형 · 합성 = 원지형 + δ.
  → 측점 위 합성 단면 = B06 설계선 · δ=0 인 곳과 띠 밖은 원지형 그대로.
- 띠 = 두 측점 설계선 offset 범위의 합(B06 ±20 m · 사면이 안 닫히면 더 넓음)
  · 한쪽 설계선 밖은 δ=0 으로 이음 · 두 구간에 걸린 셀은 |o| 작은 쪽.

지층 경계(정함 · 보고):
- 흙 두께 T = 암 측점 |암 경계 offset|(B06 평행선) · 토사 측점 = 기본 3 m
  (국내 산지 풍화토 · 붕적토 1~3 m 의 큰 쪽).
- 띠 안 = 두 측점 T 를 t 선형 · 띠 밖 · 노선 끝 밖 = 가장 가까운 측점 값.
- 암 윗면 = 원지형 − T(깎기로 드러나면 지표) · 암 종류 = 가까운 측점 라벨
  (토사 측점 아래 = 리핑암).
- 층(위→아래) = 성토 · 풍화토 · 리핑암(풍화암 · 뚫림) · 발파암(연암 · 못 뚫음).
노면 하중 = 두 측점 노면 끝(road_edges) 사이(t 보간) 띠 안 셀에 q(kPa).
"""

from __future__ import annotations

import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.spatial import cKDTree

from E03_StabSection.E03_StabSection_Types import DEFAULT_MATERIALS, DEFAULT_ROAD_SURCHARGE_KPA
from E04_Stab3D.E04_Stab3D_Engine_Column import Strata

DEFAULT_SOIL_THICKNESS_M = 3.0
CLOSE_STEP_M = 0.5  # 안 닫힌 사면 늘리기 간격 · 상한
CLOSE_MAX_M = 60.0
ROCK_KINDS = ("ripping_rock", "blasting_rock")
#: 층 → 물성 id(E03 기본표) · 못 뚫는 층
LAYER_MATERIALS = ("fill", "weathered_soil", "weathered_rock", "soft_rock")
LAYER_HARD = (False, False, False, True)
LAYER_LABELS = ("fill", "soil", "ripping_rock", "blasting_rock")  # E02 label_map 열쇠


def station_from_b06(cross: dict, design: dict) -> dict:
    """B06 횡단 파일(`cross_NNNNNm.json`) + 설계(`cross_sections.data.design`) → 측점 한 줄."""
    ground = [[s["offset_m"], s["z"]] for s in cross["samples"] if s.get("valid", True)]
    edges = design.get("road_edges") or {}
    rb = design.get("rock_boundary_offset_m")
    return {
        "chainage_m": float(cross["chainage_m"]),
        "center": [float(cross["center_x"]), float(cross["center_y"])],
        "left": [float(v) for v in cross["frame"]["left_xy"]],
        "ground": ground,
        "design": [[p["offset_m"], p["elevation_m"]] for p in design["design_line"]],
        "ground_type": design.get("ground_type", "soil"),
        "rock_depth_m": abs(float(rb))
        if rb is not None and design.get("ground_type") in ROCK_KINDS
        else None,
        "road": sorted(float(e["offset_m"]) for e in edges.values()) if len(edges) == 2 else None,
    }


def _delta(
    st: dict, center: np.ndarray, left: np.ndarray, ground_at
) -> tuple[np.ndarray, np.ndarray]:
    """설계선 offset 위 δ = 설계 − 원지형(쌍선형) · 양 끝 밖 δ(0 · 자료 끝이면 끝 값).

    B06 이 반폭 밖으로 늘린 설계선도 받음.

    끝에서 사면이 안 닫혔으면(B06 slope_unclosed · |δ| > 5 cm) 끝 기울기로 원지형과 만날 때까지
    늘림(CLOSE_MAX_M 까지) — 안 늘리면 띠 끝에 절벽이 생겨 가짜 임계 활동면이 됨.
    """
    d = np.asarray(st["design"], float)
    d = d[np.argsort(d[:, 0], kind="stable")]
    pts = center[None] + d[:, :1] * left[None]
    delta = d[:, 1] - ground_at(pts)
    ok = np.isfinite(delta)
    o, z, delta = d[ok, 0], d[ok, 1], delta[ok]
    parts_o, parts_d = [o], [delta]
    hold = {-1: False, 1: False}
    for side in (-1, 1):
        k0, k1 = (0, 1) if side < 0 else (-1, -2)
        if len(o) < 2 or abs(delta[k0]) <= 0.05 or o[k0] == o[k1]:
            continue
        slope = (z[k0] - z[k1]) / (o[k0] - o[k1])
        ext = o[k0] + side * np.arange(CLOSE_STEP_M, CLOSE_MAX_M + 1e-9, CLOSE_STEP_M)
        de = z[k0] + slope * (ext - o[k0]) - ground_at(center[None] + ext[:, None] * left[None])
        cross = np.flatnonzero(~np.isfinite(de) | (np.sign(de) != np.sign(delta[k0])))
        n = cross[0] if len(cross) else len(ext)
        if n < len(ext) and np.isfinite(de[n]):
            ext, de = ext[: n + 1], np.append(de[:n], 0.0)  # 원지형과 만남 → 0
        else:
            ext, de = ext[:n], de[:n]  # 자료 끝 · 상한 — 끝 값을 이어 감(build_grid hold)
            hold[side] = True
        parts_o.append(ext)
        parts_d.append(de)
    o = np.concatenate(parts_o)
    delta = np.concatenate(parts_d)
    order = np.argsort(o, kind="stable")
    o, delta = o[order], delta[order]
    return o, delta, (delta[0] if hold[-1] else 0.0), (delta[-1] if hold[1] else 0.0)


def _segment_map(xx, yy, c0, c1, l0, l1):
    """선직면 P = C0 + tD + o(L0 + tM) 에서 (t, o[m]) — t∈[0,1] 아닌 곳은 nan."""
    d = c1 - c0
    m = l1 - l0
    qx, qy = xx - c0[0], yy - c0[1]

    def cross(ax, ay, bx, by):
        return ax * by - ay * bx

    a = -cross(d[0], d[1], m[0], m[1])
    b = cross(qx, qy, m[0], m[1]) - cross(d[0], d[1], l0[0], l0[1])
    c = cross(qx, qy, l0[0], l0[1])
    if abs(a) < 1e-12:
        t = np.where(np.abs(b) > 1e-12, -c / np.where(b == 0, 1, b), np.nan)
    else:
        disc = np.sqrt(np.maximum(b * b - 4 * a * c, 0.0))
        t1 = (-b + disc) / (2 * a)
        t2 = (-b - disc) / (2 * a)
        ok1 = (t1 >= -1e-9) & (t1 <= 1 + 1e-9)
        t = np.where(ok1, t1, t2)
    t = np.where((t >= -1e-9) & (t <= 1 + 1e-9), np.clip(t, 0, 1), np.nan)
    lx, ly = l0[0] + t * m[0], l0[1] + t * m[1]
    norm = np.hypot(lx, ly)
    o = ((qx - t * d[0]) * lx + (qy - t * d[1]) * ly) / norm
    return t, o


def build_grid(
    dem: dict,
    stations: list[dict],
    *,
    bbox: tuple[float, float, float, float] | None = None,
    default_soil_thickness_m: float = DEFAULT_SOIL_THICKNESS_M,
    road_q_kpa: float = DEFAULT_ROAD_SURCHARGE_KPA,
    cell_m: float | None = 1.0,
) -> dict:
    """원지형 DEM(x · y · z · valid_mask) + 측점 → 합성 · 지층 격자. bbox=(x0, x1, y0, y1) 로 자름.

    cell_m — 정사각 격자로 다시 뜸(쌍선형 · B04 축은 float32 라 간격이 흔들림)
    · None = DEM 축 그대로.
    """
    x = np.asarray(dem["x"], float)
    y = np.asarray(dem["y"], float)
    z = np.asarray(dem["z"], float)
    valid = np.asarray(dem.get("valid_mask", np.isfinite(z)), bool)
    if x[0] > x[-1]:
        x, z, valid = x[::-1], z[:, ::-1], valid[:, ::-1]
    if y[0] > y[-1]:
        y, z, valid = y[::-1], z[::-1], valid[::-1]
    z_orig = np.where(valid, z, np.nan)
    # δ 는 원래 DEM 에서 뜸(B06 지반선과 같은 면) — 다시 뜬 격자로 뜨면 cm 잡음이 끝 판정을 흔듦
    interp = RegularGridInterpolator((y, x), z_orig, bounds_error=False, fill_value=np.nan)
    if cell_m:
        x0, x1, y0, y1 = bbox or (x[0], x[-1], y[0], y[-1])
        x0, y0 = max(x0, x[0]), max(y0, y[0])
        xs = x0 + np.arange(int((min(x1, x[-1]) - x0) // cell_m) + 1) * cell_m
        ys = y0 + np.arange(int((min(y1, y[-1]) - y0) // cell_m) + 1) * cell_m
        gy_, gx_ = np.meshgrid(ys, xs, indexing="ij")
        z_orig = interp(np.column_stack([gy_.ravel(), gx_.ravel()])).reshape(gy_.shape)
        x, y = xs, ys
    elif bbox is not None:
        sx = (x >= bbox[0]) & (x <= bbox[1])
        sy = (y >= bbox[2]) & (y <= bbox[3])
        x, y, z_orig = x[sx], y[sy], z_orig[np.ix_(sy, sx)]
    z = z_orig
    xx, yy = np.meshgrid(x, y)
    sts = sorted(stations, key=lambda s: s["chainage_m"])
    centers = np.array([s["center"] for s in sts], float)
    lefts = np.array([s["left"] for s in sts], float)
    lefts /= np.hypot(lefts[:, 0], lefts[:, 1])[:, None]

    def ground_at(pts: np.ndarray) -> np.ndarray:
        return interp(np.column_stack([pts[:, 1], pts[:, 0]]))

    profiles = [_delta(s, centers[i], lefts[i], ground_at) for i, s in enumerate(sts)]
    thick = np.array(
        [
            s["rock_depth_m"] if s.get("rock_depth_m") is not None else default_soil_thickness_m
            for s in sts
        ]
    )
    kind = np.array([1 if s.get("ground_type") == "blasting_rock" else 0 for s in sts])

    seg = np.full(z.shape, -1)
    t_map = np.full(z.shape, np.nan)
    o_map = np.full(z.shape, np.nan)
    dz = np.zeros(z.shape)
    road = np.zeros(z.shape, bool)
    for i in range(len(sts) - 1):
        o0, d0, l0_, r0_ = profiles[i]
        o1, d1, l1_, r1_ = profiles[i + 1]
        lo, hi = min(o0[0], o1[0]), max(o0[-1], o1[-1])
        # 자료 끝에서 멈춘 쪽은 띠를 더 넓혀 끝 값을 이어 감(빈칸 셀은 합성도 빈칸)
        lo -= CLOSE_MAX_M if (l0_ or l1_) else 0.0
        hi += CLOSE_MAX_M if (r0_ or r1_) else 0.0
        reach = max(abs(lo), abs(hi))
        corners = np.array(
            [centers[i] + s * reach * lefts[i] for s in (-1, 1)]
            + [centers[i + 1] + s * reach * lefts[i + 1] for s in (-1, 1)]
        )
        jx = np.flatnonzero((x >= corners[:, 0].min() - 1) & (x <= corners[:, 0].max() + 1))
        jy = np.flatnonzero((y >= corners[:, 1].min() - 1) & (y <= corners[:, 1].max() + 1))
        if not len(jx) or not len(jy):
            continue
        sub = np.ix_(jy, jx)
        t, o = _segment_map(xx[sub], yy[sub], centers[i], centers[i + 1], lefts[i], lefts[i + 1])
        inside = np.isfinite(t) & (o >= lo) & (o <= hi)
        better = inside & (
            ~np.isfinite(o_map[sub]) | (np.abs(o) < np.abs(np.nan_to_num(o_map[sub], nan=np.inf)))
        )
        if not better.any():
            continue
        tb, ob = t[better], o[better]
        val = (1 - tb) * np.interp(ob, o0, d0, l0_, r0_) + tb * np.interp(ob, o1, d1, l1_, r1_)
        for arr, v in ((seg, i), (t_map, tb), (o_map, ob), (dz, val)):
            blk = arr[sub]
            blk[better] = v
            arr[sub] = blk
        r0, r1 = sts[i].get("road"), sts[i + 1].get("road")
        if r0 and r1:
            a = (1 - tb) * r0[0] + tb * r1[0]
            b = (1 - tb) * r0[1] + tb * r1[1]
            blk = road[sub]
            blk[better] = (ob >= a) & (ob <= b)
            road[sub] = blk

    in_band = seg >= 0
    dz = np.where(np.isfinite(z_orig), dz, 0.0)
    z_syn = z_orig + dz

    # 흙 두께 · 암 종류 — 띠 안 = t 보간 · 밖 = 가까운 측점
    near = cKDTree(centers).query(np.column_stack([xx.ravel(), yy.ravel()]))[1].reshape(z.shape)
    t_safe = np.nan_to_num(t_map)
    s_safe = np.maximum(seg, 0)
    t_band = (1 - t_safe) * thick[s_safe] + t_safe * thick[np.minimum(s_safe + 1, len(sts) - 1)]
    soil_t = np.where(in_band, t_band, thick[near])
    k_band = np.where(t_safe < 0.5, kind[s_safe], kind[np.minimum(s_safe + 1, len(sts) - 1)])
    rock_kind = np.where(in_band, k_band, kind[near])
    return {
        "x": x,
        "y": y,
        "z_orig": z_orig,
        "z": z_syn,
        "dz": dz,
        "in_band": in_band,
        "segment": seg,
        "t": t_map,
        "offset": o_map,
        "rock_top": z_orig - soil_t,
        "soil_thickness": soil_t,
        "rock_kind": rock_kind,  # 0 리핑암 · 1 발파암
        "q": np.where(road & in_band, road_q_kpa, 0.0),
        "stations": sts,
    }


def strata_from_grid(
    grid: dict,
    materials: list[dict] | None = None,
    label_map: dict[str, str] | None = None,
    *,
    water_depth_m: float | None = None,
    with_surcharge: bool = True,
) -> Strata:
    """격자 → 기둥 솔버 지층(층 차례 = 머리 주석) · 물성은 id 로 고름(E03 꼴 · 빠지면 기본표).

    label_map — E02 층 이름 → 물성 id(soil · ripping_rock · blasting_rock · fill) · 빠지면 기본 짝.

    water_depth_m — 지하수위면 = 합성 지표 − 깊이(0 = 지표까지 포화 · KDS 우기 규칙)
    · None = 지하수 없음.
    """
    by_id = {m["id"]: m for m in DEFAULT_MATERIALS}
    by_id.update({m["id"]: m for m in materials or []})
    z = grid["z"]
    fill_bot = np.minimum(z, grid["z_orig"])
    soil_bot = np.minimum(grid["rock_top"], fill_bot)
    blasting = grid["rock_kind"] == 1
    rip_bot = np.where(blasting, soil_bot, -1e6)
    return Strata(
        grid["x"],
        grid["y"],
        z,
        np.stack([fill_bot, soil_bot, rip_bot]),
        [by_id[(label_map or {}).get(lab, k)] for lab, k in zip(LAYER_LABELS, LAYER_MATERIALS)],
        hard=list(LAYER_HARD),
        z_water=None if water_depth_m is None else z - water_depth_m,
        q=grid["q"] if with_surcharge else None,
    )


def section_check(grid: dict, dem_sampler, station: dict) -> dict:
    """측점 횡단선 위 합성 DEM(쌍선형) − B06 설계선 · 원지형 − B06 지반선 — 오차 기록용."""
    d = np.asarray(station["design"], float)
    c = np.asarray(station["center"], float)
    left = np.asarray(station["left"], float)
    left = left / np.hypot(*left)
    pts = c[None] + d[:, :1] * left[None]
    f = RegularGridInterpolator(
        (grid["y"], grid["x"]), grid["z"], bounds_error=False, fill_value=np.nan
    )
    z_syn = f(np.column_stack([pts[:, 1], pts[:, 0]]))
    err = z_syn - d[:, 1]
    g = np.asarray(station["ground"], float)
    gp = c[None] + g[:, :1] * left[None]
    z0, _ = dem_sampler(gp)
    gerr = z0 - g[:, 1]
    return {
        "max_abs": float(np.nanmax(np.abs(err))),
        "rms": float(np.sqrt(np.nanmean(err**2))),
        "ground_max_abs": float(np.nanmax(np.abs(gerr))),
    }