diffraction.py

# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2025 SWGY, Inc
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
First-order edge diffraction: the Biot-Tolstoy-Medwin (BTM) DC edge gain.

The ray tracer is impulse-only, and impulse is the zero-frequency (DC) component
of the pressure. The diffracted pressure is a convolution ``p_inc * h_edge``, so
the diffracted *impulse* is ``I_inc * H_edge(0)`` with ``H_edge(0) = int h dt``
the edge's DC diffraction gain -- purely geometric. This is why BTM, not GTD/UTD:
the high-frequency coefficient scales ``1/sqrt(k)`` and diverges at DC, while the
time-domain response is integrable, giving a finite DC gain.

For a rigid (Neumann) wedge of open angle ``theta_w`` (``nu = pi/theta_w``), with
source S and receiver R described in edge-cylindrical coordinates (edge along z,
azimuth measured from the reference face), the diffracted DC field of a unit
Laplace point source is

    D = KAPPA * (-nu) * integral over the edge of  beta(z) / (m(z) l(z)) dz

    beta = sum_{k=1..4} sin(nu phi_k) / (cosh(nu eta) - cos(nu phi_k))
    phi_1..4 = pi +/- theta_S +/- theta_R          (all four sign combinations)
    cosh(eta) = (m l + (z-z_S)(z-z_R)) / (r_S r_R)
    m = sqrt(r_S^2 + (z-z_S)^2),  l = sqrt(r_R^2 + (z-z_R)^2)

``KAPPA = 1/(16 pi^2)`` is fixed, not fitted: it is validated in
``tests/test_diffraction.py`` against the exact static rigid-wedge Green's
function (a Legendre-Q modal series, itself checked against image sums) and
independently by the shadow-boundary continuity condition (the edge integral
jumps by 4*pi across the boundary, and continuity of geometric-optics + diffracted
field forces the diffracted jump to equal the incident (1/4pi)/R, so
KAPPA = 1/(4pi * 4pi)).

The engine (:mod:`warp_core`) does not use Laplace-normalized fields; it converts
this gain into its own ``p0/eff^2 * K`` impulse currency by matching the direct
ray's normalization (see :func:`build_sensor_coupling`), which also makes
geometric-optics + diffraction continuous across the shadow boundary.

Scope: first-order edge diffraction only -- no double diffraction, no
reflection x diffraction, no smooth-surface creeping waves. Impulse-only (DC),
linear acoustics.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np
from scipy import integrate

from sampling import GridArrays

# Validated DC coupling constant (see module docstring and tests).
KAPPA = 1.0 / (16.0 * np.pi**2)


@dataclass(frozen=True)
class EdgeArrays:
    """Compact diffracting-edge description (BTM secondary sources).

    Each row is one edge segment. The device builds the source/receiver angles
    from these + the blast and sensor positions, exactly as the ray fan is built
    from :class:`sampling.GridArrays`.
    """
    apex: np.ndarray      # (N, 3) segment midpoint, world coords
    tangent: np.ndarray   # (N, 3) unit edge tangent
    n_ref: np.ndarray     # (N, 3) unit reference-face direction (azimuth zero)
    nu: np.ndarray        # (N,) wedge index pi/theta_w  (1/2 for a screen)
    seg_len: np.ndarray   # (N,) segment length along the edge


def _beta_sum(nu, theta_S, theta_R, cosh_eta):
    """The 4-term BTM beta (rigid wedge), vectorised over ``cosh_eta``."""
    eta = np.arccosh(np.maximum(cosh_eta, 1.0))
    out = np.zeros_like(np.asarray(cosh_eta, dtype=float))
    for phi in (np.pi + theta_S + theta_R, np.pi + theta_S - theta_R,
                np.pi - theta_S + theta_R, np.pi - theta_S - theta_R):
        out += np.sin(nu * phi) / (np.cosh(nu * eta) - np.cos(nu * phi))
    return out


def _cosh_eta(z, r_S, z_S, r_R, z_R):
    m = np.sqrt(r_S**2 + (z - z_S)**2)
    l = np.sqrt(r_R**2 + (z - z_R)**2)
    return (m * l + (z - z_S) * (z - z_R)) / (r_S * r_R), m, l


