"""Modal solution of the cymbal2 shallow-shell model on the same discretization as fem_shell.py.

The semi-discrete system is the one fem_shell.py builds: bicubic Hermite elements for u, v, w on the Cartesian grid,
3 x 3 Gauss points, lumped mass (with the same in-plane derivative-mass scaling). Here the linear part of that system
is diagonalized once, K phi = omega^2 M phi, and only the modes below a cutoff frequency f_max are kept; the six rigid
fields are excluded exactly, by a penalty in the eigenproblem that lifts them out of range (the rigid constraint of
MODEL.md). The state is the vector of modal amplitudes a(t), and the equations are

    a_k'' + 2 alpha_k a_k' + omega_k^2 a_k = F_k(a, a'),      alpha_k = sigma0 + eta omega_k^2 / 2,

where F is the modal projection of the nonlinear remainder of the internal force, evaluated at the Gauss points of the
mesh by the same routine as the finite element solver (so the nonlinear terms are not truncated to any tensor of modal
couplings), plus the stick loads. The time stepping is the two-step scheme of the v1 modal solver: the damped linear
oscillators are propagated exactly over each step and the force is applied with a static-exact weight, which is
second-order accurate and free of dispersion in the linear limit, so the step can be a small multiple of the output
sample. The sticks obey the Hertz law with the footprint-averaged displacement as in fem_shell.py. The Rayleigh
integral is a per-step weight table folded with a ring buffer of modal accelerations (cubic Lagrange interpolation in
time at the retarded time of every Gauss point).

The two solvers therefore solve the same discrete equations and differ only in the modal truncation and the time
integration; their difference measures the effect of the modes above f_max and of the step.

Usage:
  python modal_shell.py <challenge.json> <out-dir> [--spacing 0.008] [--fmax 20000] [--m 4] [--seconds S] [--device cuda]
                        [--fields] [--report 60] [--tag NAME] [--inplane-scale 16]
"""
from __future__ import annotations

import json
import math
import sys
import time
from pathlib import Path

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from scipy.io import wavfile

from fem_shell import ShellFEM, hardware_string


def sliced_modes(K, Mdiag, rigid, f_max: float, k_slice: int = 250, log=print, penalty: float = 1e12):
    """Every eigenpair of K phi = lambda M phi with lambda <= (2 pi f_max)^2, except the six rigid fields, which are
    lifted to the eigenvalue `penalty` by the rank-6 term M R R^T M (R the M-orthonormal rigid fields) and thereby
    made exactly orthogonal to every returned mode. Shift-invert ARPACK in slices of k_slice eigenvalues, each slice
    factorizing K - sigma M once (the low-rank term is handled with the Woodbury identity)."""
    n = K.shape[0]
    M = sp.diags(Mdiag).tocsc()
    if rigid is None:  # a clamped shell has no rigid fields: plain shift-invert
        rigid = np.zeros((0, n))
    U = (Mdiag[:, None] * rigid.T) * math.sqrt(penalty)  # (n, r): penalty term is U U^T
    nr = U.shape[1]
    Kop = spla.LinearOperator((n, n), matvec=lambda x: K @ x + U @ (U.T @ x), dtype=np.float64)
    target = (2 * math.pi * f_max) ** 2
    vals_all, vecs_all = [], []
    lo = -np.inf
    sigma = -1e3
    width = None
    tries = 0
    while True:
        t0 = time.perf_counter()
        lu = spla.splu((K - sigma * M).tocsc())
        AiU = lu.solve(U) if nr else np.zeros((n, 0))  # (n, r)
        S = np.linalg.inv(np.eye(nr) + U.T @ AiU) if nr else np.zeros((0, 0))

        def opinv(x, lu=lu, AiU=AiU, S=S):
            y = lu.solve(x)
            return y - AiU @ (S @ (U.T @ y))

        OPinv = spla.LinearOperator((n, n), matvec=opinv, dtype=np.float64)
        vals, vecs = spla.eigsh(Kop, k=k_slice, M=M, sigma=sigma, which="LM", OPinv=OPinv, tol=1e-9)
        order = np.argsort(vals)
        vals, vecs = vals[order], vecs[:, order]
        tries += 1
        if np.isfinite(lo) and vals[0] > lo * (1 + 1e-9) + 1e-3:
            if tries > 40:
                raise RuntimeError("mode slicing did not close the gap")
            sigma = 0.5 * (sigma + lo)  # the slice did not reach back to the covered range: step back
            log(f"    slice at {math.sqrt(max(sigma, 0)) / 2 / math.pi:.0f} Hz left a gap, retrying")
            continue
        keep = vals > lo if np.isfinite(lo) else np.ones(len(vals), dtype=bool)
        vals_all.append(vals[keep])
        vecs_all.append(vecs[:, keep])
        lo = vals[-1]
        width = vals[-1] - vals[0]
        log(f"    slice: {int(keep.sum())} new modes up to {math.sqrt(max(lo, 0)) / 2 / math.pi:.0f} Hz in {time.perf_counter() - t0:.1f}s")
        if lo >= target:
            break
        sigma = lo + 0.45 * width
    vals = np.concatenate(vals_all)
    vecs = np.concatenate(vecs_all, axis=1)
    sel = (vals <= target) & (vals < 0.5 * penalty)
    vals, vecs = vals[sel], vecs[:, sel]
    # rigid content check and clean-up of the near-null modes of the flat geometry (exact zeros the penalty removed)
    rc = np.abs(rigid @ (Mdiag[:, None] * vecs)).max() if nr else 0.0
    log(f"    {len(vals)} modes below {f_max:g} Hz" + (f"; largest rigid component {rc:.1e}" if nr else " (clamped: no rigid fields)"))
    return np.sqrt(np.maximum(vals, 0.0)), vecs


