"""Von Karman (Airy stress function) thin-plate solution of plate-strike-nl-v1 by a modal method.

The transverse displacement is expanded in the linear free-plate modes Phi_k computed with the
Bogner-Fox-Schmit elements of the linear solver (../kirchhoff-bfs/kl_bfs_plate.py), and the Airy
stress function in the clamped biharmonic eigenfunctions Psi_p on the same mesh. With
mass-normalized Phi and L2-normalized Psi, the equations are

    q_k'' + 2 alpha_k q_k' + omega_k^2 q_k = f_k psi(t/T_s) + sum_{j,p} H^p_{kj} q_j eta_p,
    eta_p = -(E h / (2 zeta_p^4)) sum_{i,j} H^p_{ij} q_i q_j,      H^p_{ij} = int Psi_p L(Phi_i, Phi_j) dA,

where L(f, g) = f_xx g_yy + f_yy g_xx - 2 f_xy g_xy (Touze et al.). Time stepping uses a two-step
Gautschi-type scheme at a step of 1/(fs m): the damped linear part is propagated exactly (the
recurrence coefficients are those of the exact damped oscillator), and the strike and the cubic
coupling force are applied with a static-exact weight, which is second-order accurate and free of
numerical dispersion in the linear limit. The Rayleigh integral is evaluated with
3 x 3 Gauss points per element; each point's retarded acceleration is a cubic Lagrange
interpolation of the modal accelerations at neighbouring time steps, and the interpolation weights
are folded into a per-step table, so the pressure at a step is a short sum over the recent modal
accelerations. Damping acts on the bending motion only (the in-plane state is quasi-static),
a further approximation of the challenge model.

Usage: python vk_modal_plate.py <challenge.json> <out-dir> [--fields tag,...] [--force F] [tag ...]
  where a tag is n<mesh>-w<transverse modes>-a<airy modes>-m<oversampling>, e.g. n24-w150-a150-m2
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np
from numpy.polynomial import legendre as leg
from scipy.io import wavfile
from scipy.linalg import eigh

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "kirchhoff-bfs"))
from kl_bfs_plate import BFSPlate, load_challenge  # noqa: E402


class VonKarmanModal:
    def __init__(self, ch: dict, n: int, n_w: int, n_airy: int, force_override: float | None = None):
        self.ch = ch
        p = dict(ch["parameters"])
        if force_override is not None:
            p["strike_peak_force"] = force_override
        ch2 = dict(ch)
        ch2["parameters"] = p
        self.p = p
        self.plate = BFSPlate(ch2, n)  # linear modes, modal forces, alpha, omega
        pl = self.plate
        # transverse modes: the n_w lowest (rigid ones included; they carry no membrane coupling)
        self.n_w = min(n_w, pl.n_modes)
        self.omega = pl.omega[: self.n_w]
        self.alpha = pl.alpha[: self.n_w]
        self.f = pl.f[: self.n_w]
        V = pl.V[:, : self.n_w]
        # Airy modes: clamped biharmonic eigenfunctions, L2-normalized: K_b psi = zeta^4 M_1 psi on interior dofs
        Kb = pl.assemble(pl.element_matrix(np.array([[1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]])))
        M1 = pl.M / pl.m  # unit-density mass matrix
        interior = np.setdiff1d(np.arange(pl.ndof), pl.boundary_dofs())
        z4, psi_int = eigh(Kb[np.ix_(interior, interior)], M1[np.ix_(interior, interior)])
        self.n_airy = min(n_airy, len(z4))
        self.zeta4 = z4[: self.n_airy]
        Psi = np.zeros((pl.ndof, self.n_airy))
        Psi[interior] = psi_int[:, : self.n_airy]
        # coupling tensor H[p, i, j] by 3 x 3 Gauss quadrature per element
        g, wg = leg.leggauss(3)
        s = (g + 1) / 2
        ws = wg / 2
        ex, ey = np.meshgrid(np.arange(n), np.arange(n), indexing="ij")
        XI, ETA = np.meshgrid(s, s, indexing="ij")
        xq = (ex.ravel()[:, None] + XI.ravel()[None, :]).ravel() * pl.a
        yq = (ey.ravel()[:, None] + ETA.ravel()[None, :]).ravel() * pl.b
        WQ = np.tile((ws[:, None] * ws[None, :]).ravel() * pl.a * pl.b, n * n)
        Bxx, Byy, Bxy = pl.basis_derivs_at(xq, yq)
        Pxx, Pyy, Pxy = Bxx @ V, Byy @ V, Bxy @ V  # (Q, n_w)
        PsiQ = pl.basis_at(xq, yq) @ Psi  # (Q, n_airy)
        H = np.zeros((self.n_airy, self.n_w, self.n_w))
        for pi in range(self.n_airy):
            wp = (WQ * PsiQ[:, pi])[:, None]
            H[pi] = (Pxx * wp).T @ Pyy + (Pyy * wp).T @ Pxx - 2.0 * (Pxy * wp).T @ Pxy
        self.H = H
        self.Hflat = H.reshape(self.n_airy * self.n_w, self.n_w)
        self.coef = -(p["E"] * p["h"]) / (2.0 * self.zeta4)
        # Rayleigh integral: quadrature weights, delays, modal shapes at the points
        r = np.sqrt((p["listener_x"] - xq) ** 2 + (p["listener_y"] - yq) ** 2 + p["listener_height"] ** 2)
        self.tau = r / p["c_air"]
        self.cq = p["rho_air"] / (2 * np.pi) * WQ / r
        self.PhiQ = pl.basis_at(xq, yq) @ V  # (Q, n_w)

    def nonlinear_force(self, q: np.ndarray) -> np.ndarray:
        v = (self.Hflat @ q).reshape(self.n_airy, self.n_w)  # v[p, k] = sum_j H[p,k,j] q_j
        eta = self.coef * (v @ q)  # eta_p
        return v.T @ eta  # F_k = sum_p v[p,k] eta_p

    def run(self, fs: int, nsamp: int, m: int, field_grid=None, field_rate: float | None = None):
        p = self.p
        dt = 1.0 / (fs * m)
        T, Om = p["strike_duration"], 2 * np.pi / p["strike_duration"]
        nsteps = nsamp * m
        w2, al, f = self.omega**2, self.alpha, self.f
        # exact two-step recurrence of the damped linear oscillator: q+ = a1 q + a2 q- + b F
        rigid = self.omega == 0.0
        disc = np.sqrt(al**2 - w2 + 0j)
        E1, E2 = np.exp((-al + disc) * dt), np.exp((-al - disc) * dt)
        a1 = (E1 + E2).real
        a2 = -(E1 * E2).real
        with np.errstate(divide="ignore", invalid="ignore"):
            b = np.where(rigid, 0.0, (1.0 - a1 - a2) / np.where(rigid, 1.0, w2))
        if rigid.any():
            e2 = np.exp(-2.0 * al[rigid] * dt)
            a1[rigid] = 1.0 + e2
            a2[rigid] = -e2
            b[rigid] = dt * (1.0 - e2) / (2.0 * al[rigid])
        # fractional delays: each quadrature point's retarded acceleration is a 4-point Lagrange
        # interpolation of the modal accelerations at steps d0-1 .. d0+2; the interpolation weights
        # are folded into a per-step weight table G[step, mode]
        tau_steps = self.tau / dt
        d0 = np.floor(tau_steps).astype(int)
        fr = tau_steps - d0
        lag = np.stack([
            -fr * (fr - 1.0) * (fr - 2.0) / 6.0,  # weight of step d0-1  (interpolating at d0+fr)
            (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,  # d0
            -(fr + 1.0) * fr * (fr - 2.0) / 2.0,  # d0+1
            (fr + 1.0) * fr * (fr - 1.0) / 6.0,  # d0+2
        ], axis=1)
        dmax = int(d0.max()) + 2
        G = np.zeros((dmax + 1, self.n_w))
        for j, off in enumerate([-1, 0, 1, 2]):
            dd = d0 + off
            for u in np.unique(dd):
                sel = dd == u
                G[u] += (self.cq[sel] * lag[sel, j]) @ self.PhiQ[sel]
        ring = np.zeros((dmax + 1, self.n_w))  # ring[n mod (dmax+1)] = qdd at step n
        out = np.zeros(nsamp)
        q_prev = np.zeros(self.n_w)
        q = np.zeros(self.n_w)
        frames = None
        if field_grid is not None:
            fx, fy, ft = field_grid
            XX, YY = np.meshgrid(fx, fy, indexing="xy")
            PhiG = self.plate.basis_at(XX.ravel(), YY.ravel()) @ self.plate.V[:, : self.n_w]
            frame_every = int(round(fs * m / field_rate))
            frames = []
        for nstep in range(nsteps + 1):
            t = nstep * dt
            psi = 0.5 * (1.0 - np.cos(Om * t)) if 0.0 <= t <= T else 0.0
            Fext = self.nonlinear_force(q) + f * psi  # forces other than the linear restoring and damping terms
            q_next = a1 * q + a2 * q_prev + b * Fext
            # acceleration from the equation of motion (velocity by centered difference)
            qdd = Fext - w2 * q - 2.0 * al * (q_next - q_prev) / (2.0 * dt)
            ring[nstep % (dmax + 1)] = qdd
            if nstep % m == 0:  # the ring buffer starts at zero, which is the state before the strike
                k = nstep // m
                if k < nsamp:
                    idx = (nstep - np.arange(dmax + 1)) % (dmax + 1)
                    out[k] = np.sum(G * ring[idx])
            if frames is not None and nstep % frame_every == 0 and len(frames) < len(ft):
                frames.append((PhiG @ q).astype(np.float32))
            q_prev, q = q, q_next
        fields = None
        if frames is not None:
            fields = np.stack(frames).reshape(len(frames), len(field_grid[1]), len(field_grid[0]))
        return out, fields


def parse_tag(tag: str):
    parts = dict((kv[0], int(kv[1:])) for kv in tag.split("-"))
    return parts["n"], parts["w"], parts["a"], parts["m"]


def main():
    ch = load_challenge(Path(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]
    fields_for = set()
    force = None
    if "--fields" in args:
        i = args.index("--fields")
        fields_for = set(args[i + 1].split(","))
        args = args[:i] + args[i + 2:]
    if "--force" in args:
        i = args.index("--force")
        force = float(args[i + 1])
        args = args[:i] + args[i + 2:]
    tags = args or ["n12-w50-a50-m1", "n16-w100-a100-m1", "n24-w150-a150-m2", "n24-w200-a200-m2"]
    out = ch["output"]
    fs, nsamp = out["sample_rate_hz"], out["num_samples"]
    meta = []
    for tag in tags:
        n, n_w, n_a, m = parse_tag(tag)
        t0 = time.perf_counter()
        model = VonKarmanModal(ch, n, n_w, n_a, force)
        t1 = time.perf_counter()
        grid = None
        if tag in fields_for:
            fg = ch["fields"]
            nt, ny, nx = fg["shape"]
            grid = (np.linspace(0, ch["parameters"]["Lx"], nx), np.linspace(0, ch["parameters"]["Ly"], ny), np.arange(nt) / fg["frame_rate_hz"])
        s, fields = model.run(fs, nsamp, m, grid, ch["fields"]["frame_rate_hz"] if grid else None)
        dt_total = time.perf_counter() - t0
        sid = f"vk-{tag}"
        out_dir.mkdir(parents=True, exist_ok=True)
        wavfile.write(out_dir / f"{sid}.wav", fs, (s / out["signal_scale"]).astype(np.float32))
        entry = {
            "solution_id": sid,
            "output": f"{out_dir.name}/{sid}.wav",
            "method": "von Karman plate (Airy form), modal method on Bogner-Fox-Schmit modes, Gautschi-type integrator exact in the linear part, retarded Rayleigh integral with cubic fractional delays",
            "parameters": {"elements_per_side": n, "transverse_modes": model.n_w, "airy_modes": model.n_airy, "oversampling": m, "dt_s": 1.0 / (fs * m), "highest_transverse_mode_hz": round(float(model.omega.max() / (2 * np.pi)), 1)},
            "timing_seconds": round(dt_total, 3),
            "peak_si": float(np.max(np.abs(s))),
        }
        if fields is not None:
            import h5py

            fg = ch["fields"]
            with h5py.File(out_dir / f"{sid}.h5", "w") as fh:
                fh.attrs["challenge_id"] = ch["challenge_id"]
                fh.attrs["solution_id"] = sid
                fh.attrs["quantity"] = fg["quantity"]
                fh.attrs["units"] = fg["units"]
                fh.create_dataset("x", data=grid[0])
                fh.create_dataset("y", data=grid[1])
                fh.create_dataset("t", data=grid[2][: fields.shape[0]])
                fh.create_dataset("u", data=fields, compression="gzip", compression_opts=4, chunks=(8, fields.shape[1], fields.shape[2]))
            entry["fields"] = {"path": f"{out_dir.name}/{sid}.h5", "format": "hdf5"}
        meta.append(entry)
        print(f"{tag}: modes w={model.n_w} a={model.n_airy} setup={t1 - t0:.1f}s total={dt_total:.1f}s peak={np.max(np.abs(s)):.4e} Pa", file=sys.stderr)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()