import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
from skfem import Basis, ElementTriP1, Mesh, asm, solve
from skfem.assembly import BilinearForm, FacetBasis, LinearForm
from skfem.helpers import dot, grad

from .domain import create_domain
from .fem_util import get_aggregation_per_element, get_element_midpoints


class HelmholtzProblem:
    """2D Helmholtz scattering FEM solver with Sommerfeld BC."""

    def __init__(
        self,
        *,
        fem_config: Dict[Union[str, int], Any],
        random_state: np.random.RandomState = np.random.RandomState(),
    ):
        helmholtz_config = fem_config.get("helmholtz", {})

        # ── 1. 波数 k ──
        wave_number_mode = helmholtz_config.get("wave_number_mode", "fixed")
        if wave_number_mode == "random_uniform":
            k_min = helmholtz_config.get("wave_number_min", 2.0)
            k_max = helmholtz_config.get("wave_number_max", 8.0)
            self._k = float(random_state.uniform(k_min, k_max))
        else:
            self._k = float(helmholtz_config.get("wave_number", 10.0))

        # ── 2. 介质散射体参数 ──
        sc = helmholtz_config.get("scatterer", {})
        scatterer_mode = sc.get("mode", "fixed")

        if scatterer_mode == "random_uniform":
            self._cx = float(
                random_state.uniform(sc.get("cx_min", 0.3), sc.get("cx_max", 0.7))
            )
            self._cy = float(
                random_state.uniform(sc.get("cy_min", 0.3), sc.get("cy_max", 0.7))
            )
            self._radius = float(
                random_state.uniform(
                    sc.get("radius_min", 0.1), sc.get("radius_max", 0.25)
                )
            )
            self._eps_r = float(
                random_state.uniform(
                    sc.get("eps_r_min", 2.0), sc.get("eps_r_max", 7.0)
                )
            )
        else:
            self._cx = float(sc.get("cx", 0.5))
            self._cy = float(sc.get("cy", 0.5))
            self._radius = float(sc.get("radius", 0.2))
            self._eps_r = float(sc.get("eps_r", 2.0))

        # ── 3. 组装 FEM 双线性和线性形式 ──
        self._bilin_form = self._make_bilinear_form()
        self._lin_form_real = self._make_linear_form_real()
        self._lin_form_imag = self._make_linear_form_imag()

        # ── 4. 初始化域（k^exponent 自适应网格密度 × domain area 线性缩放）──
        # exponent 和 k_ref 均可通过 helmholtz config 配置
        # exponent=2: P1 Helmholtz 理论最优 (污染误差 ∝ (kh)^2, N ∝ k^2)
        # exponent=1.5: 工程折中，避免高 k 初始过密
        # domain area 缩放: 保证不同域尺寸下每单位面积单元数一致 → h 不变
        domain_cfg = copy.deepcopy(fem_config.get("domain"))
        boundary = domain_cfg.get("boundary", [0, 0, 1, 1])
        domain_area = (boundary[2] - boundary[0]) * (boundary[3] - boundary[1])
        k_ref = helmholtz_config.get("k_ref", 6.0)
        k_exponent = helmholtz_config.get("k_exponent", 1.5)
        base_elements = domain_cfg.get("initial_num_elements", 400)
        scaled_elements = int(base_elements * (self._k / k_ref) ** k_exponent * domain_area)
        domain_cfg["initial_num_elements"] = max(scaled_elements, int(base_elements * domain_area))
        self._domain = create_domain(
            domain_config=domain_cfg,
            random_state=copy.deepcopy(random_state),
        )

        # ── 4.5. 介质区前渐近区边缘约束 ──
        # 放宽 Nyquist (N=4) → 前渐近区边缘 (N=1~1.5)，赋予介质内初始网格基本相位解析能力
        # 约束: h_init ≤ λ_local / N，λ_local = 2π/(k√ε_r)
        # N=1.5 对应约 1.5 点/波长，刚好跨过渐近区门槛，不撑爆物理预算
        pre_asymptotic_N = helmholtz_config.get("pre_asymptotic_N", 1.5)
        pre_asymptotic_mesh = self._enforce_nyquist_in_dielectric(
            self._domain.initial_mesh, N=pre_asymptotic_N
        )
        self._domain.replace_initial_mesh(pre_asymptotic_mesh)

        # ── 5. PDE 特征名称 ──
        pde_config = fem_config.get(fem_config.get("pde_type", "helmholtz"), {})
        self._element_feature_names = [
            name
            for name, include in pde_config.get("element_features", {}).items()
            if include
        ]

    # ── Public interface ─────────────────────────────────────

    def mesh_to_basis(self, mesh: Mesh) -> Basis:
        return Basis(mesh, ElementTriP1())

    def calculate_solution(self, basis: Basis, cache: bool = False) -> np.ndarray:
        K = asm(self._bilin_form, basis)
        f = asm(self._lin_form_real, basis) + 1j * asm(self._lin_form_imag, basis)

        boundary_facets = basis.mesh.boundary_facets()
        facet_basis = FacetBasis(basis.mesh, basis.elem, facets=boundary_facets)

        @BilinearForm
        def boundary_mass(u, v, w):
            return u * v

        M_boundary = asm(boundary_mass, facet_basis)
        K_total = K.astype(np.complex128) - 1j * self._k * M_boundary
        u_scat = solve(K_total, f)

        return u_scat

    def get_error_estimate_per_element(
        self, basis: Basis, solution: np.ndarray
    ) -> Dict[str, np.ndarray]:
        eps_r_arr = _compute_eps_r_at_midpoints(basis.mesh, self._cx, self._cy, self._radius, self._eps_r)
        return {"indicator": _compute_residual_indicator(basis.mesh, solution, k=self._k, eps_r=eps_r_arr)}

    def element_features(self, mesh: Mesh, element_feature_names: List[str]) -> Optional[np.ndarray]:
        features_list = []
        if "epsilon_r" in element_feature_names:
            features_list.append(
                _compute_eps_r_at_midpoints(mesh, self._cx, self._cy, self._radius, self._eps_r)[:, None]
            )
        return np.concatenate(features_list, axis=1) if features_list else None

    def _enforce_nyquist_in_dielectric(self, mesh: Mesh, N: float = 1.5, max_iter: int = 10) -> Mesh:
        """Iteratively refine elements inside the dielectric until h_K ≤ λ_d/N.

        λ_d = 2π/(k√ε_r) is the wavelength inside the dielectric.
        N=1.5 corresponds to the edge of the pre-asymptotic regime (~1.5 points
        per wavelength) — just enough for the wave field to exhibit basic phase
        resolution without exhausting the physical element budget. This relaxes
        the old Nyquist N=4 constraint, leaving headroom for the RL agent to
        selectively refine where residual indicators demand it.
        """
        lambda_d = 2.0 * np.pi / (self._k * np.sqrt(self._eps_r))
        h_max = lambda_d / N

        for _ in range(max_iter):
            i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
            x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
            x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
            x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

            e01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
            e12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            e20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
            h_K = np.maximum(np.maximum(e01, e12), e20)

            midpoints = np.mean(mesh.p[:, mesh.t], axis=1).T
            in_dielectric = (
                (midpoints[:, 0] - self._cx) ** 2
                + (midpoints[:, 1] - self._cy) ** 2
                <= self._radius**2
            )

            to_refine = np.where(in_dielectric & (h_K > h_max))[0]
            if len(to_refine) == 0:
                break
            mesh = mesh.refined(to_refine)

        return mesh

    # ── Properties ───────────────────────────────────────────

    @property
    def initial_mesh(self) -> Mesh:
        return self._domain.initial_mesh

    @property
    def boundary_line_segments(self) -> np.ndarray:
        return self._domain.boundary_line_segments


    @staticmethod
    def project_to_scalar(values: np.ndarray) -> np.ndarray:
        return values

    def additional_plots_from_mesh(self, mesh: Mesh) -> Dict:
        return {}

    # ── FEM form assembly ────────────────────────────────────

    def _eps_r_at_quad_points(self, x, y):
        in_cyl = (x - self._cx) ** 2 + (y - self._cy) ** 2 <= self._radius**2
        return np.where(in_cyl, self._eps_r, 1.0)

    def _make_bilinear_form(self):
        k = self._k

        @BilinearForm
        def bilin(u, v, w):
            eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
            return dot(grad(u), grad(v)) - k**2 * eps_r * u * v

        return bilin

    def _make_linear_form_real(self):
        k = self._k

        @LinearForm
        def lin(v, w):
            eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
            return k**2 * (eps_r - 1.0) * np.cos(k * w.x[0]) * v

        return lin

    def _make_linear_form_imag(self):
        k = self._k

        @LinearForm
        def lin(v, w):
            eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
            return k**2 * (eps_r - 1.0) * np.sin(k * w.x[0]) * v

        return lin


