"""Kirchhoff-Love thin-plate solution of plate-strike-v1 with Bogner-Fox-Schmit finite elements.

Space: a uniform n x n mesh of C1-conforming bicubic Hermite rectangles (Bogner-Fox-Schmit),
four degrees of freedom per node (w, w_x, w_y, w_xy). Free edges are natural boundary
conditions of the Kirchhoff energy, so nothing is imposed. Element matrices are integrated
exactly with 4 x 4 Gauss-Legendre points. The generalized eigenproblem gives mass-normalized
modes, three of them rigid-body modes at zero frequency.

Time: each modal equation

    q'' + 2 alpha_n q' + omega_n^2 q = f_n psi(t / T_s),   alpha_n = sigma0 + eta omega_n^2 / 2,

is solved in closed form: a particular solution for the Hann pulse (a constant plus a cosine)
during the strike, and free decay afterwards.

Radiation: the Rayleigh integral is evaluated with 3 x 3 Gauss-Legendre points per element.
The retardation is exact: once the strike has ended at every retarded time the pressure is a sum
of complex exponentials whose amplitudes absorb the delays; the few earlier samples are evaluated
directly at the retarded times.

The transverse displacement w is positive in +z, so the strike enters as -p and w is reported
directly as u_z. Modes above 200 kHz are discarded; the pulse excites them at a negligible level
and their extreme damping rates would otherwise overflow the delay factors.

Usage: python kl_bfs_plate.py <challenge.json> <out-dir> [--fields n,...] [n ...]
"""
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

MODE_CUTOFF_HZ = 200_000.0
RIGID_OMEGA = 1.0  # rad/s; eigenvalues below this are the three rigid-body modes


def load_challenge(path: Path) -> dict:
    with open(path) as f:
        return json.load(f)


# ---------------------------------------------------------------- Hermite cubics on [0, 1]
def hermite(s: np.ndarray):
    """Values and first two derivatives (w.r.t. s) of the four cubic Hermite functions. Shapes (len(s), 4)."""
    s = np.asarray(s, dtype=float)
    H = np.stack([1 - 3 * s**2 + 2 * s**3, s - 2 * s**2 + s**3, 3 * s**2 - 2 * s**3, -(s**2) + s**3], axis=-1)
    dH = np.stack([-6 * s + 6 * s**2, 1 - 4 * s + 3 * s**2, 6 * s - 6 * s**2, -2 * s + 3 * s**2], axis=-1)
    d2H = np.stack([-6 + 12 * s, -4 + 6 * s, 6 - 12 * s, -2 + 6 * s], axis=-1)
    return H, dH, d2H


def shape_functions(xi: np.ndarray, eta: np.ndarray, a: float, b: float):
    """BFS shape functions and derivatives at local coordinates (xi, eta) in [0,1]^2 for an a x b element.

    Returns N, N_xx, N_yy, N_xy, each of shape (npts, 16). Local dof order: for corner c in
    [(0,0), (1,0), (0,1), (1,1)]: w, w_x, w_y, w_xy.
    """
    Hx, dHx, d2Hx = hermite(xi)
    Hy, dHy, d2Hy = hermite(eta)
    # per-corner 1D function indices: corner at xi=0 uses (H0, H1), at xi=1 uses (H2, H3)
    N = np.zeros((len(xi), 16))
    Nxx = np.zeros_like(N)
    Nyy = np.zeros_like(N)
    Nxy = np.zeros_like(N)
    corners = [(0, 0), (1, 0), (0, 1), (1, 1)]
    for c, (cx, cy) in enumerate(corners):
        ix0, ix1 = (0, 1) if cx == 0 else (2, 3)
        iy0, iy1 = (0, 1) if cy == 0 else (2, 3)
        # (1D function in x, scale for derivative dof) pairs for dofs w, w_x, w_y, w_xy
        combos = [(ix0, iy0, 1.0), (ix1, iy0, a), (ix0, iy1, b), (ix1, iy1, a * b)]
        for d, (fx, fy, sc) in enumerate(combos):
            k = 4 * c + d
            N[:, k] = sc * Hx[:, fx] * Hy[:, fy]
            Nxx[:, k] = sc * d2Hx[:, fx] * Hy[:, fy] / a**2
            Nyy[:, k] = sc * Hx[:, fx] * d2Hy[:, fy] / b**2
            Nxy[:, k] = sc * dHx[:, fx] * dHy[:, fy] / (a * b)
    return N, Nxx, Nyy, Nxy


