"""Von Karman (Airy stress function) solution of cymbal-flat-disk-v1 by a modal method on a Cartesian grid.

The disk is approximated by the square cells of a uniform Cartesian grid whose centers lie inside
it (see bfs_disk.py). The transverse displacement is expanded in the linear free-plate modes Phi_k
of that staircase domain, computed with Bogner-Fox-Schmit elements, 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 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 cell; each point's retarded acceleration is a cubic Lagrange
interpolation of the modal accelerations at neighbouring time steps, folded into a per-step
weight table. Damping acts on the bending motion only (the in-plane state is quasi-static),
a further approximation of the challenge model.

The convergence chain has three axes: the grid spacing (8 mm, 4 mm), the number of modes kept
(transverse and Airy), and the time step. The mode truncation caps the frequency the nonlinear
cascade can reach.

Because the strike is centered on the x axis, the solution is even in y, and by default the modal bases
(transverse and Airy) are restricted to that symmetric subspace, which is exact and halves the number of
modes needed for a given frequency reach (the coupling cost scales with the cube of the count). --full
uses the unrestricted bases instead.

With --device cuda the coupling tensor is assembled and the time stepping runs on a GPU through PyTorch
(the modes and quadrature stay in numpy/scipy); the CPU path is the reference and the two agree to round-off.

Usage: python vk_modal_disk.py <challenge.json> <out-dir> [--fields tag,...] [--force F] [--duration T] [--seconds S] [--full] [--device cuda] [tag ...]
  where a tag is g<grid spacing in mm>-w<transverse modes>-a<airy modes>-m<oversampling>, e.g. g8-w200-a200-m1
"""
from __future__ import annotations

import json
import os
import sys
import time
from pathlib import Path

import numpy as np
from scipy.io import wavfile
sys.path.insert(0, str(Path(__file__).resolve().parent))
from bfs_disk import BFSDisk, load_challenge  # noqa: E402


