"""Reference implementation (Python, cymsim) of the modal precompute that cymplay does in the browser.

Used only to validate the browser numbers; the app itself needs no Python.

The in-plane DOFs (u, v) of the linearised shell are eliminated by static condensation, which
leaves a dense stiffness on the (w, theta_x, theta_y) DOFs. Its mass-normalised eigenmodes below
``--fmax`` give, for each mode m: the frequency f_m, the decay rate sigma_m (sigma0 plus the
projected sigma1 term), the complex listener gains (front Rayleigh sum and diffracted back sum,
with the element delays folded into a phase at f_m), and the transverse mode shape on the nodes,
quantised to int16 with one scale per mode. The mesh, the nodal surface height, the full config
and the exciter defaults are stored alongside so that the app can locate strike points, run the
contact law and write a config back out.

Output format: magic ``CYMB``, a uint32 little-endian header length, a UTF-8 JSON header, then
the raw little-endian arrays at the byte offsets the header lists.

Usage (on a machine with enough memory for a dense (3N)^2 float64 matrix, 1.5 GB at 6 mm):
    python make_modes.py ../cym2/examples/crash_mallet.yaml public/modes/crash_mallet.bin
"""
from __future__ import annotations

import argparse
import dataclasses
import json
import struct
import sys
import time

import numpy as np
import scipy.linalg as sla
import scipy.sparse as sp
import scipy.sparse.linalg as spla

try:
    import shim  # noqa: F401  (office machine: Python 3.9 compatibility, adds cym2/src to the path)
except ImportError:
    pass

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from cymsim.config import MalletExciter, SimulationConfig, StickExciter   # noqa: E402
from cymsim.elements import element_energy                                # noqa: E402
from cymsim.geometry import CymbalProfile                                 # noqa: E402
from cymsim.mesh import disk_mesh                                         # noqa: E402
from cymsim.radiation import RadiationModel                               # noqa: E402
from cymsim.shell import N_DOF, ShellModel                                # noqa: E402

FORMAT_VERSION = 2


def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