class BFSPlate:
    def __init__(self, ch: dict, n: int):
        p = ch["parameters"]
        self.p = p
        self.n = n
        Lx, Ly, h = p["Lx"], p["Ly"], p["h"]
        D = p["E"] * h**3 / (12.0 * (1.0 - p["nu"] ** 2))
        m = p["rho"] * h
        nu = p["nu"]
        self.D, self.m = D, m
        a, b = Lx / n, Ly / n
        self.a, self.b = a, b

        # element matrices (uniform mesh: one of each)
        g, wg = leg.leggauss(4)
        s = (g + 1) / 2
        ws = wg / 2
        XI, ETA = np.meshgrid(s, s, indexing="ij")
        WQ = (ws[:, None] * ws[None, :]).ravel() * a * b
        N, Nxx, Nyy, Nxy = shape_functions(XI.ravel(), ETA.ravel(), a, b)
        C = np.array([[1.0, nu, 0.0], [nu, 1.0, 0.0], [0.0, 0.0, 2.0 * (1.0 - nu)]])
        Ke = np.zeros((16, 16))
        Me = np.zeros((16, 16))
        for q in range(len(WQ)):
            B2 = np.stack([Nxx[q], Nyy[q], Nxy[q]])  # (3, 16)
            Ke += WQ[q] * D * (B2.T @ C @ B2)
            Me += WQ[q] * m * np.outer(N[q], N[q])

        # assembly into dense global matrices
        nn = n + 1
        ndof = 4 * nn * nn
        self.ndof = ndof

        def node(i, j):
            return i * nn + j

        self.node = node
        self._elem_quad = (XI.ravel(), ETA.ravel(), WQ)
        self._shape = (N, Nxx, Nyy, Nxy)
        K = self.assemble(Ke)
        M = self.assemble(Me)
        self.K, self.M = K, M

        evals, evecs = eigh(K, M)  # mass-normalized eigenvectors
        omega2 = np.maximum(evals, 0.0)
        omega = np.sqrt(omega2)
        omega[omega < RIGID_OMEGA] = 0.0
        keep = omega <= 2 * np.pi * MODE_CUTOFF_HZ
        self.omega = omega[keep]
        self.V = evecs[:, keep]
        self.n_modes = int(keep.sum())
        self.alpha = p["sigma0"] + 0.5 * p["eta"] * self.omega**2
        self.rigid = self.omega == 0.0
        self.n_rigid = int(self.rigid.sum())

        # modal forces: f_n = -int p_hat Phi_n dA over the strike disk
        R, xs, ys = p["strike_radius"], p["strike_x"], p["strike_y"]
        c_phi = 0.5 - 2.0 / np.pi**2
        gr, gw = leg.leggauss(24)
        r = R * (gr + 1.0) / 2.0
        wr = gw * R / 2.0
        nth = 64
        th = 2 * np.pi * np.arange(nth) / nth
        rr, tt = np.meshgrid(r, th, indexing="ij")
        ww = (wr[:, None] * rr * (2 * np.pi / nth)).ravel()
        px = (xs + rr * np.cos(tt)).ravel()
        py = (ys + rr * np.sin(tt)).ravel()
        phat = p["strike_peak_force"] / (c_phi * np.pi * R**2) * 0.5 * (1.0 + np.cos(np.pi * rr.ravel() / R))
        self.f = -(ww * phat) @ self.modes_at(px, py)

        # closed-form response constants
        T = p["strike_duration"]
        Om = 2 * np.pi / T
        self.T, self.Om = T, Om
        al, w, f = self.alpha, self.omega, self.f
        with np.errstate(divide="ignore", invalid="ignore"):
            disc = np.sqrt(al**2 - w**2 + 0j)
            self.s1 = -al + disc
            self.s2 = -al - disc
            H = 1.0 / (w**2 - Om**2 + 2j * al * Om)
            self.qconst = np.where(self.rigid, 0.0, f / (2.0 * w**2))
            self.Hc = np.where(self.rigid, 0.0, (f / 2.0) * H)
            rhs1 = -(self.qconst - self.Hc.real)
            rhs2 = (1j * Om * self.Hc).real
            self.B = (rhs2 - self.s1 * rhs1) / (self.s2 - self.s1)
            self.A = rhs1 - self.B
            self.Gc = (f / 2.0) / (2.0 * al + 1j * Om)
            self.vconst = f / (4.0 * al)
            self.Ar = -(self.vconst - self.Gc.real)
        qT, vT, _ = self._pulse_phase(np.array([T]))
        self.qT, self.vT = qT[:, 0], vT[:, 0]
        with np.errstate(divide="ignore", invalid="ignore"):
            self.B2 = (self.vT - self.s1 * self.qT) / (self.s2 - self.s1)
            self.A2 = self.qT - self.B2

    # ------------------------------------------------------------ assembly helpers
    def element_dofs(self, ex: int, ey: int) -> np.ndarray:
        nodes = [self.node(ex, ey), self.node(ex + 1, ey), self.node(ex, ey + 1), self.node(ex + 1, ey + 1)]
        return np.array([4 * nd + d for nd in nodes for d in range(4)])

    def assemble(self, Ae: np.ndarray) -> np.ndarray:
        """Scatter one 16 x 16 element matrix (uniform mesh) into a dense global matrix."""
        A = np.zeros((self.ndof, self.ndof))
        for ex in range(self.n):
            for ey in range(self.n):
                dofs = self.element_dofs(ex, ey)
                A[np.ix_(dofs, dofs)] += Ae
        return A

    def element_matrix(self, Cmat: np.ndarray) -> np.ndarray:
        """Element 'stiffness' for the bilinear form int [w_xx, w_yy, w_xy] Cmat [v_xx, v_yy, v_xy]^T dA (no D factor)."""
        _, Nxx, Nyy, Nxy = self._shape
        WQ = self._elem_quad[2]
        Ae = np.zeros((16, 16))
        for q in range(len(WQ)):
            B2 = np.stack([Nxx[q], Nyy[q], Nxy[q]])
            Ae += WQ[q] * (B2.T @ Cmat @ B2)
        return Ae

    def boundary_dofs(self) -> np.ndarray:
        n = self.n
        b = []
        for i in range(n + 1):
            for j in range(n + 1):
                if i in (0, n) or j in (0, n):
                    b.extend(4 * self.node(i, j) + d for d in range(4))
        return np.array(sorted(set(b)))

    # ------------------------------------------------------------ evaluation at points
    def _locate(self, x: np.ndarray, y: np.ndarray):
        n, a, b = self.n, self.a, self.b
        ex = np.clip(np.floor(x / a).astype(int), 0, n - 1)
        ey = np.clip(np.floor(y / b).astype(int), 0, n - 1)
        return ex, ey, x / a - ex, y / b - ey

    def _scatter(self, vals: np.ndarray, ex: np.ndarray, ey: np.ndarray) -> np.ndarray:
        out = np.zeros((len(ex), self.ndof))
        rows = np.arange(len(ex))
        for c, (cx, cy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
            nd = self.node(ex + cx, ey + cy)
            for d in range(4):
                out[rows, 4 * nd + d] += vals[:, 4 * c + d]
        return out

    def basis_derivs_at(self, x: np.ndarray, y: np.ndarray):
        """Global shape-function second derivatives (xx, yy, xy) at points, each (npts, ndof)."""
        ex, ey, xi, eta = self._locate(x, y)
        _, Nxx, Nyy, Nxy = shape_functions(xi, eta, self.a, self.b)
        return self._scatter(Nxx, ex, ey), self._scatter(Nyy, ex, ey), self._scatter(Nxy, ex, ey)

    def basis_at(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Global shape-function values at points, shape (npts, ndof), dense."""
        n, a, b = self.n, self.a, self.b
        ex = np.clip(np.floor(x / a).astype(int), 0, n - 1)
        ey = np.clip(np.floor(y / b).astype(int), 0, n - 1)
        xi = x / a - ex
        eta = y / b - ey
        N, _, _, _ = shape_functions(xi, eta, a, b)
        out = np.zeros((len(x), self.ndof))
        rows = np.arange(len(x))
        for c, (cx, cy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]):
            nd = self.node(ex + cx, ey + cy)
            for d in range(4):
                out[rows, 4 * nd + d] += N[:, 4 * c + d]
        return out

    def modes_at(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        return self.basis_at(x, y) @ self.V

    # ------------------------------------------------------------ modal time response
    def _pulse_phase(self, t: np.ndarray):
        t = t[None, :]
        Om = self.Om
        e = np.exp(1j * Om * t)
        s1, s2, A, B = self.s1[:, None], self.s2[:, None], self.A[:, None], self.B[:, None]
        E1, E2 = np.exp(s1 * t), np.exp(s2 * t)
        Hc = self.Hc[:, None]
        q = (A * E1 + B * E2).real + self.qconst[:, None] - (Hc * e).real
        v = (A * s1 * E1 + B * s2 * E2).real - (1j * Om * Hc * e).real
        acc = (A * s1**2 * E1 + B * s2**2 * E2).real + (Om**2 * Hc * e).real
        rg = self.rigid
        if rg.any():
            al = self.alpha[rg][:, None]
            Ar, Gc, vc = self.Ar[rg][:, None], self.Gc[rg][:, None], self.vconst[rg][:, None]
            Ed = np.exp(-2.0 * al * t)
            q[rg] = Ar * (1.0 - Ed) / (2.0 * al) + vc * t - ((e - 1.0) * Gc / (1j * Om)).real
            v[rg] = Ar * Ed + vc - (Gc * e).real
            acc[rg] = -2.0 * al * Ar * Ed - (1j * Om * Gc * e).real
        return q, v, acc

    def _free_phase(self, tau: np.ndarray):
        tau = tau[None, :]
        s1, s2, A, B = self.s1[:, None], self.s2[:, None], self.A2[:, None], self.B2[:, None]
        E1, E2 = np.exp(s1 * tau), np.exp(s2 * tau)
        q = (A * E1 + B * E2).real
        v = (A * s1 * E1 + B * s2 * E2).real
        acc = (A * s1**2 * E1 + B * s2**2 * E2).real
        rg = self.rigid
        if rg.any():
            al = self.alpha[rg][:, None]
            qT, vT = self.qT[rg][:, None], self.vT[rg][:, None]
            Ed = np.exp(-2.0 * al * tau)
            q[rg] = qT + vT * (1.0 - Ed) / (2.0 * al)
            v[rg] = vT * Ed
            acc[rg] = -2.0 * al * vT * Ed
        return q, v, acc

    def response(self, t: np.ndarray):
        q = np.zeros((self.n_modes, len(t)))
        v = np.zeros_like(q)
        acc = np.zeros_like(q)
        m1 = (t >= 0) & (t <= self.T)
        m2 = t > self.T
        if m1.any():
            q[:, m1], v[:, m1], acc[:, m1] = self._pulse_phase(t[m1])
        if m2.any():
            q[:, m2], v[:, m2], acc[:, m2] = self._free_phase(t[m2] - self.T)
        return q, v, acc

    # ------------------------------------------------------------ Rayleigh integral
    def pressure(self, fs: int, nsamp: int) -> np.ndarray:
        p = self.p
        n, a, b = self.n, self.a, self.b
        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() * a
        yq = (ey.ravel()[:, None] + ETA.ravel()[None, :]).ravel() * b
        WW = np.tile((ws[:, None] * ws[None, :]).ravel() * a * b, n * n)
        r = np.sqrt((p["listener_x"] - xq) ** 2 + (p["listener_y"] - yq) ** 2 + p["listener_height"] ** 2)
        tau = r / p["c_air"]
        cq = p["rho_air"] / (2 * np.pi) * WW / r
        Phi = self.modes_at(xq, yq)  # (nq, n_modes)
        t = np.arange(nsamp) / fs
        out = np.zeros(nsamp)
        t_split = self.T + tau.max()
        early = (t < t_split) & (t >= tau.min())
        late = t >= t_split
        ne = int(early.sum())
        if ne:
            chunk = max(1, 4_000_000 // max(1, self.n_modes * ne))
            for i0 in range(0, len(xq), chunk):
                sl = slice(i0, i0 + chunk)
                tt = (t[early][None, :] - tau[sl][:, None]).ravel()
                _, _, acc = self.response(tt)
                acc = acc.reshape(self.n_modes, -1, ne)
                out[early] += np.einsum("q,qn,nqk->k", cq[sl], Phi[sl], acc, optimize=True)
        if late.any():
            tl = t[late] - self.T
            rg = self.rigid
            G1 = (cq[:, None] * Phi * np.exp(-np.outer(tau, self.s1))).sum(axis=0)
            G2 = (cq[:, None] * Phi * np.exp(-np.outer(tau, self.s2))).sum(axis=0)
            c1 = self.A2 * self.s1**2 * G1
            c2 = self.B2 * self.s2**2 * G2
            s1, s2 = self.s1.copy(), self.s2.copy()
            if rg.any():
                al = self.alpha[rg]
                Gr = (cq[:, None] * Phi[:, rg] * np.exp(np.outer(tau, 2.0 * al))).sum(axis=0)
                c1[rg] = -2.0 * al * self.vT[rg] * Gr
                c2[rg] = 0.0
                s1[rg] = -2.0 * al
                s2[rg] = 0.0
            coef = np.concatenate([c1, c2])
            rates = np.concatenate([s1, s2])
            res = np.zeros(len(tl))
            step = max(1, 20_000_000 // len(coef))
            for k0 in range(0, len(tl), step):
                tk = tl[k0:k0 + step]
                res[k0:k0 + step] = (np.exp(np.outer(tk, rates)) @ coef).real
            out[late] = res
        return out

    def field(self, x: np.ndarray, y: np.ndarray, t: np.ndarray) -> np.ndarray:
        XX, YY = np.meshgrid(x, y, indexing="xy")  # (ny, nx)
        Phi = self.modes_at(XX.ravel(), YY.ravel())
        q, _, _ = self.response(t)
        return (Phi @ q).T.reshape(len(t), len(y), len(x))


def main():
    ch = load_challenge(Path(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]
    fields_for = set()
    if "--fields" in args:
        i = args.index("--fields")
        fields_for = {int(v) for v in args[i + 1].split(",")}
        args = args[:i] + args[i + 2:]
    meshes = [int(a) for a in args] or [4, 6, 8, 12, 16, 24, 32]
    out = ch["output"]
    fs, nsamp = out["sample_rate_hz"], out["num_samples"]
    meta = []
    for n in meshes:
        t0 = time.perf_counter()
        plate = BFSPlate(ch, n)
        t1 = time.perf_counter()
        s = plate.pressure(fs, nsamp)
        dt = time.perf_counter() - t0
        sid = f"bfs-n{n}"
        out_dir.mkdir(parents=True, exist_ok=True)
        wavfile.write(out_dir / f"{sid}.wav", fs, (s / out["signal_scale"]).astype(np.float32))
        low = plate.omega[plate.omega > 0][:6] / (2 * np.pi)
        entry = {
            "solution_id": sid,
            "output": f"{out_dir.name}/{sid}.wav",
            "method": "Kirchhoff-Love plate, Bogner-Fox-Schmit C1 finite elements, closed-form modal time response, exact retarded Rayleigh integral",
            "parameters": {"elements_per_side": n, "element_size_m": round(plate.a, 6), "n_dof": plate.ndof, "n_modes": plate.n_modes, "rigid_modes": plate.n_rigid, "mode_cutoff_hz": MODE_CUTOFF_HZ},
            "timing_seconds": round(dt, 3),
            "peak_si": float(np.max(np.abs(s))),
            "lowest_frequencies_hz": [round(float(v), 3) for v in low],
        }
        if n in fields_for:
            import h5py

            fg = ch["fields"]
            nt, ny, nx = fg["shape"]
            x = np.linspace(0, ch["parameters"]["Lx"], nx)
            y = np.linspace(0, ch["parameters"]["Ly"], ny)
            tt = np.arange(nt) / fg["frame_rate_hz"]
            u = plate.field(x, y, tt)
            with h5py.File(out_dir / f"{sid}.h5", "w") as f:
                f.attrs["challenge_id"] = ch["challenge_id"]
                f.attrs["solution_id"] = sid
                f.attrs["quantity"] = fg["quantity"]
                f.attrs["units"] = fg["units"]
                f.create_dataset("x", data=x)
                f.create_dataset("y", data=y)
                f.create_dataset("t", data=tt)
                f.create_dataset("u", data=u.astype(np.float32), compression="gzip", compression_opts=4, chunks=(8, ny, nx))
            entry["fields"] = {"path": f"{out_dir.name}/{sid}.h5", "format": "hdf5"}
        meta.append(entry)
        print(f"n={n:3d} dof={plate.ndof:5d} modes={plate.n_modes:5d} rigid={plate.n_rigid} setup={t1 - t0:6.1f}s total={dt:6.1f}s peak={np.max(np.abs(s)):.4e} Pa  f_low={np.round(low, 2)}", file=sys.stderr)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()