class VonKarmanDisk:
    def __init__(self, ch: dict, spacing: float, n_w: int, n_airy: int, force_override: float | None = None, log=None, symmetric_y: bool | None = None, device: str = "cpu", duration_override: float | None = None):
        self.ch = ch
        self.device = device
        self.torch = None
        if device != "cpu":
            import torch

            self.torch = torch
        p = dict(ch["parameters"])
        if force_override is not None:
            p["strike_peak_force"] = force_override
        if duration_override is not None:
            p["strike_duration"] = duration_override
        ch2 = dict(ch)
        ch2["parameters"] = p
        self.p = p
        self.log = log or (lambda s: None)
        if symmetric_y is None:
            symmetric_y = p["strike_y"] == 0.0
        self.symmetric_y = symmetric_y
        t0 = time.perf_counter()
        self.plate = BFSDisk(ch2, spacing, n_w, symmetric_y=symmetric_y)  # linear modes, modal forces, alpha, omega
        pl = self.plate
        self.log(f"  linear modes: {pl.n_modes}{' (even in y)' if symmetric_y else ''} on {pl.n_cells} cells ({pl.ndof} dof) in {time.perf_counter() - t0:.1f}s, f_max={pl.omega[-1] / (2 * np.pi):.0f} Hz, rigid={pl.n_rigid}")
        self.n_w = pl.n_modes
        self.omega = pl.omega
        self.alpha = pl.alpha
        self.f = pl.f
        V = pl.V
        # Airy modes: clamped biharmonic eigenfunctions, L2-normalized: K_b psi = zeta^4 M_1 psi on interior dofs
        t0 = time.perf_counter()
        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
        z4, Psi = pl.eigenmodes(Kb, M1, n_airy, sigma=0.0, dofs=pl.interior_dofs)
        self.n_airy = len(z4)
        self.zeta4 = z4
        self.log(f"  airy modes: {self.n_airy} in {time.perf_counter() - t0:.1f}s")
        # coupling tensor H[p, i, j] by 3 x 3 Gauss quadrature per cell
        t0 = time.perf_counter()
        xq, yq, WQ = pl.quadrature(3)
        Bxx, Byy, Bxy = pl.basis_derivs_at(xq, yq)
        Pxx, Pyy, Pxy = Bxx @ V, Byy @ V, Bxy @ V  # (Q, n_w)
        Bn = pl.basis_at(xq, yq)
        PsiQ = Bn @ Psi  # (Q, n_airy)
        self.coef = -(p["E"] * p["h"]) / (2.0 * self.zeta4)
        if self.torch is None:
            H = np.zeros((self.n_airy, self.n_w, self.n_w))
            for pi in range(self.n_airy):
                wp = (WQ * PsiQ[:, pi])[:, None]
                A = (Pxx * wp).T @ Pyy
                H[pi] = A + A.T - 2.0 * (Pxy * wp).T @ Pxy
            self.H = H
            self.Hflat = H.reshape(self.n_airy * self.n_w, self.n_w)
            shape = H.shape
        else:
            # the same assembly on the device; the tensor stays there for the time stepping
            t = self.torch
            dev = t.device(self.device)
            Pxx_t, Pyy_t, Pxy_t = (t.as_tensor(a, device=dev) for a in (Pxx, Pyy, Pxy))
            WPsi = t.as_tensor(WQ[:, None] * PsiQ, device=dev)  # (Q, n_airy)
            H = t.empty((self.n_airy, self.n_w, self.n_w), dtype=t.float64, device=dev)
            for pi in range(self.n_airy):
                wp = WPsi[:, pi : pi + 1]
                A = (Pxx_t * wp).T @ Pyy_t
                H[pi] = A + A.T - 2.0 * (Pxy_t * wp).T @ Pxy_t
            t.cuda.synchronize() if self.device.startswith("cuda") else None
            self.H = None
            self.Hflat_t = H.reshape(self.n_airy * self.n_w, self.n_w)
            shape = tuple(H.shape)
            del Pxx_t, Pyy_t, Pxy_t, WPsi
        self.log(f"  coupling tensor {shape} on {self.device} in {time.perf_counter() - t0:.1f}s")
        # Rayleigh integral: quadrature weights, delays, modal shapes at the points
        _, _, self.tau, self.cq = pl.listener_quadrature(3)
        self.PhiQ = Bn @ 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 _prepare(self, fs: int, nsamp: int, m: int, field_grid, field_rate):
        p = self.p
        dt = 1.0 / (fs * m)
        w2, al = self.omega**2, self.alpha
        # 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 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,
            (fr + 1.0) * (fr - 1.0) * (fr - 2.0) / 2.0,
            -(fr + 1.0) * fr * (fr - 2.0) / 2.0,
            (fr + 1.0) * fr * (fr - 1.0) / 6.0,
        ], 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]
        fld = None
        if field_grid is not None:
            fx, fy, ft = field_grid
            XX, YY = np.meshgrid(fx, fy, indexing="xy")
            inside = self.plate.inside(XX.ravel(), YY.ravel())
            PhiG = self.plate.basis_at(XX.ravel(), YY.ravel()) @ self.plate.V
            fld = dict(PhiG=PhiG, inside=inside, every=int(round(fs * m / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx))
        return dict(dt=dt, T=p["strike_duration"], Om=2 * np.pi / p["strike_duration"], nsteps=nsamp * m, a1=a1, a2=a2, b=b, w2=w2, al=al, G=G, dmax=dmax, fld=fld)

    def run(self, fs: int, nsamp: int, m: int, field_grid=None, field_rate: float | None = None):
        pr = self._prepare(fs, nsamp, m, field_grid, field_rate)
        if self.torch is not None:
            return self._run_device(pr, m, nsamp)
        dt, T, Om, nsteps, a1, a2, b, w2, al, G, dmax, fld = (pr[k] for k in ("dt", "T", "Om", "nsteps", "a1", "a2", "b", "w2", "al", "G", "dmax", "fld"))
        f = self.f
        ring = np.zeros((dmax + 1, self.n_w))
        out = np.zeros(nsamp)
        q_prev = np.zeros(self.n_w)
        q = np.zeros(self.n_w)
        frames = [] if fld else None
        t_start = time.perf_counter()
        report = max(1, nsteps // 10)
        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
            q_next = a1 * q + a2 * q_prev + b * Fext
            qdd = Fext - w2 * q - 2.0 * al * (q_next - q_prev) / (2.0 * dt)
            ring[nstep % (dmax + 1)] = qdd
            if nstep % m == 0:
                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 % fld["every"] == 0 and len(frames) < fld["nt"]:
                fr_ = (fld["PhiG"] @ q).astype(np.float32)
                fr_[~fld["inside"]] = np.nan
                frames.append(fr_)
            q_prev, q = q, q_next
            if nstep % report == 0 and nstep:
                if not np.all(np.isfinite(q)):
                    raise RuntimeError(f"solution diverged at step {nstep} (t = {t:.4f} s); increase the oversampling m")
                el = time.perf_counter() - t_start
                self.log(f"  step {nstep}/{nsteps} t={t:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(q).max():.3e}")
        if not np.all(np.isfinite(out)):
            raise RuntimeError("output is not finite; increase the oversampling m")
        fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
        return out, fields

    def _run_device(self, pr, m: int, nsamp: int):
        """The same loop with the tensor contraction and the recurrence on the device (PyTorch)."""
        t = self.torch
        dev = t.device(self.device)
        T64 = lambda a: t.as_tensor(np.asarray(a, dtype=np.float64), device=dev)
        dt, T, Om, nsteps, dmax, fld = pr["dt"], pr["T"], pr["Om"], pr["nsteps"], pr["dmax"], pr["fld"]
        a1, a2, b, w2, al, G, f, coef = (T64(x) for x in (pr["a1"], pr["a2"], pr["b"], pr["w2"], pr["al"], pr["G"], self.f, self.coef))
        Hf = self.Hflat_t
        L = dmax + 1
        ar = t.arange(L, device=dev)
        ring = t.zeros((L, self.n_w), dtype=t.float64, device=dev)
        out = t.zeros(nsamp, dtype=t.float64, device=dev)
        q = t.zeros(self.n_w, dtype=t.float64, device=dev)
        q_prev = t.zeros_like(q)
        frames = [] if fld else None
        PhiG = T64(fld["PhiG"]) if fld else None
        invdt = 1.0 / dt  # 2 alpha (q+ - q-) / (2 dt)
        t_start = time.perf_counter()
        report = max(1, nsteps // 10)
        for nstep in range(nsteps + 1):
            tt = nstep * dt
            psi = 0.5 * (1.0 - np.cos(Om * tt)) if 0.0 <= tt <= T else 0.0
            v = (Hf @ q).reshape(self.n_airy, self.n_w)
            eta = coef * (v @ q)
            Fext = v.T @ eta
            if psi:
                Fext = Fext + f * psi
            q_next = a1 * q + a2 * q_prev + b * Fext
            qdd = Fext - w2 * q - al * ((q_next - q_prev) * invdt)
            ring[nstep % L] = qdd
            if nstep % m == 0:
                k = nstep // m
                if k < nsamp:
                    idx = (nstep - ar) % L
                    out[k] = t.sum(G * ring[idx])
            if frames is not None and nstep % fld["every"] == 0 and len(frames) < fld["nt"]:
                fr_ = (PhiG @ q).cpu().numpy().astype(np.float32)
                fr_[~fld["inside"]] = np.nan
                frames.append(fr_)
            q_prev, q = q, q_next
            if nstep % report == 0 and nstep:
                qc = q.cpu().numpy()
                if not np.all(np.isfinite(qc)):
                    raise RuntimeError(f"solution diverged at step {nstep} (t = {tt:.4f} s); increase the oversampling m")
                el = time.perf_counter() - t_start
                self.log(f"  step {nstep}/{nsteps} t={tt:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(qc).max():.3e}")
        out = out.cpu().numpy()
        if not np.all(np.isfinite(out)):
            raise RuntimeError("output is not finite; increase the oversampling m")
        fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
        return out, fields


def hardware_string(device: str) -> str:
    cpu = "unknown CPU"
    try:
        for line in open("/proc/cpuinfo"):
            if line.startswith("model name"):
                cpu = " ".join(line.split(":", 1)[1].split())
                break
    except OSError:
        pass
    s = f"{cpu} ({os.cpu_count()} threads), numpy/scipy for the setup"
    if device != "cpu":
        import torch

        s += f"; coupling tensor and time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
    else:
        s += " and the time stepping"
    return s


def parse_tag(tag: str):
    parts = dict((kv[0], int(kv[1:])) for kv in tag.split("-"))
    return parts["g"] * 1e-3, 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
    seconds = None

    def take(flag, conv):
        nonlocal args
        if flag in args:
            i = args.index(flag)
            v = conv(args[i + 1])
            args = args[:i] + args[i + 2:]
            return v
        return None

    v = take("--fields", lambda s: set(s.split(",")))
    fields_for = v or set()
    force = take("--force", float)
    duration = take("--duration", float)
    seconds = take("--seconds", float)
    device = take("--device", str) or "cpu"
    full = "--full" in args
    if full:
        args.remove("--full")
    tags = args or ["g8-w150-a150-m1", "g8-w300-a300-m1", "g4-w300-a300-m2", "g4-w400-a400-m2"]
    out = ch["output"]
    fs = out["sample_rate_hz"]
    nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
    log = lambda s: print(s, file=sys.stderr, flush=True)
    meta = []
    for tag in tags:
        spacing, n_w, n_a, m = parse_tag(tag)
        log(f"{tag}: force={force if force is not None else ch['parameters']['strike_peak_force']} N, duration={duration if duration is not None else ch['parameters']['strike_duration']} s")
        t0 = time.perf_counter()
        model = VonKarmanDisk(ch, spacing, n_w, n_a, force, log=log, symmetric_y=False if full else None, device=device, duration_override=duration)
        t1 = time.perf_counter()
        grid = None
        if tag in fields_for:
            fg = ch["fields"]
            nt, ny, nx = fg["shape"]
            R = ch["parameters"]["disk_radius"]
            grid = (np.linspace(-R, R, nx), np.linspace(-R, R, 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) on a Cartesian grid masked to the disk, modal method on Bogner-Fox-Schmit modes, Gautschi-type integrator exact in the linear part, retarded Rayleigh integral with cubic fractional delays",
            "parameters": {
                "grid_spacing_m": spacing,
                "cells": model.plate.n_cells,
                "n_dof": model.plate.ndof,
                "transverse_modes": model.n_w,
                "basis_even_in_y": model.symmetric_y,
                "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),
            "hardware": hardware_string(device),
            "peak_si": float(np.max(np.abs(s))),
            "setup_seconds": round(t1 - t0, 3),
        }
        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)
        log(f"{tag}: done, setup={t1 - t0:.1f}s total={dt_total:.1f}s peak={np.max(np.abs(s)):.4e} Pa")
        with open(out_dir / f"{sid}.meta.json", "w") as fh:
            json.dump(entry, fh, indent=1)
    print(json.dumps(meta, indent=1))


if __name__ == "__main__":
    main()