"""The time loop: leapfrog integration of the shell with contact, damping, rigid-mode projection
and radiation, block-compiled with JAX.

One sub-step advances the state (q, v) of the shell, the state (zs, vs) of every exciter and the
two retarded-time pressure rings:

    f      = f_int(q) - contact(q, zs) - 2 sigma1 rho h Lap(v_w)
    v'     = P_rigid [ (v + dt M^-1 f) / (1 + 2 sigma0 dt) ]
    q'     = q + dt v'
    vs'    = vs + dt F / m,    zs' = zs + dt vs'

where P_rigid removes the six rigid-body modes. The sub-step rate is an integer multiple m of the
audio rate chosen from the linear stability bound. The acceleration that radiates, and that the
optional record keeps, is the normal component of (v' - v) / dt, which for this scheme is the
exact central second difference of q. It is decimated to the audio rate inside the loop by the
polyphase accumulators of ``decimation.py``, so that no second difference is ever taken on
stored displacements. At the optional probe points (``output.probes``) the same acceleration is
interpolated, kept at the sub-step rate for one block, and decimated on the host with the same
kernel as the listener pressure. See docs/model.md, section "Time integration", and docs/format.md.
"""
from __future__ import annotations

import logging
import math
import time
from dataclasses import dataclass, field

import jax
import jax.numpy as jnp
import numpy as np
from jax import lax
from scipy.io import wavfile

from .config import SimulationConfig
from .decimation import DECIMATION_SPAN, StreamingDecimator, decimation_taps
from .exciter import HitSchedule, build_patch
from .geometry import CymbalProfile
from .mesh import Mesh, disk_mesh
from .radiation import RadiationModel
from .shell import N_DOF, ShellModel

log = logging.getLogger("cymsim")


@dataclass
class HitSummary:
    time: float
    kind: str
    peak_force: float           # [N]
    contact_time: float         # total time with F > 0 [s]
    rebound_velocity: float     # exciter velocity at the end of the run [m/s], positive = upward
    energy_in: float            # kinetic energy at release [J]
    energy_out: float           # kinetic energy left in the exciter [J]


@dataclass
class EnergyRecord:
    time: float
    kinetic: float
    membrane: float
    bending: float
    exciters: float
    max_w_over_h: float         # nonlinearity meter: < 1 effectively linear, 3 to 5 clear cascade

    @property
    def total(self) -> float:
        return self.kinetic + self.membrane + self.bending + self.exciters


