"""Modal synthesis of a strike. Each mode m of the linearized plate is a damped
oscillator q'' + 2 sigma q' + omega^2 q = f, with sigma the sum of its three
loss rates. The stick is the Hertz contact of the full model, stepped against
the modal system with an exact update of every oscillator under a piecewise
constant force. Once the stick has left, every mode is a damped sinusoid with a
known complex amplitude, and the delayed Rayleigh sum over the plate collapses
to one complex gain per mode. During the attack, while the modes are being
driven, the delayed sum is evaluated directly from a table of the nodal
accelerations.
"""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np

from ..constants import STICK_RELEASE
from ..params import Listener, Strike
from ..plate import Plate
from ..sound import Sound
from .modes import Modes


def strike_gains(modes: Modes, strike: Strike) -> np.ndarray:
    """The contact: modal gains g_m = sum over the patch of (m_k / M_patch) phi_m(k),
    so that the modal force is -F g_m and the contact coordinate is sum_m g_m q_m."""
    pl = modes.plate
    patch = pl.strike_patch(strike)
    frac = pl.node_mass[patch.mask] / patch.mass
    return modes.w[:, patch.mask] @ frac


@dataclass
class Contact:
    """The state of every mode at the end of the contact phase and the tables recorded during it."""

    t_end: float
    """Time at which the stick is withdrawn, STICK_RELEASE after the contact last opened [s]."""
    t_contact: float
    """Time at which the contact last opened [s]; every mode decays freely from here on."""
    amplitude: np.ndarray
    """Complex modal amplitudes A at ``t_end``: q(t) = Re[A exp(s (t - t_end))], s = -sigma + i omega_d."""
    qdd: np.ndarray
    """Modal accelerations sampled at ``table_rate`` from t = 0 to t_end, ``(table_len, M)``."""
    table_rate: float
    q_frames: np.ndarray
    """Modal displacements at the movie frame times before ``t_end``, ``(n_frames, M)``."""
    force: np.ndarray
    """Stick force history at the fine step [N]."""
    dt: float
    peak_force: float
    contact_time: float
    """Total time the contact was closed [s]."""


def simulate_contact(modes: Modes, sigma: np.ndarray, g: np.ndarray, strike: Strike, fps: float, fs: float, dt: float = 2e-6, table_rate: float | None = None) -> Contact:
    """Steps the stick against the modal system until the contact has been open
    for STICK_RELEASE, recording the tables the synthesis needs. Every mode is
    advanced by its exact state transition under a constant force over the step;
    the stick by symplectic Euler; the force is the Hertz law of the penetration
    at the start of the step."""
    table_rate = table_rate or 2 * fs
    omega = modes.omega
    wd = np.sqrt(omega**2 - sigma**2)
    e, c, s = np.exp(-sigma * dt), np.cos(wd * dt), np.sin(wd * dt)
    p11 = e * (c + sigma * s / wd)
    p12 = e * s / wd
    p21 = -e * omega**2 * s / wd
    p22 = e * (c - sigma * s / wd)
    M = modes.count
    q, qd = np.zeros(M), np.zeros(M)
    t_max = 0.2
    max_steps = math.ceil(t_max / dt)
    table_cap = math.ceil(t_max * table_rate) + 2
    frame_cap = math.ceil(t_max * fps) + 2
    qdd = np.zeros((table_cap, M))
    q_frames = np.zeros((frame_cap, M))
    force = np.zeros(max_steps + 1)
    table_len = n_frames = 0
    peak_force = contact_time = 0.0
    z, zd = 0.0, -strike.velocity  # stick tip, positive away from the plate; it arrives moving toward it
    last_open, ever, F = 0.0, False, 0.0
    step = 0
    for step in range(max_steps):
        t = step * dt
        delta = float(g @ q) - z  # penetration: the plate has come past the tip
        F = strike.stiffness * delta**strike.exponent if delta > 0 else 0.0
        if F > 0:
            ever = True
            last_open = t + dt
            contact_time += dt
        peak_force = max(peak_force, F)
        force[step] = F
        while table_len < table_cap and table_len / table_rate <= t + 1e-12:
            qdd[table_len] = -2 * sigma * qd - omega**2 * q - F * g
            table_len += 1
        while n_frames < frame_cap and n_frames / fps <= t + 1e-12:
            q_frames[n_frames] = q
            n_frames += 1
        if ever and t - last_open >= STICK_RELEASE:
            break
        qp = -F * g / omega**2
        x = q - qp
        q = p11 * x + p12 * qd + qp
        qd = p21 * x + p22 * qd
        zd += F / strike.mass * dt
        z += zd * dt
    else:
        step = max_steps
    t_end = step * dt
    amplitude = q - 1j * (qd + sigma * q) / wd
    return Contact(t_end, min(last_open, t_end), amplitude, qdd[:table_len], table_rate, q_frames[:n_frames], force[: step + 1], dt, peak_force, contact_time)


def listener_weights(pl: Plate, listener: Listener) -> tuple[np.ndarray, np.ndarray, float]:
    """Listener terms with the common delay removed, so that only the spread of the delays across the plate remains."""
    weight, delay = pl.listener_terms(listener)
    sel = weight > 0
    delay = np.where(sel, delay - delay[sel].min(), 0.0)
    return weight, delay, float(delay.max())


def tail_start(c: Contact, max_delay: float) -> float:
    """The time from which the tail formula holds at the listener: every node's
    retarded time is past the last contact. The table path covers the output before it."""
    return min(c.t_end, c.t_contact + max_delay + 2 / c.table_rate)