# ── 辅助函数 ──────────────────────────────────────────────────


def _compute_eps_r_at_midpoints(
    mesh: Mesh,
    cx: float = 0.5,
    cy: float = 0.5,
    radius: float = 0.2,
    eps_r_in: float = 2.0,
) -> np.ndarray:
    """
    计算每个单元中点处的相对介电常数 ε_r。

    判断单元中点是否落在介质圆柱内:
        - 在圆柱内 → ε_r = eps_r_in (如 2.0)
        - 在圆柱外 → ε_r = 1.0 (真空)

    Returns:
        eps_r: shape (num_elements,)
    """
    midpoints = get_element_midpoints(mesh)
    x_mid, y_mid = midpoints[:, 0], midpoints[:, 1]
    in_cylinder = (x_mid - cx) ** 2 + (y_mid - cy) ** 2 <= radius**2
    return np.where(in_cylinder, eps_r_in, 1.0)


def _compute_residual_indicator(
    mesh: Mesh,
    u_h: np.ndarray,
    k: float = 10.0,
    eps_r: Union[float, np.ndarray] = 1.0,
) -> np.ndarray:
    """
    基于残差的逐单元后验误差估计 — 无量纲化版本。

    引入局部波数 k_local = k√ε_r 消除纯几何尺度 h 带来的特征偏差，
    使误差指示子反映"相位分辨率残差"而非"网格粗疏程度"。

    P1 单元三项:
        1. r_int  = (h_K/k_local)·√V_K · |k²ε_r·u_h + k²(ε_r-1)·u_inc|
        2. r_jump = √(½ Σ_{e∈∂K} (h_e/k_local)·|[[∇u_h·n]]|²)
        3. r_sbc  = (h_bnd/k_local)·|∂u/∂n - i·k_local·u|

    Returns:
        eta_elements: shape (num_elements,) 的逐单元误差指标
    """
    n_elements = mesh.t.shape[1]
    eps_r = np.asarray(eps_r)
    k_local = k * np.sqrt(np.maximum(eps_r, 1.0))

    # ── 1. 单元几何量 ──
    i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
    x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
    x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
    x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

    det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    element_areas = np.abs(det_J) / 2.0

    edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
    edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
    edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
    h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)

    # ── 2. 梯度（常数，因为是 P1 单元）──
    u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
    inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
    du10, du20 = u1 - u0, u2 - u0

    grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
    grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
    grad_per_element = np.column_stack([grad_x, grad_y])

    # ── 3. 单元内部残差 ──
    u_mid = (u0 + u1 + u2) / 3.0
    x_mid = (x0 + x1 + x2) / 3.0
    u_inc_mid = np.exp(1j * k * x_mid)
    f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
    r_mid = f_mid + (k**2) * eps_r * u_mid

    cell_residual_sq = (h_K**2) * element_areas * np.abs(r_mid) ** 2 / (k_local ** 2)
    cell_residual_sq[element_areas < 1e-15] = 0.0

    # ── 4. 内部边梯度跳变 ──
    interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
    elem_left = mesh.f2t[0, interior_facets_idx]
    elem_right = mesh.f2t[1, interior_facets_idx]

    edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
    edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
    edge_vectors = edges_p2 - edges_p1
    h_e = np.linalg.norm(edge_vectors, axis=1)
    n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)

    grad_left = grad_per_element[elem_left]
    grad_right = grad_per_element[elem_right]
    jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))
    jump_val_sq = jump_val ** 2

    jump_residual_sq = np.zeros(n_elements)
    np.add.at(jump_residual_sq, elem_left, 0.5 * h_e * jump_val_sq / k_local[elem_left])
    np.add.at(jump_residual_sq, elem_right, 0.5 * h_e * jump_val_sq / k_local[elem_right])

    # ── 5. 合并 ──
    eta_sq = cell_residual_sq + jump_residual_sq

    # ── 6. SBC 边界残差 ──
    boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
    if len(boundary_facets_idx) > 0:
        bnd_elem = mesh.f2t[0, boundary_facets_idx]
        bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
        bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
        bnd_vectors = bnd_p2 - bnd_p1
        h_bnd = np.linalg.norm(bnd_vectors, axis=1)
        n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)

        grad_bnd = grad_per_element[bnd_elem]
        du_dn = np.sum(grad_bnd * n_bnd, axis=1)

        if eps_r.ndim == 1:
            k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
        else:
            k_local = k

        u_edge_mean = (
            u_h[mesh.facets[0, boundary_facets_idx]]
            + u_h[mesh.facets[1, boundary_facets_idx]]
        ) / 2.0
        sbc_residual = du_dn - 1j * k_local * u_edge_mean
        sbc_residual_sq = (h_bnd ** 2) * np.abs(sbc_residual) ** 2 / (k_local ** 2)
        np.add.at(eta_sq, bnd_elem, sbc_residual_sq)

    eta_sq = np.maximum(eta_sq, 0.0)
    return np.sqrt(eta_sq)


