"""Linear check of the assembly: the lowest natural frequencies of the linearized shell with the lumped and the
consistent mass, from the same element operators the time stepping uses. For the flat annulus the free disk's
values are 21.9 (2,0), 38.0 (0,0), 50.9 (3,0), 85.9 (1,0), 89.6 (4,0) Hz (Leissa), the 6 mm hole changing them
by well under a percent. Usage: python check_modes.py <challenge.json> [spacing_m] [n_modes]"""
import json
import sys

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

from fem_shell import ShellFEM, element_basis

ch = json.load(open(sys.argv[1]))
a = float(sys.argv[2]) if len(sys.argv) > 2 else 0.008
k = int(sys.argv[3]) if len(sys.argv) > 3 else 14
fem = ShellFEM(ch, a, log=lambda s: print(s))
p = fem.p
C, G, nu = fem.n_cells, fem.G, p["nu"]
B = fem.B.reshape(G, 9, 48)
# L: 9 derivative quantities -> [exx, eyy, gxy, kxx, kyy, 2kxy] at the reference state
L = np.zeros((C, G, 6, 9))
L[:, :, 0, 0] = fem.Zx; L[:, :, 0, 5] = 1.0
L[:, :, 1, 1] = fem.Zy; L[:, :, 1, 8] = 1.0
L[:, :, 2, 0] = fem.Zy; L[:, :, 2, 1] = fem.Zx; L[:, :, 2, 6] = 1.0; L[:, :, 2, 7] = 1.0
L[:, :, 3, 2] = 1.0; L[:, :, 4, 3] = 1.0; L[:, :, 5, 4] = 2.0
Q = np.array([[1, nu, 0], [nu, 1, 0], [0, 0, (1 - nu) / 2]])
Dm = np.zeros((C, G, 6, 6))
Dm[:, :, :3, :3] = fem.Cm[:, :, None, None] * Q
Dm[:, :, 3:, 3:] = fem.Db[:, :, None, None] * Q
LB = np.einsum("cgij,gjk->cgik", L, B)  # (C, G, 6, 48)
Ke = np.einsum("g,cgik,cgij,cgjl->ckl", fem.wq, LB, Dm, LB, optimize=True)
rows = np.repeat(fem.cell_dofs, 48, axis=1).ravel()
cols = np.tile(fem.cell_dofs, (1, 48)).ravel()
K = sp.coo_matrix((Ke.ravel(), (rows, cols)), shape=(fem.n_dof, fem.n_dof)).tocsr()
K = 0.5 * (K + K.T)
# consistent mass
N = element_basis(fem.sx, fem.sy, a)["N"]  # (G, 16)
Me16 = np.einsum("cg,gi,gj->cij", p["rho"] * fem.hq * fem.wq[None, :], N, N)  # (C, 16, 16)
Me = np.zeros((C, 48, 48))
for b in range(3):
    Me[:, 16 * b:16 * b + 16, 16 * b:16 * b + 16] = Me16
Mc = sp.coo_matrix((Me.ravel(), (rows, cols)), shape=(fem.n_dof, fem.n_dof)).tocsr()
Ml = sp.diags(fem.Mdiag)
if fem.clamped:
    fr = fem.free_dofs
    K, Mc, Ml = K[fr][:, fr].tocsr(), Mc[fr][:, fr].tocsr(), sp.diags(fem.Mdiag[fr])
    print(f"clamped at the mount: {len(fem.fixed_dofs)} unknowns fixed")
print(f"total mass: consistent {Mc.sum():.5f}, lumped {fem.Mdiag[0:4 * fem.n_nodes:4].sum():.5f} (w block)")
for name, M in [("lumped", Ml), ("consistent", Mc)]:
    vals = spla.eigsh(K, k=k, M=M.tocsc(), sigma=-1e3, which="LM", return_eigenvectors=False)
    f = np.sqrt(np.maximum(np.sort(vals), 0)) / (2 * np.pi)
    print(f"{name:10s}: " + " ".join(f"{x:7.2f}" for x in f))
# highest frequency with the lumped mass (stability)
vals = spla.eigsh(K, k=1, M=Ml.tocsc(), which="LM", return_eigenvectors=False)
print(f"omega_max (lumped) = {np.sqrt(vals[0]):.4e} rad/s")
# which block sets the explicit step: bending-only and in-plane-only top frequencies
nw = 4 * fem.n_nodes
Kw = K[:nw, :nw]; Kp = K[nw:, nw:]
for name, Kb, Mb in [("bending block", Kw, sp.diags(fem.Mdiag[:nw])), ("in-plane block", Kp, sp.diags(fem.Mdiag[nw:]))]:
    v = spla.eigsh(Kb, k=1, M=Mb.tocsc(), which="LM", return_eigenvectors=False)
    print(f"omega_max {name}: {np.sqrt(v[0]):.4e} rad/s")
# and the bending block with the derivative-dof masses scaled up by 4 (rotational mass scaling)
Ms = fem.Mdiag[:nw].copy(); Ms[1::4] *= 4; Ms[2::4] *= 4; Ms[3::4] *= 4
v = spla.eigsh(Kw, k=1, M=sp.diags(Ms).tocsc(), which="LM", return_eigenvectors=False)
print(f"omega_max bending with derivative masses x4: {np.sqrt(v[0]):.4e} rad/s")
# in-plane block with its derivative-dof masses scaled, and the effect on the low modes of the whole system
for sc in (4, 16, 64):
    Ms = fem.Mdiag.copy()
    for b in (1, 2):
        base = nw * b
        Ms[base + 1::4][: fem.n_nodes] *= sc; Ms[base + 2::4][: fem.n_nodes] *= sc; Ms[base + 3::4][: fem.n_nodes] *= sc
    Msp = sp.diags(Ms)
    v = spla.eigsh(Kp, k=1, M=sp.diags(Ms[nw:]).tocsc(), which="LM", return_eigenvectors=False)
    low = spla.eigsh(K, k=14, M=Msp.tocsc(), sigma=-1e3, which="LM", return_eigenvectors=False)
    fl = np.sqrt(np.maximum(np.sort(low), 0)) / (2 * np.pi)
    print(f"in-plane derivative masses x{sc}: omega_max in-plane {np.sqrt(v[0]):.4e} rad/s; low modes " + " ".join(f"{x:6.2f}" for x in fl[6:]))
# lowest in-plane modes with and without the derivative-mass scaling (in-plane block alone; 6 -> 3 rigid modes there)
for sc in (1, 16):
    Ms = fem.Mdiag[nw:].copy()
    for b in (0, 1):
        base = nw * b
        Ms[base + 1::4][: fem.n_nodes] *= sc; Ms[base + 2::4][: fem.n_nodes] *= sc; Ms[base + 3::4][: fem.n_nodes] *= sc
    low = spla.eigsh(Kp, k=9, M=sp.diags(Ms).tocsc(), sigma=-1e3, which="LM", return_eigenvectors=False)
    fl = np.sqrt(np.maximum(np.sort(low), 0)) / (2 * np.pi)
    print(f"in-plane modes, derivative masses x{sc}: " + " ".join(f"{x:7.1f}" for x in fl[3:]))