class ModalShell:
    def __init__(self, fem: ShellFEM, f_max: float, log=print):
        self.fem = fem
        self.log = log
        t0 = time.perf_counter()
        K = fem.assemble_linear()
        self.log(f"  linear stiffness assembled ({K.nnz} nonzeros) in {time.perf_counter() - t0:.1f}s")
        t0 = time.perf_counter()
        cache_dir = Path(__file__).resolve().parent / "cache"
        cache_dir.mkdir(exist_ok=True)
        import hashlib

        geom_key = hashlib.sha1(json.dumps({k: v for k, v in fem.ch["parameters"].items() if k in ("disk_radius", "hole_radius", "mount_radius", "bow_height", "bell_radius", "bell_height", "thickness", "thickness_center", "thickness_rim", "E", "nu", "rho")}, sort_keys=True).encode()).hexdigest()[:8]
        cache = cache_dir / f"modes-{fem.ch['challenge_id']}-g{fem.a * 1e3:g}-f{f_max / 1e3:g}k-s{fem.inplane_scale:g}-{geom_key}.npz"
        if cache.exists():
            z = np.load(cache)
            self.omega, self.Phi = z["omega"], z["Phi"]
            self.log(f"  modes loaded from {cache.name}")
        else:
            if fem.clamped:  # eigenproblem on the free unknowns only; the modes are zero at the mount
                fr = fem.free_dofs
                om, Pf = sliced_modes(K[fr][:, fr].tocsr(), fem.Mdiag[fr], None, f_max, log=log)
                Phi = np.zeros((fem.n_dof, len(om)))
                Phi[fr] = Pf
                self.omega, self.Phi = om, Phi
            else:
                self.omega, self.Phi = sliced_modes(K, fem.Mdiag, fem.rigid, f_max, log=log)
            np.savez(cache, omega=self.omega, Phi=self.Phi)
        self.n_modes = len(self.omega)
        self.log(f"  {self.n_modes} modes in {time.perf_counter() - t0:.1f}s: f = {np.round(self.omega[:5] / 2 / np.pi, 2)} ... {self.omega[-1] / 2 / np.pi:.0f} Hz")
        p = fem.p
        self.alpha = p["sigma0"] + 0.5 * p["eta"] * self.omega**2

    def run(self, fs: int, nsamp: int, m: int, device: str = "cpu", field_grid=None, field_rate=None, report: float = 60.0,
            out_dir: Path | None = None, tag: str = "modal", scale: float = 1.0, compile_: bool = True, use_graph: bool = True):
        import torch

        fem, p = self.fem, self.fem.p
        t_setup = time.perf_counter()
        dev = torch.device(device)
        force, ops = fem.build_force(torch, dev, compile_)
        T64, T64i = ops["T64"], ops["T64i"]
        N = self.n_modes
        dt = 1.0 / (fs * m)
        nsteps = nsamp * m
        Phi = T64(self.Phi)  # (n_dof, N)
        PhiT = Phi.T.contiguous()
        w2 = T64(self.omega**2)
        al = T64(self.alpha)
        eta, sigma0 = p["eta"], p["sigma0"]
        ms, kc, v0 = p["stick_mass"], p["contact_stiffness"], fem.v0
        # exact two-step recurrence of the damped linear oscillators
        disc = np.sqrt(self.alpha**2 - self.omega**2 + 0j)
        E1, E2 = np.exp((-self.alpha + disc) * dt), np.exp((-self.alpha - disc) * dt)
        a1 = T64((E1 + E2).real)
        a2 = T64(-(E1 * E2).real)
        b = T64((1.0 - (E1 + E2).real + (E1 * E2).real) / self.omega**2)
        # Rayleigh integral: per-step weight table G[d, k] from the modes at the Gauss points
        tau = fem.ray_tau.reshape(-1)
        cq = fem.ray_c.reshape(-1)
        tau_steps = tau / dt
        d0 = np.floor(tau_steps).astype(np.int64)
        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
        L = dmax + 1
        Nw = T64(fem.Nw)  # (G, 16)
        Psi = torch.einsum("gi,cin->cgn", Nw, Phi[ops["cell_dofs_w"]]).reshape(-1, N)  # (C*G, N) w of the modes at the Gauss points
        Psi = Psi - (T64(np.tile(fem.wq, fem.n_cells) / fem.area) @ Psi)[None, :]  # remove each mode's area mean: the output of MODEL.md
        Gt = torch.zeros((L, N), dtype=torch.float64, device=dev)
        for j, off in enumerate([-1, 0, 1, 2]):
            Gt.index_add_(0, T64i(d0 + off), T64(cq * lag[:, j])[:, None] * Psi)
        del Psi
        # sticks
        nst = len(fem.strikes)
        gains = T64(np.stack([s["phi"] for s in fem.strikes])) @ Phi  # (nst, N): footprint average of the modes
        n_on = T64i([int(round(s["t"] / dt)) for s in fem.strikes])
        z = torch.zeros(nst, dtype=torch.float64, device=dev)
        zd = torch.zeros(nst, dtype=torch.float64, device=dev)
        active = torch.zeros(nst, dtype=torch.float64, device=dev)
        Fst = torch.zeros(nst, dtype=torch.float64, device=dev)
        imp = torch.zeros(nst, dtype=torch.float64, device=dev)
        dur = torch.zeros(nst, dtype=torch.float64, device=dev)
        fmax_t = torch.zeros(nst, dtype=torch.float64, device=dev)
        # state
        a = torch.zeros(N, dtype=torch.float64, device=dev)
        a_prev = torch.zeros(N, dtype=torch.float64, device=dev)
        q = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
        q_prev = torch.zeros(fem.n_dof, dtype=torch.float64, device=dev)
        ring = torch.zeros((L, N), dtype=torch.float64, device=dev)
        out = torch.zeros(nsamp, dtype=torch.float64, device=dev)
        nstep = torch.zeros((), dtype=torch.int64, device=dev)
        ar = torch.arange(L, device=dev)
        fld = None
        if field_grid is not None:
            fx, fy, ft = field_grid
            idx, wts, inside = fem.field_interpolator(fx, fy)
            Wg = (T64(wts)[:, :, None] * Phi[T64i(idx)]).sum(1)  # (P, N)
            fld = dict(Wg=Wg, inside=inside, every=int(round(fs / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx), frames=[])

        def step():
            n = nstep
            q.copy_(Phi @ a)
            v = (q - q_prev) / dt
            adot = (a - a_prev) / dt
            # sticks
            wbar = gains @ a
            wbar_dot = gains @ adot
            starting = (n == n_on).to(torch.float64)
            z.copy_(z * (1 - starting) + wbar * starting)
            zd.copy_(zd * (1 - starting) + (wbar_dot - v0) * starting)
            active.copy_(torch.maximum(active, starting))
            pen = torch.clamp(wbar - z, min=0.0)
            F = active * kc * pen * torch.sqrt(pen)
            Fst.copy_(F)
            imp.add_(dt * F)
            dur.add_(dt * (F > 0).to(torch.float64))
            fmax_t.copy_(torch.maximum(fmax_t, F))
            # modal force: nonlinear remainder of the internal force plus the sticks
            f = force(q, v)
            Fm = -(PhiT @ f) + w2 * a + eta * w2 * adot - F @ gains
            a_next = a1 * a + a2 * a_prev + b * Fm
            add = Fm - w2 * a - al * ((a_next - a_prev) / dt)
            ring.index_copy_(0, (n % L).view(1), add.view(1, -1))
            zd.add_(dt * F / ms)
            z.add_(dt * zd)
            q_prev.copy_(q)
            a_prev.copy_(a)
            a.copy_(a_next)
            nstep.add_(1)

        def sample(k):
            # sample k is at step k m: the first step below computes and stores the acceleration there, the output sample
            # folds the table with the ring (delays counted back from that step), then the remaining steps of the sample
            step()
            idx = (nstep - 1 - ar) % L
            out.index_copy_(0, k.view(1), torch.sum(Gt * ring[idx]).view(1))
            for _ in range(m - 1):
                step()

        graph = None
        k_t = torch.zeros((), dtype=torch.int64, device=dev)

        def sample_graph():
            sample(k_t)

        def reset():
            for t in (z, zd, active, Fst, imp, dur, fmax_t, a, a_prev, q, q_prev, ring, out, nstep, k_t):
                t.zero_()

        if use_graph and device.startswith("cuda"):
            s = torch.cuda.Stream()
            s.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(s):
                for _ in range(2):
                    sample_graph()
            torch.cuda.current_stream().wait_stream(s)
            reset()
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                sample_graph()
            reset()
        self.log(f"  setup {time.perf_counter() - t_setup:.1f}s; stepping {nsteps} steps ({nsamp / fs:.2f} s, m = {m}) on {device}{' with a CUDA graph' if graph else ''}")

        def take_frame():
            fr_ = (fld["Wg"] @ a).cpu().numpy().astype(np.float32)
            fr_[~fld["inside"]] = np.nan
            fld["frames"].append(fr_)

        stick_stats = lambda: dict(impulse_Ns=[round(float(x), 5) for x in imp.cpu().numpy()], contact_ms=[round(1e3 * float(x), 4) for x in dur.cpu().numpy()],
                                   peak_force_N=[round(float(x), 2) for x in fmax_t.cpu().numpy()], rebound_mps=[round(float(x), 4) for x in zd.cpu().numpy()])
        t_start = time.perf_counter()
        last_report = t_start
        if fld is not None:
            take_frame()
        for k in range(nsamp):
            k_t.fill_(k)
            if graph is not None:
                graph.replay()
            else:
                sample(k_t)
            if fld is not None and (k + 1) % fld["every"] == 0 and len(fld["frames"]) < fld["nt"]:
                take_frame()
            now = time.perf_counter()
            if (now - last_report >= report or k == nsamp - 1) and k > 0:
                last_report = now
                el = now - t_start
                an = a.cpu().numpy()
                if not np.all(np.isfinite(an)):
                    raise RuntimeError(f"solution diverged at sample {k} (t = {k / fs:.4f} s); increase m")
                sig = out.cpu().numpy()
                prog = dict(tag=tag, step=(k + 1) * m, steps=nsteps, t=round((k + 1) / fs, 6), duration=nsamp / fs, elapsed_s=round(el, 1),
                            eta_s=round(el * (nsamp - k - 1) / (k + 1), 1), steps_per_s=round((k + 1) * m / el, 1), max_abs_a=float(np.abs(an).max()),
                            stick_z_m=[round(float(x), 6) for x in z.cpu().numpy()], stick_v_mps=[round(float(x), 4) for x in zd.cpu().numpy()],
                            contact_force_N=[round(float(x), 2) for x in Fst.cpu().numpy()], **stick_stats(), peak_pa=float(np.abs(sig).max()), m=m, dt=dt, device=device, n_modes=N)
                self.log(f"  t={prog['t']:.3f}s elapsed {el:.0f}s eta {prog['eta_s']:.0f}s ({prog['steps_per_s']:.0f} steps/s) max|a|={prog['max_abs_a']:.3e} "
                         f"peak={prog['peak_pa']:.3f} Pa sticks z={prog['stick_z_m']} F={prog['contact_force_N']}")
                if out_dir is not None:
                    with open(out_dir / f"{tag}.progress.json", "w") as fh:
                        json.dump(prog, fh, indent=1)
                    wavfile.write(out_dir / f"{tag}.partial.wav", fs, (sig / scale).astype(np.float32))
        elapsed = time.perf_counter() - t_start
        sig = out.cpu().numpy()
        if not np.all(np.isfinite(sig)):
            raise RuntimeError("output is not finite; increase m")
        fields = np.stack(fld["frames"]).reshape(len(fld["frames"]), fld["ny"], fld["nx"]) if fld is not None else None
        return sig, fields, dict(m=m, dt=dt, elapsed=elapsed, sticks=stick_stats())


def main():
    ch = json.load(open(sys.argv[1]))
    out_dir = Path(sys.argv[2])
    args = sys.argv[3:]

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

    spacing = take("--spacing", float, 0.008)
    f_max = take("--fmax", float, 20000.0)
    m = take("--m", int, 4)
    seconds = take("--seconds", float)
    device = take("--device", str, "cpu")
    report = take("--report", float, 60.0)
    speed = take("--speed", float)
    tag = take("--tag", str)
    inplane_scale = take("--inplane-scale", float, 1.0)
    fields = "--fields" in args
    compile_ = "--no-compile" not in args
    use_graph = "--no-graph" not in args
    out = ch["output"]
    fs = out["sample_rate_hz"]
    nsamp = out["num_samples"] if seconds is None else int(round(fs * seconds))
    tag = tag or f"modal-g{spacing * 1e3:g}-f{f_max / 1e3:g}k-m{m}"
    log = lambda s: print(s, file=sys.stderr, flush=True)
    out_dir.mkdir(parents=True, exist_ok=True)
    t0 = time.perf_counter()
    log(f"{tag}: {ch['challenge_id']}, spacing {spacing * 1e3:g} mm, modes to {f_max:g} Hz, m = {m}, {nsamp / fs:.2f} s, device {device}")
    fem = ShellFEM(ch, spacing, log=log, speed=speed, no_drift=True, inplane_scale=inplane_scale)
    model = ModalShell(fem, f_max, log=log)
    t1 = time.perf_counter()
    scale = out.get("signal_scale") or 1.0
    grid = None
    if fields:
        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"])
    sig, fld, info = model.run(fs, nsamp, m, device=device, field_grid=grid, field_rate=ch["fields"]["frame_rate_hz"] if grid else None,
                               report=report, out_dir=out_dir, tag=tag, scale=scale, compile_=compile_, use_graph=use_graph)
    total = time.perf_counter() - t0
    wavfile.write(out_dir / f"{tag}.wav", fs, (sig / scale).astype(np.float32))
    meta = {
        "solution_id": tag,
        "output": f"{out_dir.name}/{tag}.wav",
        "method": "modes of the bicubic Hermite finite element system on a Cartesian grid masked to the annulus (lumped mass), truncated at a cutoff frequency, "
                  "exact linear propagation with the nonlinear force evaluated at the Gauss points, Hertz stick, Rayleigh integral by a retarded weight table",
        "parameters": {"grid_spacing_m": spacing, "cells": fem.n_cells, "n_dof": fem.n_dof, "n_modes": model.n_modes, "mode_cutoff_hz": f_max,
                       "highest_mode_hz": round(float(model.omega.max() / 2 / np.pi), 1), "oversampling": m, "dt_s": info["dt"], "inplane_derivative_mass_scale": inplane_scale},
        "timing_seconds": round(info["elapsed"], 3),
        "setup_seconds": round(t1 - t0, 3),
        "hardware": hardware_string(device),
        "peak_si": float(np.max(np.abs(sig))),
        "sticks": info["sticks"],
        **({"stick_speed_override": speed} if speed is not None else {}),
    }
    if fld is not None:
        import h5py

        fg = ch["fields"]
        with h5py.File(out_dir / f"{tag}.h5", "w") as fh:
            fh.attrs["challenge_id"] = ch["challenge_id"]
            fh.attrs["solution_id"] = tag
            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][: fld.shape[0]])
            fh.create_dataset("u", data=fld, compression="gzip", compression_opts=4, chunks=(8, fld.shape[1], fld.shape[2]))
        meta["fields"] = {"path": f"{out_dir.name}/{tag}.h5", "format": "hdf5"}
    with open(out_dir / f"{tag}.meta.json", "w") as fh:
        json.dump(meta, fh, indent=1)
    log(f"{tag}: done in {total:.0f}s, peak {meta['peak_si']:.4f} Pa, sticks {info['sticks']}")


if __name__ == "__main__":
    main()