def _compute_residual_components(
    mesh: Mesh,
    u_h: np.ndarray,
    k: float = 10.0,
    eps_r: Union[float, np.ndarray] = 1.0,
    apply_log: bool = True,
) -> Dict[str, np.ndarray]:
    """
    计算逐单元的三项 PDE 物理残差（分离版，无量纲化）。

    引入 k_local = k√ε_r 消除几何尺度偏差，使 GNN 跨介质公平感知"相位分辨率残差"。
    保留源项信息（k²(ε_r-1)·u_inc），确保极粗网格下介质内部巨大物理激励仍可被网络捕捉。

    P1 单元返回:
        internal_residual: (h_K/k_local)·√V_i · |k²ε_r·u + k²(ε_r-1)·u_inc|
        gradient_jump:     √(½ Σ_{e∈∂K_i} (h_e/k_local)·|[[∇u·n]]|²)
        sbc_residual:      (h_bnd/k_local)·|∂u/∂n - i·k_local·u|
        element_areas:     单元面积
        is_sbc_boundary:   该单元是否与 SBC 边界相邻 (0/1)

    Args:
        apply_log: True → log10 压缩（喂 GNN）；False → 原始值（喂 reward）
    """
    n_elements = mesh.t.shape[1]
    eps_r = np.asarray(eps_r)
    k_local = k * np.sqrt(np.maximum(eps_r, 1.0))

    # ── 1. 单元几何量 ──
    i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
    x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
    x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
    x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

    det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
    element_areas = np.abs(det_J) / 2.0

    edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
    edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
    edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
    h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)

    # ── 2. 梯度（常数，因为是 P1 单元）──
    u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
    inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
    du10, du20 = u1 - u0, u2 - u0

    grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
    grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
    grad_per_element = np.column_stack([grad_x, grad_y])

    # P1 单元内部残差: ∇²u_h = 0（线性元二阶导为零），故仅含反应项
    # 完整强形式: |∇²u + k²·ε_r·u + k²·(ε_r-1)·u_inc|
    # 对 P1:  ∇²u_h ≡ 0 → 残差 = |k²·ε_r·u + k²·(ε_r-1)·u_inc|
    u_mid = (u0 + u1 + u2) / 3.0
    x_mid = (x0 + x1 + x2) / 3.0
    u_inc_mid = np.exp(1j * k * x_mid)
    f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
    r_mid = f_mid + (k**2) * eps_r * u_mid
    internal_residual = (h_K / k_local) * np.sqrt(element_areas) * np.abs(r_mid)
    internal_residual[element_areas < 1e-15] = 0.0

    # ── 4. 内部边梯度跳变 (逐单元) ──
    interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
    elem_left = mesh.f2t[0, interior_facets_idx]
    elem_right = mesh.f2t[1, interior_facets_idx]

    edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
    edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
    edge_vectors = edges_p2 - edges_p1
    h_e = np.linalg.norm(edge_vectors, axis=1)
    n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)

    grad_left = grad_per_element[elem_left]
    grad_right = grad_per_element[elem_right]
    jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))

    gradient_jump = np.zeros(n_elements, dtype=np.float64)
    jump_sq_per_edge = jump_val ** 2
    np.add.at(gradient_jump, elem_left, 0.5 * h_e * jump_sq_per_edge / k_local[elem_left])
    np.add.at(gradient_jump, elem_right, 0.5 * h_e * jump_sq_per_edge / k_local[elem_right])
    gradient_jump = np.sqrt(gradient_jump)

    # ── 5. SBC 边界残差 + 边界标记 ──
    sbc_residual = np.zeros(n_elements, dtype=np.float64)
    is_sbc_boundary = np.zeros(n_elements, dtype=np.float32)
    boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
    if len(boundary_facets_idx) > 0:
        bnd_elem = mesh.f2t[0, boundary_facets_idx]
        bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
        bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
        bnd_vectors = bnd_p2 - bnd_p1
        h_bnd = np.linalg.norm(bnd_vectors, axis=1)
        n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)

        grad_bnd = grad_per_element[bnd_elem]
        du_dn = np.sum(grad_bnd * n_bnd, axis=1)

        if eps_r.ndim == 1:
            k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
        else:
            k_local = k

        u_edge_mean = (
            u_h[mesh.facets[0, boundary_facets_idx]]
            + u_h[mesh.facets[1, boundary_facets_idx]]
        ) / 2.0
        sbc_val = np.abs(du_dn - 1j * k_local * u_edge_mean)
        np.add.at(sbc_residual, bnd_elem, (h_bnd / k_local) * sbc_val)
        is_sbc_boundary[bnd_elem] = 1.0

    # ── 对数预处理：压缩跨数量级动态范围（仅 GNN 特征需要）──
    if apply_log:
        _log_eps = 1e-8
        internal_residual = np.log10(np.maximum(internal_residual, _log_eps))
        gradient_jump = np.log10(np.maximum(gradient_jump, _log_eps))
        sbc_residual = np.log10(np.maximum(sbc_residual, _log_eps))

    return {
        "internal_residual": internal_residual.astype(np.float32),
        "gradient_jump": gradient_jump.astype(np.float32),
        "sbc_residual": sbc_residual.astype(np.float32),
        "element_areas": element_areas.astype(np.float32),
        "is_sbc_boundary": is_sbc_boundary,
    }