def diffracted_dc_field(theta_w, r_S, theta_S, z_S, r_R, theta_R, z_R,
                        zmax=200.0):
    """Exact BTM diffracted DC field of a unit Laplace source (edge -inf..inf).

    Adaptive edge-coordinate integral; the reference against which the cheaper
    per-segment kernel path is checked. Returns the Laplace-normalized diffracted
    field (incident direct field of the same source is ``1/(4 pi R)``).
    """
    nu = np.pi / theta_w

    def integrand(z):
        cosh_eta, m, l = _cosh_eta(z, r_S, z_S, r_R, z_R)
        return _beta_sum(nu, theta_S, theta_R, cosh_eta) / (m * l)

    # flag the stationary point (apex) so quad resolves the near-boundary spike
    zc = 0.5 * (z_S + z_R)
    val, _ = integrate.quad(integrand, -zmax, zmax,
                            points=[z_S, z_R, zc], limit=500)
    return KAPPA * (-nu) * val


def segment_dc_gain(nu, r_S, theta_S, z_S, r_R, theta_R, z_R, seg_len):
    """Per-segment (midpoint) diffracted DC field -- the kernel's quadrature term.

    Summing this over the segments of an edge approximates
    :func:`diffracted_dc_field`; each Warp thread evaluates one segment.
    """
    cosh_eta, m, l = _cosh_eta(0.0, r_S, z_S, r_R, z_R)   # apex at edge origin
    beta = _beta_sum(nu, theta_S, theta_R, np.asarray(cosh_eta))
    return KAPPA * (-nu) * float(beta) / (m * l) * seg_len


def to_edge_frame(P, apex, tangent, n_ref):
    """World point -> (r, theta, z) in the edge-cylindrical frame.

    z along ``tangent`` from ``apex``; azimuth from ``n_ref`` toward
    ``tangent x n_ref`` (right-handed), so theta increases into the fluid region.
    """
    d = np.asarray(P, float) - np.asarray(apex, float)
    t = tangent / np.linalg.norm(tangent)
    e0 = n_ref / np.linalg.norm(n_ref)
    e1 = np.cross(t, e0)
    z = float(d @ t)
    perp = d - z * t
    r = float(np.linalg.norm(perp))
    theta = float(np.arctan2(perp @ e1, perp @ e0)) % (2.0 * np.pi)
    return r, theta, z


def build_sensor_coupling(grid: GridArrays) -> np.ndarray:
    """Per-explosion direct-path solid-angle weight S_fan(e) = sum cos(el) res^2.

    This is the identical weight the ray fan applies in ``warp_core._trace``.
    Reusing it verbatim is what makes geometric-optics + diffraction continuous
    across the shadow boundary in the engine's own units.
    """
    res2 = grid.res_rad**2
    cos_el = np.cos(grid.el_off) * grid.valid_el          # (E, Gel)
    n_az = grid.valid_az.sum(axis=1)                       # (E,)
    return (cos_el.sum(axis=1) * n_az * res2).astype(np.float64)


# --- Edge sourcing (trimesh; exercised in the container, not on the musl host) -

def rim_from_plane(mesh, normal=(0.0, 0.0, 1.0), offset_m=-0.09,
                   seg_len_m=0.01, nu=0.5) -> EdgeArrays:
    """Procedural rim: intersect the (closed, smooth) shell with a plane.

    MICH.obj has no rim edge (watertight, max dihedral ~28deg), so the helmet's
    dominant diffractor is defined as the shell's intersection with a plane -- a
    clean closed curve lying on the shell, which keeps geometric-optics +
    diffraction continuous. ``nu=0.5`` treats the thin shell edge as a screen
    (theta_w = 2pi), the conservative maximal-diffraction bound.
    """
    import trimesh

    normal = np.asarray(normal, float)
    origin = normal / (normal @ normal) * offset_m
    seg, _ = trimesh.intersections.mesh_plane(
        mesh, plane_normal=normal, plane_origin=origin, return_faces=True)
    if len(seg) == 0:
        raise ValueError("rim plane does not intersect the mesh; check offset")
    return _segments_to_edges(seg, seg_len_m, nu)


def feature_edges(mesh, angle_deg=30.0, seg_len_m=0.01) -> EdgeArrays:
    """Convex sharp edges (e.g. the SAPI plate perimeter chamfer).

    Uses trimesh face-adjacency dihedral angles; wedge index nu = pi/theta_w with
    theta_w = pi + dihedral for a convex edge. Concave edges are dropped.
    """
    ang = mesh.face_adjacency_angles
    mask = ang > np.radians(angle_deg)
    edges = mesh.face_adjacency_edges[mask]
    theta_w = np.pi + ang[mask]
    verts = mesh.vertices
    rows = []
    for (a, b), tw in zip(edges, theta_w):
        rows.append(_one_segment(verts[a], verts[b], np.pi / tw))
    return _split_segments(_stack_rows(rows), seg_len_m)


