"""The discrete potential energy of the plate in NumPy (float64), whose negative
gradient is the force of the full model's kernel, and the largest frequency of
the linearized system, which sets the time step. Neither is used in a time loop;
they serve the tests and the set-up of a run."""

from __future__ import annotations

import numpy as np
import scipy.linalg

from ..grid import Grid
from ..linear.assemble import assemble, corners, shape_derivatives
from ..material import Material
from ..plate import Plate


def potential_energy(pl: Plate, u: np.ndarray, v: np.ndarray, w: np.ndarray) -> float:
    """Bending plus membrane energy [J] of the displacement fields ``(u, v, w)``
    ``(ny, nx)``, measured from the stress-free shape."""
    g, m = pl.grid, pl.material
    hx, hy, nu = g.hx, g.hy, m.nu
    left, right, down, up = g.neighbor_masks()
    wp = np.pad(w, 1)
    a = np.where(left & right, (wp[1:-1, 2:] - 2 * w + wp[1:-1, :-2]) / hx**2, 0.0)
    b = np.where(down & up, (wp[2:, 1:-1] - 2 * w + wp[:-2, 1:-1]) / hy**2, 0.0)
    Ub = (g.node_weight * m.bending_stiffness(pl.h) * (a * a + b * b + 2 * nu * a * b)).sum()
    cc = (w[1:, 1:] - w[1:, :-1] - w[:-1, 1:] + w[:-1, :-1]) / (hx * hy)
    Ub += (2 * (1 - nu) * m.bending_stiffness(pl.cell_thickness) * cc * cc * g.cell_mask).sum()
    Ub *= 0.5 * hx * hy

    hc = pl.cell_thickness
    C, Gs = m.membrane_stiffness(hc), m.shear_stiffness(hc)
    Uc, Vc, Wc, W0c = corners(u), corners(v), corners(w), corners(pl.w0)
    Um = 0.0
    q = hx * hy / 4
    for gp in range(4):
        dNdx, dNdy = shape_derivatives(gp, hx, hy)
        ux = sum(dNdx[k] * Uc[k] for k in range(4))
        uy = sum(dNdy[k] * Uc[k] for k in range(4))
        vx = sum(dNdx[k] * Vc[k] for k in range(4))
        vy = sum(dNdy[k] * Vc[k] for k in range(4))
        dwx = sum(dNdx[k] * Wc[k] for k in range(4))
        dwy = sum(dNdy[k] * Wc[k] for k in range(4))
        w0x = sum(dNdx[k] * W0c[k] for k in range(4))
        w0y = sum(dNdy[k] * W0c[k] for k in range(4))
        exx = ux + dwx * (w0x + 0.5 * dwx)
        eyy = vy + dwy * (w0y + 0.5 * dwy)
        gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy
        Um += (q * 0.5 * (C * (exx * exx + eyy * eyy + 2 * nu * exx * eyy) + Gs * gxy * gxy) * g.cell_mask).sum()
    return float(Ub + Um)


def max_angular_frequency(g: Grid, m: Material, h: float, n_small: int = 12) -> float:
    """Largest angular frequency of the linearized discrete system [rad/s], from
    the three-field stiffness of a small uniform flat plate with the same
    spacing and the thickness ``h`` (the largest of the plate, which bounds the
    bending frequency; the in-plane wave speed does not depend on the
    thickness). The extreme mode is a grid-scale mode, so neither the plate
    size nor its outline matters."""
    gs = Grid(n_small, n_small, g.hx, g.hy, np.ones((n_small, n_small), dtype=bool))
    sys = assemble(Plate.uniform(gs, m, h), fields=3)
    s = 1 / np.sqrt(sys.mass)
    A = sys.K.multiply(s[:, None]).multiply(s[None, :]).toarray()
    lam = scipy.linalg.eigvalsh(A, subset_by_index=(sys.n - 1, sys.n - 1))[0]
    return float(np.sqrt(max(lam, 0.0)))