def _compute_residual_density(
    mesh: Mesh,
    u_h: np.ndarray,
    k: float = 10.0,
    eps_r: Union[float, np.ndarray] = 1.0,
) -> Dict[str, np.ndarray]:
    """
    Compute intensive (h-free) residual density components for reward.

    Unlike _compute_residual_components which includes h-scaling
    (h_K·√V, h_e·|jump|, h_bnd·|sbc|), this returns the raw PDE residuals
    that are independent of element size — true "error densities".

    Returns:
        rho_int:  |k²·ε_r·u + k²·(ε_r-1)·u_inc|           per element
        rho_jump: √(mean_{e∈∂K_int} |[[∇u·n]]|²)           per element
        rho_sbc:  √(mean_{e∈∂K∩Γ_sbc} |∂u/∂n - i·k·u|²)   per element
    """
    n_elements = mesh.t.shape[1]
    eps_r = np.asarray(eps_r)

    # ── 1. element geometry ──
    i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
    x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
    x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
    x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

    det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)

    # ── 2. gradient (constant per P1 element) ──
    u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
    inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
    du10, du20 = u1 - u0, u2 - u0

    grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
    grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
    grad_per_element = np.column_stack([grad_x, grad_y])

    # ── 3. interior residual density: |k²·ε_r·u_mid + k²·(ε_r-1)·u_inc_mid| ──
    u_mid = (u0 + u1 + u2) / 3.0
    x_mid = (x0 + x1 + x2) / 3.0
    u_inc_mid = np.exp(1j * k * x_mid)
    r_mid = (k**2) * eps_r * u_mid + (k**2) * (eps_r - 1.0) * u_inc_mid
    rho_int = np.abs(r_mid)

    # ── 4. gradient jump density: √(mean |[[∇u·n]]|²) per element ──
    interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
    elem_left = mesh.f2t[0, interior_facets_idx]
    elem_right = mesh.f2t[1, interior_facets_idx]

    edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
    edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
    edge_vectors = edges_p2 - edges_p1
    h_e = np.linalg.norm(edge_vectors, axis=1)
    n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)

    grad_left = grad_per_element[elem_left]
    grad_right = grad_per_element[elem_right]
    jump_val_sq = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1)) ** 2

    jump_sq_sum = np.zeros(n_elements, dtype=np.float64)
    jump_count = np.zeros(n_elements, dtype=np.float64)
    np.add.at(jump_sq_sum, elem_left, jump_val_sq)
    np.add.at(jump_sq_sum, elem_right, jump_val_sq)
    np.add.at(jump_count, elem_left, 1)
    np.add.at(jump_count, elem_right, 1)

    rho_jump = np.zeros(n_elements, dtype=np.float64)
    mask_jump = jump_count > 0
    rho_jump[mask_jump] = np.sqrt(jump_sq_sum[mask_jump] / jump_count[mask_jump])

    # ── 5. SBC boundary density: √(mean |∂u/∂n - i·k·u|²) per element ──
    rho_sbc = np.zeros(n_elements, dtype=np.float64)
    boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
    if len(boundary_facets_idx) > 0:
        bnd_elem = mesh.f2t[0, boundary_facets_idx]
        bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
        bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
        bnd_vectors = bnd_p2 - bnd_p1
        h_bnd = np.linalg.norm(bnd_vectors, axis=1)
        n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)

        grad_bnd = grad_per_element[bnd_elem]
        du_dn = np.sum(grad_bnd * n_bnd, axis=1)

        if eps_r.ndim == 1:
            k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
        else:
            k_local = k

        u_edge_mean = (
            u_h[mesh.facets[0, boundary_facets_idx]]
            + u_h[mesh.facets[1, boundary_facets_idx]]
        ) / 2.0
        sbc_val_sq = np.abs(du_dn - 1j * k_local * u_edge_mean) ** 2

        sbc_sq_sum = np.zeros(n_elements, dtype=np.float64)
        sbc_count = np.zeros(n_elements, dtype=np.float64)
        np.add.at(sbc_sq_sum, bnd_elem, sbc_val_sq)
        np.add.at(sbc_count, bnd_elem, 1)

        mask_sbc = sbc_count > 0
        rho_sbc[mask_sbc] = np.sqrt(sbc_sq_sum[mask_sbc] / sbc_count[mask_sbc])

    return {
        "rho_int": rho_int.astype(np.float64),
        "rho_jump": rho_jump.astype(np.float64),
        "rho_sbc": rho_sbc.astype(np.float64),
    }


# ── 工厂函数 ──────────────────────────────────────────────────


def create_helmholtz_problem(
    *, fem_config: Dict[Union[str, int], Any], random_state: np.random.RandomState
) -> HelmholtzProblem:
    """
    创建 Helmholtz 问题实例。

    Args:
        fem_config:   FEM 配置字典
        random_state: 随机状态

    Returns:
        HelmholtzProblem 实例
    """
    return HelmholtzProblem(fem_config=fem_config, random_state=random_state)