def build(cfg: SimulationConfig, fmax: float, condense_rotations: bool = False):
    t0 = time.time()
    prof = CymbalProfile(cfg.profile, cfg.imperfections)
    mesh = disk_mesh(cfg.profile.radius, cfg.mesh.element_size, cfg.mesh.smoothing)
    sh = ShellModel(mesh, prof, cfg.material, cfg.numerics)
    N, E, tri = mesh.n_nodes, mesh.n_elements, mesh.tri
    log(f"mesh: {N} nodes, {E} elements; shell built in {time.time() - t0:.0f} s")

    # ---- global tangent stiffness at rest from the element Hessians (float64)
    Kh = np.asarray(jax.vmap(jax.hessian(element_energy), in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
        jnp.zeros((E, 3, N_DOF)), jnp.asarray(mesh.Gx), jnp.asarray(mesh.Gy), jnp.asarray(mesh.area),
        jnp.asarray(sh.z0x), jnp.asarray(sh.z0y), jnp.asarray(sh.Cm_e), jnp.asarray(sh.Ke),
        sh.nl)).reshape(E, 15, 15)
    gi = (tri[:, :, None] * 5 + np.arange(5)[None, None, :]).reshape(E, 15)
    rows = np.repeat(gi, 15, axis=1).ravel()
    cols = np.tile(gi, (1, 15)).ravel()
    K = sp.coo_matrix((Kh.ravel(), (rows, cols)), shape=(5 * N, 5 * N)).tocsr()
    M = sh.M.ravel()
    log(f"assembled K: {K.nnz} nonzeros")

    # ---- static condensation of the in-plane DOFs. Three in-plane DOFs are pinned to remove
    #      the in-plane rigid motions; the condensed operator then has exactly three zero modes
    #      (transverse translation and the two tilts), which we drop.
    if condense_rotations:
        ip = np.sort(np.concatenate([np.arange(N) * 5 + d for d in (0, 1, 3, 4)]))
    else:
        ip = np.sort(np.concatenate([np.arange(N) * 5, np.arange(N) * 5 + 1]))
    wb = np.setdiff1d(np.arange(5 * N), ip)
    # pinned in-plane DOFs u0, v0, v1 in the local numbering of ``ip``
    pins = [int(np.searchsorted(ip, g)) for g in (0, 1, 5 + 1)]
    Kuu = K[ip][:, ip].tolil()
    Kuw = K[ip][:, wb].toarray()
    Kww = K[wb][:, wb].toarray()
    for p in pins:                            # u0, v0, v1
        Kuu[p, :] = 0
        Kuu[:, p] = 0
        Kuu[p, p] = 1.0
        Kuw[p, :] = 0
    t0 = time.time()
    lu = spla.splu(Kuu.tocsc())
    X = lu.solve(Kuw)                         # (2N, 3N)
    S = Kww - Kuw.T @ X
    S = 0.5 * (S + S.T)
    del Kww, Kuw, X
    log(f"condensed stiffness {S.shape} in {time.time() - t0:.0f} s")

    # ---- mass-normalised eigenmodes below fmax (with some margin, trimmed afterwards)
    Mw = M[wb]
    s = 1.0 / np.sqrt(Mw)
    S *= s[:, None]
    S *= s[None, :]
    t0 = time.time()
    lam, V = sla.eigh(S, subset_by_value=[-np.inf, (2 * np.pi * fmax) ** 2], driver="evr",
                      overwrite_a=True, check_finite=False)
    del S
    f = np.sqrt(np.maximum(lam, 0)) / (2 * np.pi)
    keep = f > 1.0                             # drop the rigid modes
    f, V = f[keep], V[:, keep]
    Phi = V * s[:, None]                       # Phi^T M Phi = I on the (w, thx, thy) DOFs
    Nm = len(f)
    log(f"eigh: {Nm} modes between {f[0]:.1f} Hz and {f[-1]:.0f} Hz in {time.time() - t0:.0f} s; "
        f"dropped {int((~keep).sum())} rigid modes")
    stride = 1 if condense_rotations else 3
    Phi_w = np.ascontiguousarray(Phi[0::stride, :])     # (N, Nm), the w rows

    # ---- modal damping: sigma_m = sigma0 + sigma1 phi^T L phi with L = sum_e (rho h A)_e G^T G
    rhoA = cfg.material.rho * sh.h_e * mesh.area
    Lrows, Lcols, Lvals = [], [], []
    for i in range(3):
        for j in range(3):
            Lrows.append(tri[:, i]); Lcols.append(tri[:, j])
            Lvals.append(rhoA * (mesh.Gx[:, i] * mesh.Gx[:, j] + mesh.Gy[:, i] * mesh.Gy[:, j]))
    L = sp.coo_matrix((np.concatenate(Lvals), (np.concatenate(Lrows), np.concatenate(Lcols))),
                      shape=(N, N)).tocsr()
    damp_q = np.einsum("nm,nm->m", Phi_w, L @ Phi_w)
    sigma = cfg.material.sigma0 + cfg.material.sigma1 * damp_q
    log(f"damping: sigma from {sigma.min():.2f} to {sigma.max():.1f} 1/s "
        f"(T60 {6.91 / sigma.max() * 1e3:.0f} ms at the top mode)")

    # ---- listener gains: Rayleigh front sum and diffracted back sum per mode, with the element
    #      delays as phases at the mode frequency. pressure = Re(G_m * complex acceleration).
    rad = RadiationModel(mesh, prof, sh.z0_node, cfg.listener, 1e-6, 1000, jnp.float64)
    phi_e = Phi_w[tri].mean(1)                            # (E, Nm) element-mean mode shape
    c = cfg.listener.c_air
    omega = 2 * np.pi * f
    G_front = np.einsum("em,e,em->m", phi_e, np.asarray(rad.coef_f),
                        np.exp(-1j * np.outer(rad.r_front, omega) / c))
    G_back = np.einsum("em,e,em->m", phi_e, np.asarray(rad.coef_b),
                       np.exp(-1j * np.outer(rad.r_back, omega) / c))
    log("radiation gains done")

    # ---- quantise the mode shapes: int16 with one scale per mode
    scale = np.abs(Phi_w).max(0) / 32767.0
    Phi_q = np.round(Phi_w / scale[None, :]).astype(np.int16)

    arrays = {
        "xy": mesh.xy.astype(np.float32),                 # (N, 2)
        "tri": tri.astype(np.int32),                      # (E, 3)
        "z0": sh.z0_node.astype(np.float32),              # (N,)
        "h": sh.h_node.astype(np.float32),                # (N,)
        "freq": f.astype(np.float64),                     # (Nm,)
        "damp_q": damp_q.astype(np.float64),              # (Nm,) sigma = sigma0 + sigma1 * damp_q
        "gain_front": np.stack([G_front.real, G_front.imag], 1).astype(np.float32),   # (Nm, 2)
        "gain_back": np.stack([G_back.real, G_back.imag], 1).astype(np.float32),     # (Nm, 2)
        "phi_scale": scale.astype(np.float32),            # (Nm,)
        "phi_q": Phi_q,                                   # (N, Nm) int16, node-major
    }
    meta = {
        "format_version": FORMAT_VERSION,
        "n_nodes": int(N), "n_elements": int(E), "n_modes": int(Nm), "fmax": float(fmax),
        "radius": float(cfg.profile.radius),
        "E0": float(cfg.material.E), "rho0": float(cfg.material.rho),
        "mass": float(sh.mass),
        "config": cfg.to_dict(),
        "playback_element_size": float(cfg.mesh.element_size),
        "condense_rotations": bool(condense_rotations),
        "exciters": {"stick": dataclasses.asdict(StickExciter()),
                     "mallet": dataclasses.asdict(MalletExciter())},
        "notes": "phi = phi_q * phi_scale[m] is the mass-normalised transverse mode shape on the "
                 "nodes; modal EOM: a'' + 2 sigma a' + (2 pi f)^2 a = phi(p) . F; pressure = "
                 "Re(sum_m G_m * complex acceleration_m), unbaffled = front - back.",
    }
    return meta, arrays