def load_polyline(path, nu=0.5, seg_len_m=0.01) -> EdgeArrays:
    """Authored world-coordinate rim polyline (fallback sourcing mode)."""
    import trimesh

    path_obj = trimesh.load(path)
    seg = np.asarray(path_obj.vertices)[path_obj.entities[0].points]
    seg = np.stack([seg[:-1], seg[1:]], axis=1)
    return _segments_to_edges(seg, seg_len_m, nu)


def concat(edge_list):
    """Stack several EdgeArrays into one (skips None / empty)."""
    parts = [e for e in edge_list if e is not None and e.apex.shape[0] > 0]
    if not parts:
        return None
    return EdgeArrays(
        np.concatenate([p.apex for p in parts]),
        np.concatenate([p.tangent for p in parts]),
        np.concatenate([p.n_ref for p in parts]),
        np.concatenate([p.nu for p in parts]),
        np.concatenate([p.seg_len for p in parts]))


def build_edge_sets(cfg, helmet, vest, back):
    """Per-config diffracting-edge sets, mirroring the mesh plan in ``simulate``.

    The helmet rim is sourced per ``cfg.diffraction.edge_mode``; the torso plates
    always use their perimeter feature edges.
    """
    dc = cfg.diffraction
    nu = 0.5 if dc.wedge_model == "screen" else None  # dihedral -> per-edge nu
    if dc.edge_mode == "rim_plane":
        rim = rim_from_plane(helmet, dc.rim_plane_normal, dc.rim_plane_offset_m,
                             dc.segment_length_m, nu=0.5)
    elif dc.edge_mode == "polyline":
        rim = load_polyline(dc.polyline_file, nu=0.5, seg_len_m=dc.segment_length_m)
    else:  # feature_edges on the (smooth) helmet -- usually sparse
        rim = feature_edges(helmet, dc.feature_angle_deg, dc.segment_length_m)
    front = feature_edges(vest, dc.feature_angle_deg, dc.segment_length_m)
    back_e = feature_edges(back, dc.feature_angle_deg, dc.segment_length_m)
    return {
        "full_armor": concat([rim, front, back_e]),
        "front_plate_helmet": concat([rim, front]),
        "helmet_only": rim,
        "vest_only": front,
        "no_armor": None,
    }


def _one_segment(p0, p1, nu):
    p0, p1 = np.asarray(p0, float), np.asarray(p1, float)
    mid = 0.5 * (p0 + p1)
    tan = p1 - p0
    length = np.linalg.norm(tan)
    tan = tan / length if length > 0 else np.array([1.0, 0.0, 0.0])
    # reference-face direction: any unit vector perpendicular to the tangent
    ref = np.cross(tan, [0.0, 0.0, 1.0])
    if np.linalg.norm(ref) < 1e-9:
        ref = np.cross(tan, [0.0, 1.0, 0.0])
    ref = ref / np.linalg.norm(ref)
    return mid, tan, ref, nu, length


def _stack_rows(rows):
    apex = np.array([r[0] for r in rows])
    tan = np.array([r[1] for r in rows])
    ref = np.array([r[2] for r in rows])
    nu = np.array([r[3] for r in rows])
    ln = np.array([r[4] for r in rows])
    return EdgeArrays(apex, tan, ref, nu, ln)


def _segments_to_edges(seg, seg_len_m, nu):
    rows = [_one_segment(s[0], s[1], nu) for s in seg]
    return _split_segments(_stack_rows(rows), seg_len_m)


def _split_segments(edges: EdgeArrays, seg_len_m) -> EdgeArrays:
    """Subdivide any segment longer than ``seg_len_m`` into equal pieces."""
    apex, tan, ref, nu, ln = [], [], [], [], []
    for i in range(edges.apex.shape[0]):
        n = max(1, int(np.ceil(edges.seg_len[i] / seg_len_m)))
        sub = edges.seg_len[i] / n
        start = edges.apex[i] - 0.5 * edges.seg_len[i] * edges.tangent[i]
        for j in range(n):
            c = start + (j + 0.5) * sub * edges.tangent[i]
            apex.append(c)
            tan.append(edges.tangent[i])
            ref.append(edges.n_ref[i])
            nu.append(edges.nu[i])
            ln.append(sub)
    return EdgeArrays(np.array(apex), np.array(tan), np.array(ref),
                      np.array(nu), np.array(ln))