"""Check that the internal force of fem_shell is the gradient of the shell's strain energy (numpy re-implementation of
the same Gauss-point evaluation), on the crash geometry at a random state of millimeter amplitude."""
import json, sys
import numpy as np
from fem_shell import ShellFEM
ch = json.load(open(sys.argv[1])); a = float(sys.argv[2]) if len(sys.argv) > 2 else 0.008
fem = ShellFEM(ch, a, log=lambda s: None)
nu = fem.p["nu"]; C, G = fem.n_cells, fem.G
B = fem.B  # (81, 48)
def strains(q):
    D = (q[fem.cell_dofs] @ B.T).reshape(C, G, 9)
    wx, wy = D[..., 0], D[..., 1]
    exx = D[..., 5] + fem.Zx * wx + 0.5 * wx**2
    eyy = D[..., 8] + fem.Zy * wy + 0.5 * wy**2
    exy = 0.5 * (D[..., 6] + D[..., 7] + fem.Zx * wy + fem.Zy * wx + wx * wy)
    return D, exx, eyy, exy
def energy(q):
    D, exx, eyy, exy = strains(q)
    Um = 0.5 * fem.Cm * (exx**2 + eyy**2 + 2 * nu * exx * eyy + 2 * (1 - nu) * exy**2)
    Ub = 0.5 * fem.Db * (D[..., 2]**2 + D[..., 3]**2 + 2 * nu * D[..., 2] * D[..., 3] + 2 * (1 - nu) * D[..., 4]**2)
    return float(np.sum((Um + Ub) * fem.wq[None, :]))
def force(q):
    D, exx, eyy, exy = strains(q)
    wx, wy = D[..., 0], D[..., 1]; ax, ay = fem.Zx + wx, fem.Zy + wy
    Nxx = fem.Cm * (exx + nu * eyy); Nyy = fem.Cm * (eyy + nu * exx); Nxy = fem.Cm * (1 - nu) * exy
    Mxx = fem.Db * (D[..., 2] + nu * D[..., 3]); Myy = fem.Db * (D[..., 3] + nu * D[..., 2]); Mxy = fem.Db * (1 - nu) * D[..., 4]
    S = np.stack([Nxx * ax + Nxy * ay, Nyy * ay + Nxy * ax, Mxx, Myy, 2 * Mxy, Nxx, Nxy, Nxy, Nyy], axis=-1) * fem.wq[None, :, None]
    fe = S.reshape(C, G * 9) @ B
    f = np.zeros(fem.n_dof); np.add.at(f, fem.cell_dofs.ravel(), fe.ravel()); return f
rng = np.random.default_rng(1)
q = np.zeros(fem.n_dof)
# smooth random state: w of a few mm, in-plane of 0.1 mm, derivative dofs scaled accordingly
xy = fem.node_xy; R = fem.p["disk_radius"]
for blk, amp in [(0, 3e-3), (1, 1e-4), (2, 1e-4)]:
    for k in range(6):
        kx, ky, ph = rng.uniform(-40, 40, 2).tolist() + [rng.uniform(0, 2 * np.pi)]
        val = amp * np.cos(kx * xy[:, 0] + ky * xy[:, 1] + ph)
        base = blk * 4 * fem.n_nodes + 4 * np.arange(fem.n_nodes)
        q[base] += val; q[base + 1] += -amp * kx * np.sin(kx * xy[:, 0] + ky * xy[:, 1] + ph)
        q[base + 2] += -amp * ky * np.sin(kx * xy[:, 0] + ky * xy[:, 1] + ph); q[base + 3] += -amp * kx * ky * np.cos(kx * xy[:, 0] + ky * xy[:, 1] + ph)
f = force(q)
print(f"energy {energy(q):.6e} J, |f| {np.linalg.norm(f):.4e}")
for trial in range(4):
    d = rng.standard_normal(fem.n_dof) * np.abs(q).mean() * 1e-3
    d[4 * fem.n_nodes:] *= 0.03
    for eps in (1e-2, 1e-3, 1e-4):
        num = (energy(q + eps * d) - energy(q - eps * d)) / (2 * eps)
        ana = f @ d
        print(f"  dir {trial} eps {eps:g}: dE/ds numerical {num:.10e} analytic {ana:.10e} rel diff {abs(num - ana) / abs(ana):.2e}")