@dataclass
class AccelerationRecord:
    """Normal acceleration of every node at the audio rate, int16 against a scale per sample.

    ``data[j, i]`` * ``scale[j]`` is the acceleration of node i along the rest-surface normal
    (front side positive) at t = j / fs, low-pass filtered by the kernel ``taps`` (see
    decimation.py). Sample j is final once the run has passed t = (j + W/2) / fs.
    """
    fs: int
    data: np.ndarray            # (T, N) int16
    scale: np.ndarray           # (T,) float32
    taps: np.ndarray            # (W, m) float64 polyphase kernel

    @property
    def n_samples(self) -> int:
        return len(self.data)

    @property
    def time(self) -> np.ndarray:
        return np.arange(self.n_samples) / self.fs

    def as_float(self, start: int = 0, stop: int | None = None) -> np.ndarray:
        """Decoded accelerations [m/s^2] for samples [start, stop), float32 (T', N)."""
        return self.data[start:stop].astype(np.float32) * self.scale[start:stop, None]

    @staticmethod
    def encode(a: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """(int16 data, float32 scale) of a float block (T', N)."""
        peak = np.abs(a).max(axis=1)
        scale = np.where(peak > 0, peak / 32767.0, 1.0).astype(np.float32)
        return np.round(a / scale[:, None]).astype(np.int16), np.where(peak > 0, scale, 0.0).astype(np.float32)

    @classmethod
    def empty(cls, fs: int, n_nodes: int, taps: np.ndarray) -> "AccelerationRecord":
        return cls(fs, np.zeros((0, n_nodes), np.int16), np.zeros(0, np.float32), taps)

    def extend(self, data: np.ndarray, scale: np.ndarray) -> None:
        self.data = np.concatenate([self.data, data])
        self.scale = np.concatenate([self.scale, scale])


@dataclass
class DisplacementRecord:
    """Transverse displacement w of every node [m] at a low frame rate, for display."""
    rate: float
    time: np.ndarray            # (Td,)
    w: np.ndarray               # (Td, N) float32

    @property
    def n_frames(self) -> int:
        return len(self.time)

    @classmethod
    def empty(cls, rate: float, n_nodes: int) -> "DisplacementRecord":
        return cls(rate, np.zeros(0), np.zeros((0, n_nodes), np.float32))

    def extend(self, time: np.ndarray, w: np.ndarray) -> None:
        self.time = np.concatenate([self.time, time])
        self.w = np.concatenate([self.w, w])


@dataclass
class SimulatorState:
    """Everything the Python integrator needs to continue a run from step ``n_steps_done``.

    Written to the ``state`` group of a .cymsim file. It is specific to this implementation
    (the web companion keeps its own); readers of the format may ignore it.
    """
    n_steps_done: int
    q: np.ndarray                       # (N, 5)
    v: np.ndarray                       # (N, 5)
    exciter_z: np.ndarray               # (K,)
    exciter_v: np.ndarray               # (K,)
    peak_force: np.ndarray              # (K,)
    contact_steps: np.ndarray           # (K,) int32
    ring_front: np.ndarray              # (L,) pending retarded-time pressure, sub-step rate
    ring_back: np.ndarray               # (L,)
    decimator_tail_front: np.ndarray    # sub-steps still needed by non-final audio samples
    decimator_tail_back: np.ndarray
    acceleration_ring: np.ndarray | None = None   # (W, N) partial audio samples of the record


@dataclass
class SimulationResult:
    config: SimulationConfig
    fs: int
    dt: float
    substeps: int
    mesh: Mesh
    z0_node: np.ndarray               # (N,) rest height of the nodes [m]
    h_node: np.ndarray                # (N,) thickness at the nodes [m]
    normal_node: np.ndarray           # (N, 3) unit normal of the rest surface, front side
    r_front: np.ndarray               # (E,) element centroid to listener [m]
    r_back: np.ndarray                # (E,) shortest path over the rim to the listener [m]
    pressure_rayleigh: np.ndarray     # baffled Rayleigh integral at the listener [Pa]
    pressure_unbaffled: np.ndarray    # front minus diffracted back term [Pa]
    hits: list[HitSummary]
    energy_log: list[EnergyRecord] = field(default_factory=list)
    acceleration: AccelerationRecord | None = None
    displacement: DisplacementRecord | None = None
    complete: bool = True
    n_steps_done: int = 0
    state: SimulatorState | None = None
    wall_time: float = 0.0
    creator: str = ""
    probes: np.ndarray | None = None                # (P, 2) probe points [m]
    probe_acceleration: np.ndarray | None = None    # (T, P) normal acceleration at the probes [m/s^2]

    @property
    def n_nodes(self) -> int:
        return self.mesh.n_nodes

    @property
    def n_elements(self) -> int:
        return self.mesh.n_elements

    @property
    def time(self) -> np.ndarray:
        return np.arange(len(self.pressure_rayleigh)) / self.fs

    @property
    def duration_done(self) -> float:
        """Simulated time covered by the run so far [s]."""
        return self.n_steps_done * self.dt

    def write_wav(self, prefix: str, normalize: bool = True) -> list[str]:
        """Write <prefix>_rayleigh.wav and <prefix>_unbaffled.wav as 16-bit PCM. With
        ``normalize`` the peak is placed at 0.9 full scale; otherwise 1 Pa maps to 0.05 FS."""
        paths = []
        for name, p in (("rayleigh", self.pressure_rayleigh), ("unbaffled", self.pressure_unbaffled)):
            peak = float(np.abs(p).max()) if len(p) else 0.0
            scale = 0.9 / (peak or 1.0) if normalize else 0.05
            path = f"{prefix}_{name}.wav"
            wavfile.write(path, self.fs, (np.clip(p * scale, -1, 1) * 32767).astype(np.int16))
            paths.append(path)
        return paths

    def write_npz(self, prefix: str) -> str:
        path = f"{prefix}_result.npz"
        np.savez(path, fs=self.fs, dt=self.dt, pressure_rayleigh=self.pressure_rayleigh,
                 pressure_unbaffled=self.pressure_unbaffled,
                 energy_time=[e.time for e in self.energy_log],
                 energy=[[e.kinetic, e.membrane, e.bending, e.exciters, e.max_w_over_h] for e in self.energy_log],
                 **({} if self.probe_acceleration is None else
                    dict(probes=self.probes, probe_acceleration=self.probe_acceleration)))
        return path

    def save(self, path: str) -> str:
        """Write a .cymsim file (see docs/format.md)."""
        from .cymsim_file import save_cymsim
        return save_cymsim(self, path)


class Simulator:
    """Build every part of the model from a SimulationConfig; ``run()`` integrates it."""

    def __init__(self, config: SimulationConfig):
        config.validate()
        self.config = config
        num, mat, out = config.numerics, config.material, config.output
        t0 = time.time()

        self.profile = CymbalProfile(config.profile, config.imperfections)
        self.mesh = disk_mesh(config.profile.radius, config.mesh.element_size, config.mesh.smoothing)
        log.info("mesh: %d nodes, %d triangles, %d DOFs (nonlinear = %s)",
                 self.mesh.n_nodes, self.mesh.n_elements, N_DOF * self.mesh.n_nodes, num.nonlinear)

        self.shell = ShellModel(self.mesh, self.profile, mat, num)
        sh = self.shell
        rim = sh.z0_node[self.mesh.boundary]
        log.info("profile: height %.1f mm, rim warp p-p %.2f mm, %d dents, h in [%.2f, %.2f] mm, "
                 "max |grad z0| = %.3f, mass %.3f kg",
                 sh.z0_node.max() * 1e3, (rim.max() - rim.min()) * 1e3, config.imperfections.n_dents,
                 sh.h_e.min() * 1e3, sh.h_e.max() * 1e3, sh.max_slope, sh.mass)
        if sh.max_slope > config.profile.max_slope:
            log.warning("max slope %.2f exceeds the shallow-shell bound %.2f", sh.max_slope, config.profile.max_slope)

        # ---- time step: integer number of sub-steps per audio sample, inside the stability bound.
        #      The run covers W/2 audio samples beyond the duration so that every sample of the
        #      duration is final after the decimation kernel's look-ahead.
        self.fs = num.fs
        self.substeps = math.ceil(1.0 / (self.fs * num.cfl_safety * sh.dt_stable))
        self.dt = 1.0 / (self.fs * self.substeps)
        self.n_out = int(round(self.fs * num.duration))
        self.span = DECIMATION_SPAN
        self.n_blocks = math.ceil((self.n_out + self.span // 2) * self.substeps / num.steps_per_block)
        self.n_steps = self.n_blocks * num.steps_per_block
        xc, yc = self.mesh.centroid[sh.limiting_element]
        log.info("omega_max: element bound %.3e rad/s (limiting element at r = %.3f m, h = %.2f mm), "
                 "global %.3e rad/s (%d power iterations); %s bound used: dt = %.3e s, "
                 "%d sub-steps per sample, %d steps",
                 sh.omega_max, math.hypot(xc, yc), sh.h_e[sh.limiting_element] * 1e3, sh.omega_global,
                 sh.omega_global_iterations, "global" if sh.omega_step < sh.omega_max else "element",
                 self.dt, self.substeps, self.n_steps)
        self.taps = decimation_taps(self.substeps)

        # ---- hits
        self.patches = [build_patch(h, self.mesh, self.profile) for h in config.hits]
        self.schedule = HitSchedule(self.patches, self.dt, sh.dtype)
        for k, p in enumerate(self.patches):
            tau = p.free_contact_time()
            if p.prescribed:
                log.info("hit %d at t = %.3f s: prescribed raised-cosine force at r = %.3f m, peak %.1f N, "
                         "duration %.2f ms", k, p.time, p.hit.r, p.force_peak, p.force_duration * 1e3)
                continue
            log.info("hit %d at t = %.3f s: %s, %d-point patch at r = %.3f m, k = %.3g N/m^%.2f, "
                     "m = %.3f kg, v = %.1f m/s, energy %.2f J, free contact ~ %.2f ms "
                     "(direct excitation below ~%.0f Hz)",
                     k, p.time, p.hit.exciter.kind, p.n_points, p.hit.r, p.k_total, p.alpha,
                     p.mass, p.velocity, p.energy, tau * 1e3, 0.5 / tau)
            if p.time >= num.duration:
                log.warning("hit %d is scheduled after the end of the run", k)

        # ---- radiation
        self.radiation = RadiationModel(self.mesh, self.profile, sh.z0_node, config.listener,
                                        self.dt, num.steps_per_block, sh.dtype)

        # ---- probes: containing element and barycentric weights of each point
        self.probe_xy = np.array(out.probes, dtype=np.float64).reshape(-1, 2)
        loc = [self.mesh.locate(x, y) for x, y in self.probe_xy]
        self.probe_nodes = jnp.asarray(np.array([l[0] for l in loc]).reshape(-1, 3), dtype=jnp.int32)
        self.probe_w = jnp.asarray(np.array([l[1] for l in loc]).reshape(-1, 3), dtype=sh.dtype)

        # ---- records
        self.record = bool(out.record_acceleration)
        self.disp_rate = float(out.displacement_rate)
        if self.record:
            nbytes = self.n_out * (2 * self.mesh.n_nodes + 4)
            log.info("acceleration record: %d samples x %d nodes, int16, %.2f GB in memory",
                     self.n_out, self.mesh.n_nodes, nbytes / 1e9)
            if nbytes > 2e9:
                log.warning("the acceleration record needs %.1f GB; consider a shorter duration or "
                            "a coarser mesh", nbytes / 1e9)
        S = self.fs * self.substeps                         # sub-steps per second
        if self.disp_rate > 0:
            n_frames = int(math.floor(self.n_out * self.substeps * self.disp_rate / S)) + 1
            self.frame_steps = np.round(np.arange(n_frames) * S / self.disp_rate).astype(np.int64)
        else:
            self.frame_steps = np.zeros(0, np.int64)
        self._build_step()
        self.setup_time = time.time() - t0
        log.info("setup done in %.1f s; backend %s", self.setup_time, jax.devices()[0].platform)

    # ---- one sub-step and its block loop
    def _build_step(self):
        sh, sched, rad = self.shell, self.schedule, self.radiation
        dt, sigma0 = self.dt, float(self.config.material.sigma0)   # a numpy scalar would promote float32
        m, W, N = self.substeps, self.span, self.mesh.n_nodes
        spb = self.config.numerics.steps_per_block
        tri, normal = sh.tri_j, sh.normal_j
        record, rate = self.record, self.disp_rate
        taps_j = jnp.asarray(self.taps, sh.dtype)
        S = float(self.fs * m)
        self.acc_rows = spb // m + 2                        # audio samples emitted per block, at most
        self.disp_rows = int(math.ceil(spb * rate / S)) + 2 if rate > 0 else 0

        n_probes = len(self.probe_xy)
        p_nodes, p_w = self.probe_nodes, self.probe_w

        def step(n, k_block0, j_block0, st):
            q, v, zs, vs, bf, bb, Fmax, ncont, ring, acc_out, disp_out, pbuf = st
            if rate > 0:
                # frame k sits at sub-step round(k S / rate); emit q at step n if n is one of them
                k_cand = jnp.round(n * (rate / S))
                is_frame = jnp.round(k_cand * (S / rate)) == n
                row = jnp.clip(k_cand.astype(jnp.int32) - k_block0, 0, self.disp_rows - 1)
                disp_out = disp_out.at[row].add(jnp.where(is_frame, q[:, 2], 0.0))
            Fk, F, zs, vs = sched.contact(n, q, zs, vs)
            fo = sh.internal_force(q)
            fo = sched.apply_to_plate(fo, Fk)
            fo = fo.at[:, 2].add(-sh.damping_force_w(v[:, 2]))
            v1 = (v + dt * fo * sh.Minv_j) / (1 + 2 * sigma0 * dt)   # sigma0 implicit
            v1 = sh.project_rigid(v1)
            q1 = q + dt * v1
            zs1, vs1 = sched.advance(zs, vs, F, dt)
            # normal acceleration at step n: central second difference of q, exact for leapfrog
            a_node = (((v1 - v) / dt)[:, :3] * normal).sum(1)
            acc_e = a_node[tri].mean(1)
            bf, bb = rad.scatter(n, acc_e, bf, bb)
            if n_probes:
                pbuf = pbuf.at[n % spb].set((a_node[p_nodes] * p_w).sum(1))   # blocks start at multiples of spb
            if record:
                j0, kk = n // m, n % m
                emit = kk == 0
                j_e = j0 - W // 2                            # the sample that just became final
                slot_e = j_e % W
                done = ring[slot_e]
                row = jnp.clip(j_e - j_block0, 0, self.acc_rows - 1)
                acc_out = acc_out.at[row].add(jnp.where(emit, done, 0.0))
                ring = ring.at[slot_e].set(jnp.where(emit, 0.0, done))
                # slot s holds sample j0 + 1 - W/2 + w with s = (j0 + 1 - W/2 + w) mod W
                r = (j0 + 1 - W // 2) % W
                tw = jnp.take(taps_j[:, kk], (jnp.arange(W) - r) % W)
                ring = ring + tw[:, None] * a_node[None, :]
            return (q1, v1, zs1, vs1, bf, bb, jnp.maximum(Fmax, F), ncont + (F > 0), ring, acc_out, disp_out, pbuf)

        @jax.jit
        def run_block(n0, k_block0, j_block0, st):
            return lax.fori_loop(0, spb, lambda k, s: step((n0 + k).astype(jnp.int32), k_block0, j_block0, s), st)

        @jax.jit
        def energies(st):
            q, v, _, vs = st[0], st[1], st[2], st[3]
            Em, Eb = sh.energy_parts_per_element(q)
            return (sh.kinetic_energy(v), Em.sum(), Eb.sum(), sched.energy(vs),
                    jnp.abs(q[:, 2] / sh.h_node_j).max())

        @jax.jit
        def clear(buf, idx):
            return buf.at[idx].set(0.0)

        self._step, self._run_block, self._energies, self._clear = step, run_block, energies, clear

    def initial_state(self):
        N, dtype = self.mesh.n_nodes, self.shell.dtype
        zs, vs = self.schedule.initial_state()
        bf, bb = self.radiation.initial_buffers()
        K = self.schedule.n_hits
        ring = jnp.zeros((self.span, N) if self.record else (0,), dtype)
        return (jnp.zeros((N, N_DOF), dtype), jnp.zeros((N, N_DOF), dtype), zs, vs, bf, bb,
                jnp.zeros(K, dtype), jnp.zeros(K, jnp.int32), ring, self._empty_outputs()[0],
                self._empty_outputs()[1], jnp.zeros((self.config.numerics.steps_per_block, len(self.probe_xy)), dtype))

    def _empty_outputs(self):
        N, dtype = self.mesh.n_nodes, self.shell.dtype
        return (jnp.zeros((self.acc_rows, N) if self.record else (0,), dtype),
                jnp.zeros((self.disp_rows, N) if self.disp_rate > 0 else (0,), dtype))

    def _hit_summaries(self, Fmax, ncont, vs) -> list[HitSummary]:
        hits = []
        for k, p in enumerate(self.patches):
            hits.append(HitSummary(p.time, p.hit.exciter.kind, float(Fmax[k]), float(ncont[k]) * self.dt,
                                   float(vs[k]), p.energy, 0.5 * p.mass * float(vs[k]) ** 2))
        return hits

    def _result(self, state, energy_log, p_front, p_back, acc, disp, n_done, wall, dec_f, dec_b,
                probe=None) -> SimulationResult:
        q, v, zs, vs, bf, bb, Fmax, ncont, ring = jax.device_get(state[:9])
        sim_state = SimulatorState(
            n_done, np.asarray(q), np.asarray(v), np.asarray(zs), np.asarray(vs), np.asarray(Fmax),
            np.asarray(ncont), np.asarray(bf), np.asarray(bb), dec_f.tail.copy(), dec_b.tail.copy(),
            np.asarray(ring) if self.record else None)
        n = min(len(p_front), len(p_back))
        p_front, p_back = np.asarray(p_front[:n]), np.asarray(p_back[:n])
        from . import __version__
        return SimulationResult(
            self.config, self.fs, self.dt, self.substeps, self.mesh, self.shell.z0_node, self.shell.h_node,
            self.shell.normal_node, self.radiation.r_front, self.radiation.r_back, p_front, p_front - p_back,
            self._hit_summaries(Fmax, ncont, vs), list(energy_log), acc, disp, n_done >= self.n_steps, n_done,
            sim_state, wall, f"cymsim {__version__}",
            self.probe_xy if len(self.probe_xy) else None, probe)

    def run(self, save_path: str | None = None, save_every: float | None = None) -> SimulationResult:
        """Integrate the run. With ``save_path`` a .cymsim file is written at the end and, if
        ``save_every`` (seconds of wall time) is given, also periodically during the run, each
        time as a consistent partial file that can be loaded (see docs/format.md)."""
        num = self.config.numerics
        m, spb, N = self.substeps, num.steps_per_block, self.mesh.n_nodes
        state = self.initial_state()
        energy_log: list[EnergyRecord] = []
        dec_f, dec_b = StreamingDecimator(m), StreamingDecimator(m)
        audio_f, audio_b = [], []
        n_audio = 0
        acc = AccelerationRecord.empty(self.fs, N, self.taps) if self.record else None
        disp = DisplacementRecord.empty(self.disp_rate, N) if self.disp_rate > 0 else None
        n_probes = len(self.probe_xy)
        dec_p = [StreamingDecimator(m) for _ in range(n_probes)]
        audio_p = []
        probe_now = lambda: np.concatenate(audio_p) if n_probes else None
        writer = None
        if save_path is not None:
            from .cymsim_file import CymsimWriter
            writer = CymsimWriter(save_path, self._result(state, [], np.zeros(0), np.zeros(0), acc, disp, 0, 0.0,
                                                          dec_f, dec_b))
        t_run = time.time()
        t_saved = t_run
        report_every = max(1, self.n_blocks // max(1, num.n_progress))
        log.info("running %d blocks of %d steps (nonlinearity meter: max|w|/h < 1 effectively "
                 "linear; 3 to 5 clear cascade)", self.n_blocks, spb)
        for blk in range(self.n_blocks):
            n0, n1 = blk * spb, (blk + 1) * spb
            k0 = int(np.searchsorted(self.frame_steps, n0))
            k1 = int(np.searchsorted(self.frame_steps, n1))
            j_block0 = -(-n0 // m) - self.span // 2
            acc_out, disp_out = self._empty_outputs()
            state = state[:9] + (acc_out, disp_out) + state[11:]
            state = self._run_block(jnp.int32(n0), jnp.int32(k0), jnp.int32(j_block0), state)

            # ---- pressure: the ring slots of this block's sub-steps are final; read, clear, decimate
            slots = jnp.asarray(self.radiation.slots(n0, n1))
            pf, pb = np.asarray(state[4][slots], np.float64), np.asarray(state[5][slots], np.float64)
            state = state[:4] + (self._clear(state[4], slots), self._clear(state[5], slots)) + state[6:]
            new_f, new_b = dec_f.push(pf), dec_b.push(pb)
            keep = max(0, min(len(new_f), self.n_out - n_audio))
            audio_f.append(new_f[:keep])
            audio_b.append(new_b[:keep])
            if n_probes:
                rows = np.asarray(state[11], np.float64)
                new_p = np.stack([d.push(rows[:, i]) for i, d in enumerate(dec_p)], 1)
                audio_p.append(new_p[:keep] if len(new_p) else np.zeros((0, n_probes)))
            n_audio += keep

            # ---- records
            if self.record:
                n_emit = -(-n1 // m) - -(-n0 // m)
                rows = np.asarray(state[9][:n_emit], np.float32)
                j = j_block0 + np.arange(n_emit)
                sel = (j >= 0) & (j < self.n_out)
                if sel.any():
                    acc.extend(*AccelerationRecord.encode(rows[sel]))
            if self.disp_rate > 0 and k1 > k0:
                frames = np.asarray(state[10][:k1 - k0], np.float32)
                t_frames = self.frame_steps[k0:k1] * self.dt
                sel = self.frame_steps[k0:k1] <= self.n_out * m
                disp.extend(t_frames[sel], frames[sel])

            if blk % report_every == 0 or blk == self.n_blocks - 1:
                ke, em, eb, ks, woh = (float(x) for x in self._energies(state))
                t = n1 * self.dt
                energy_log.append(EnergyRecord(t, ke, em, eb, ks, woh))
                log.info("  t = %6.3f s  KE %.3e  membrane %.3e  bending %.3e  exciters %.3e J  "
                         "max|w|/h = %5.2f  [%.0f s wall]", t, ke, em, eb, ks, woh, time.time() - t_run)
                if not math.isfinite(ke + em + eb):
                    raise FloatingPointError("energy is not finite: the run has blown up "
                                             "(lower numerics.cfl_safety)")
            if writer is not None and save_every is not None and time.time() - t_saved >= save_every \
                    and blk < self.n_blocks - 1:
                res = self._result(state, energy_log, np.concatenate(audio_f), np.concatenate(audio_b),
                                   acc, disp, n1, time.time() - t_run, dec_f, dec_b, probe_now())
                writer.update(res)
                t_saved = time.time()
                log.info("  saved partial run (%.3f s) to %s", res.duration_done, save_path)

        result = self._result(state, energy_log, np.concatenate(audio_f), np.concatenate(audio_b),
                              acc, disp, self.n_steps, time.time() - t_run, dec_f, dec_b, probe_now())
        for h in result.hits:
            log.info("hit (%s): peak force %.1f N, contact %.2f ms, rebound velocity %+.2f m/s "
                     "(energy left in exciter %.2f of %.2f J)", h.kind, h.peak_force,
                     h.contact_time * 1e3, h.rebound_velocity, h.energy_out, h.energy_in)
        for name, p in (("rayleigh", result.pressure_rayleigh), ("unbaffled", result.pressure_unbaffled)):
            peak = float(np.abs(p).max())
            log.info("%-10s: peak %.3f Pa (%.1f dB SPL peak)", name, peak, 20 * np.log10(max(peak, 1e-12) / 2e-5))
        if writer is not None:
            writer.update(result)
            writer.close()
            log.info("wrote %s", save_path)
        return result


def simulate(config: SimulationConfig) -> SimulationResult:
    return Simulator(config).run()