def write_bin(path: str, meta: dict, arrays: dict[str, np.ndarray]) -> None:
    offset = 0
    entries = {}
    blobs = []
    for name, a in arrays.items():
        a = np.ascontiguousarray(a)
        if a.dtype.byteorder == ">":
            a = a.astype(a.dtype.newbyteorder("<"))
        b = a.tobytes()
        entries[name] = {"dtype": a.dtype.name, "shape": list(a.shape), "offset": offset, "nbytes": len(b)}
        blobs.append(b)
        offset += len(b)
    header = json.dumps({**meta, "arrays": entries}).encode("utf-8")
    with open(path, "wb") as fh:
        fh.write(b"CYMB")
        fh.write(struct.pack("<I", len(header)))
        fh.write(header)
        for b in blobs:
            fh.write(b)
    log(f"wrote {path}: {(8 + len(header) + offset) / 1e6:.1f} MB")


def read_bin(path: str):
    """Reader for the same format (used by the validation script)."""
    with open(path, "rb") as fh:
        data = fh.read()
    assert data[:4] == b"CYMB", "not a cymplay modes file"
    (hlen,) = struct.unpack("<I", data[4:8])
    meta = json.loads(data[8:8 + hlen].decode("utf-8"))
    base = 8 + hlen
    arrays = {}
    for name, e in meta.pop("arrays").items():
        arrays[name] = np.frombuffer(data, dtype=np.dtype(e["dtype"]).newbyteorder("<"),
                                     count=int(np.prod(e["shape"])), offset=base + e["offset"]).reshape(e["shape"])
    return meta, arrays


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("config", help="cymsim YAML config")
    ap.add_argument("out", help="output .bin file")
    ap.add_argument("--fmax", type=float, default=20000.0, help="keep modes below this frequency [Hz]")
    ap.add_argument("--element-size", type=float, help="override mesh.element_size [m]")
    ap.add_argument("--condense-rotations", action="store_true",
                    help="also condense theta_x, theta_y (w-only modes, as the browser does)")
    ap.add_argument("--no-imperfections", action="store_true",
                    help="zero every seeded imperfection (deterministic geometry, comparable across RNGs)")
    args = ap.parse_args(argv)
    cfg = SimulationConfig.from_yaml(args.config)
    if args.element_size:
        cfg.mesh.element_size = args.element_size
    if args.no_imperfections:
        im = cfg.imperfections
        im.n_dents, im.rim_warp, im.shoulder_ovality, im.bell_offset = 0, 0.0, 0.0, 0.0
    cfg.validate()
    meta, arrays = build(cfg, args.fmax, args.condense_rotations)
    write_bin(args.out, meta, arrays)


if __name__ == "__main__":
    sys.exit(main())