def radiation_gains(modes: Modes, sigma: np.ndarray, weight: np.ndarray, delay: np.ndarray, batch: int = 256) -> np.ndarray:
    """The complex radiation gain of each mode for the listener:
    G_m = sum_k weight_k phi_m(k) exp(-s_m delay_k), with s_m = -sigma + i omega_d."""
    sel = weight > 0
    w, tau = weight[sel], delay[sel]
    wd = np.sqrt(modes.omega**2 - sigma**2)
    G = np.zeros(modes.count, dtype=complex)
    for a in range(0, modes.count, batch):
        b = min(modes.count, a + batch)
        phi = modes.w[a:b][:, sel]
        E = np.exp((sigma[a:b, None] - 1j * wd[a:b, None]) * tau[None, :])
        G[a:b] = (phi * w[None, :] * E).sum(axis=1)
    return G


def attack_pressure(modes: Modes, c: Contact, weight: np.ndarray, delay: np.ndarray, fs: float, t_start: float, out: np.ndarray) -> None:
    """The listener pressure over the attack from the nodal acceleration table:
    p(t) = sum_k weight_k wdd_k(t - delay_k), linearly interpolated, added to
    ``out`` for the samples before ``t_start``."""
    n_out = min(len(out), math.ceil(t_start * fs - 1e-9))
    if n_out <= 0:
        return
    length = min(len(c.qdd), math.ceil(t_start * c.table_rate) + 2)
    sel = weight > 0
    table = modes.w[:, sel].T @ c.qdd[:length].T  # (R, length) nodal accelerations
    x = (np.arange(n_out)[None, :] / fs - delay[sel][:, None]) * c.table_rate  # (R, n_out)
    i = np.floor(x).astype(int)
    valid = (x >= 0) & (i + 1 < length)
    i0 = np.clip(i, 0, length - 1)
    i1 = np.clip(i + 1, 0, length - 1)
    fr = x - i
    rows = np.arange(sel.sum())[:, None]
    val = (1 - fr) * table[rows, i0] + fr * table[rows, i1]
    out[:n_out] += (weight[sel][:, None] * np.where(valid, val, 0.0)).sum(axis=0)


def modal_tail(modes: Modes, sigma: np.ndarray, c: Contact, gains: np.ndarray, fs: float, t_start: float, out: np.ndarray, block: int = 4096) -> None:
    """The listener pressure from ``t_start`` on: p(t) = sum_m Re[G_m s_m^2 A_m
    exp(s_m (t - t_end))], added to ``out`` in blocks. Modes that have decayed
    by 140 dB are dropped."""
    n = len(out)
    n0 = math.ceil(t_start * fs - 1e-9)
    if n0 >= n:
        return
    wd = np.sqrt(modes.omega**2 - sigma**2)
    s = -sigma + 1j * wd
    coef = gains * s**2 * c.amplitude * np.exp(s * (n0 / fs - c.t_end))  # the value at sample n0
    Z = np.exp(np.outer(np.arange(block) / fs, s))  # (block, M): z^b
    scale = np.abs(coef)
    for start in range(n0, n, block):
        stop = min(n, start + block)
        cur = coef * np.exp(s * ((start - n0) / fs))
        active = np.abs(cur) > 1e-7 * scale
        if not active.any():
            break
        out[start:stop] += (Z[: stop - start, active] @ cur[active]).real


def movie_frames(modes: Modes, sigma: np.ndarray, c: Contact, fps: float, n_frames: int) -> np.ndarray:
    """Movie frames of the transverse displacement, ``(n_frames, ny, nx)`` float32:
    sum_m phi_m(k) q_m(t_f), with q from the contact tables before t_end and from the free decay after it."""
    wd = np.sqrt(modes.omega**2 - sigma**2)
    t = np.arange(n_frames) / fps - c.t_end
    Q = np.exp(-sigma[None, :] * t[:, None]) * (c.amplitude.real[None, :] * np.cos(wd[None, :] * t[:, None]) - c.amplitude.imag[None, :] * np.sin(wd[None, :] * t[:, None]))
    k = min(len(c.q_frames), n_frames)
    Q[:k] = c.q_frames[:k]
    g = modes.plate.grid
    sel = g.mask.reshape(-1)
    frames = np.zeros((n_frames, g.n), dtype=np.float32)
    frames[:, sel] = Q @ modes.w.reshape(modes.count, -1)[:, sel]
    return frames.reshape(n_frames, g.ny, g.nx)


def synthesize(modes: Modes, strike: Strike = Strike(), listener: Listener = Listener(), *, duration: float = 4.0, fs: int = 48000, fps: float = 120.0, movie: bool = False) -> Sound:
    """The sound of a strike on the plate by modal synthesis.

    Returns a ``Sound`` with the listener pressure [Pa] at ``fs`` over
    ``duration`` seconds after the impact and, with ``movie``, the transverse
    displacement at ``fps`` frames per second. The strike and the listener can
    be changed at no cost beyond this call; the modes are the expensive part.
    """
    pl = modes.plate
    sigma = modes.sigma
    g = strike_gains(modes, strike)
    c = simulate_contact(modes, sigma, g, strike, fps, fs)
    weight, delay, max_delay = listener_weights(pl, listener)
    t_start = tail_start(c, max_delay)
    n = round(duration * fs)
    audio = np.zeros(n)
    attack_pressure(modes, c, weight, delay, fs, t_start, audio)
    modal_tail(modes, sigma, c, radiation_gains(modes, sigma, weight, delay), fs, t_start, audio)
    frames = movie_frames(modes, sigma, c, fps, round(duration * fps)) if movie else None
    info = {
        "model": "linear",
        "modes": modes.count,
        "peak_force": c.peak_force,
        "contact_time": c.contact_time,
        "t_end": c.t_end,
        "force": c.force,
        "force_dt": c.dt,
    }
    return Sound(audio.astype(np.float32), fs, pl, strike, listener, frames, fps, info)