"""The discrete shell: element data, lumped mass, rigid-body modes, internal force and the
stable time-step bound.

Given a mesh and a CymbalProfile, this module samples the thickness and the surface slope on the
elements, builds the DKT bending stiffness and the membrane stiffness of every element, lumps the
mass, and exposes the internal nodal force f_int(q) = -dU/dq obtained by reverse-mode automatic
differentiation of the summed element strain energies. It also provides the two operators the
integrator needs besides f_int: the sigma1 damping force on w and the projection that removes
the six rigid-body modes of the free shell. See docs/model.md, sections "Spatial discretization"
and "Time integration".
"""
from __future__ import annotations

import logging
import math
from dataclasses import dataclass

import jax
import jax.numpy as jnp
import numpy as np

from .config import Material, Numerics
from .elements import dkt_stiffness, element_energy, element_energy_parts
from .geometry import CymbalProfile
from .mesh import Mesh

N_DOF = 5   # (u, v, w, theta_x, theta_y) per node

log = logging.getLogger(__name__)


def global_omega_max(K_e: np.ndarray, M: np.ndarray, tri: np.ndarray, tol: float = 1e-7,
                     max_iter: int = 20000) -> tuple[float, int, bool]:
    """sqrt(lambda_max(M^-1 K)) of the stiffness assembled from the element matrices ``K_e``
    (E, 15, 15), node-major DOFs, by power iteration on M^-1/2 K M^-1/2 (lumped M, (N, 5)).

    The start vector x_i = sin(i + 1) and the stopping rule (the Rayleigh quotient changes by at
    most ``tol`` relative over 10 iterations) are fixed, so that cymsimweb, which implements the
    same iteration, arrives at the same value. The Rayleigh quotient approaches lambda_max from
    below. Returns (omega, iterations, converged)."""
    E, L = len(tri), M.size
    dof = (N_DOF * tri[:, :, None] + np.arange(N_DOF)).reshape(E, 15)
    isq = 1.0 / np.sqrt(M.ravel())
    x = np.sin(np.arange(L) + 1.0)
    x /= np.linalg.norm(x)
    hist = []
    for it in range(1, max_iter + 1):
        ye = np.einsum("eab,eb->ea", K_e, (x * isq)[dof])
        y = np.bincount(dof.ravel(), weights=ye.ravel(), minlength=L) * isq
        lam = float(x @ y)
        hist.append(lam)
        x = y / np.linalg.norm(y)
        if it >= 20 and it % 10 == 0 and abs(lam - hist[-11]) <= tol * lam:
            return math.sqrt(lam), it, True
    return math.sqrt(hist[-1]), max_iter, False


@dataclass
class ElementChecks:
    """Residuals of the element self-tests; all should be small (1e-6 or better)."""
    rigid_tilt_residual: float          # |K_e U_tilt| / |K_e| for a rigid tilt w = y, thx = 1
    constant_curvature_error: float     # relative error of the energy of w = x^2 / 2
    zero_energy_modes_min: int          # zero eigenvalues of M_e^-1 K_e per element (expect 6)
    zero_energy_modes_max: int
    rigid_mode_ratio: float             # max |K r| / |K random| over the 6 rigid modes of the shell


class ShellModel:
    """Assembled shell on a mesh. All device arrays are in the simulation precision ``dtype``."""

    def __init__(self, mesh: Mesh, profile: CymbalProfile, material: Material, numerics: Numerics):
        self.mesh = mesh
        self.material = material
        self.numerics = numerics
        self.dtype = jnp.float32 if numerics.dtype == "float32" else jnp.float64
        self.nl = 1.0 if numerics.nonlinear else 0.0
        xy, tri = mesh.xy, mesh.tri
        N, E = mesh.n_nodes, mesh.n_elements
        x, y = xy[tri, 0], xy[tri, 1]
        area = mesh.area

        # ---- profile sampled on the mesh: thickness per element and node, height per node.
        #      The slope that enters the membrane strain is the CST gradient of the nodal
        #      heights, i.e. the initial surface is represented in the same piecewise-linear
        #      space as w. This keeps a rigid tilt of the shell exactly force-free; the analytic
        #      slope at the centroid would not (see docs/model.md).
        self.h_e = profile.thickness_at(mesh.centroid)
        self.h_node = profile.thickness_at(xy)
        self.z0_node = profile.z0_at(xy)
        z0_e = self.z0_node[tri]
        self.z0x = (mesh.Gx * z0_e).sum(1)
        self.z0y = (mesh.Gy * z0_e).sum(1)
        self.max_slope = float(np.hypot(self.z0x, self.z0y).max())
        # ---- unit normal of the rest surface at the nodes, from the area-weighted mean of the
        #      CST slopes of the adjacent elements; +z side is the front (bell up). The normal
        #      acceleration n . (u_dd, v_dd, w_dd) is what is recorded and what radiates.
        gx, gy = np.zeros(N), np.zeros(N)
        for k in range(3):
            np.add.at(gx, tri[:, k], area * self.z0x)
            np.add.at(gy, tri[:, k], area * self.z0y)
        a_node = np.zeros(N)
        for k in range(3):
            np.add.at(a_node, tri[:, k], area)
        gx, gy = gx / a_node, gy / a_node
        nrm = np.sqrt(1 + gx ** 2 + gy ** 2)
        self.normal_node = np.column_stack([-gx / nrm, -gy / nrm, 1 / nrm])

        # ---- element stiffness: bending rigidity D, membrane stiffness C, DKT matrix
        E_, nu, rho = material.E, material.nu, material.rho
        self.D_e = E_ * self.h_e ** 3 / (12 * (1 - nu ** 2))
        self.Cm_e = (E_ * self.h_e / (1 - nu ** 2))[:, None] * np.array([1.0, nu, 0.5 * (1 - nu)])
        self.Ke = np.asarray(jax.vmap(dkt_stiffness, in_axes=(0, 0, 0, None))(
            jnp.asarray(x), jnp.asarray(y), jnp.asarray(self.D_e), nu))

        # ---- lumped mass per DOF: a third of the element mass to each node; rotational inertia
        #      I = m (h^2 / 12 + f A_e), the second term being the mass-scaling device discussed
        #      in the documentation
        m_e = rho * self.h_e * area / 3
        I_e = m_e * (self.h_e ** 2 / 12 + numerics.rot_inertia_factor * area)
        M = np.zeros((N, N_DOF))
        for k in range(3):
            np.add.at(M[:, 0], tri[:, k], m_e * numerics.inplane_mass_scale)
            np.add.at(M[:, 1], tri[:, k], m_e * numerics.inplane_mass_scale)
            np.add.at(M[:, 2], tri[:, k], m_e)
            np.add.at(M[:, 3], tri[:, k], I_e)
            np.add.at(M[:, 4], tri[:, k], I_e)
        self.M = M
        self.m_e, self.I_e = m_e, I_e
        self.mass = float(M[:, 2].sum())

        # ---- rigid-body modes of the curved shell, M-orthonormalised (see docs). A tilt drags
        #      the in-plane displacements along: rotation about x by an angle a gives v = -a z0,
        #      w = a y; rotation about y gives u = +a z0, w = -a x (first order, shallow shell).
        X, Y = xy[:, 0] - xy[:, 0].mean(), xy[:, 1] - xy[:, 1].mean()
        Rm = np.zeros((6, N, N_DOF))
        Rm[0, :, 0] = 1
        Rm[1, :, 1] = 1
        Rm[2, :, 0], Rm[2, :, 1] = -Y, X
        Rm[3, :, 2] = 1
        Rm[4, :, 2], Rm[4, :, 3], Rm[4, :, 1] = Y, 1, -self.z0_node
        Rm[5, :, 2], Rm[5, :, 4], Rm[5, :, 0] = -X, 1, self.z0_node
        for k in range(6):
            for l in range(k):
                Rm[k] -= (Rm[l] * M * Rm[k]).sum() * Rm[l]
            Rm[k] /= math.sqrt((Rm[k] * M * Rm[k]).sum())
        self.Rm = Rm

        # ---- stability bound (Irons): lambda_max(M^-1 K) <= max_e lambda_max(M_e^-1 K_e), with
        #      K_e the tangent stiffness at rest (Hessian of the element energy at q = 0)
        Khess = 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(area),
            jnp.asarray(self.z0x), jnp.asarray(self.z0y), jnp.asarray(self.Cm_e), jnp.asarray(self.Ke),
            self.nl).reshape(E, 15, 15)
        Me = np.tile(np.stack([m_e * numerics.inplane_mass_scale, m_e * numerics.inplane_mass_scale,
                               m_e, I_e, I_e], 1), (1, 3))               # (E, 15), node-major
        S = np.asarray(Khess) / np.sqrt(Me[:, :, None] * Me[:, None, :])
        self._lam = np.linalg.eigvalsh(S)                                 # (E, 15) ascending
        self.limiting_element = int(np.argmax(self._lam[:, -1]))
        self.omega_max = math.sqrt(self._lam[:, -1].max())
        # ---- the global lambda_max(M^-1 K) of the assembled tangent, which the element bound
        #      overestimates (by about 1.45 at 6 mm elements); see docs/model.md
        self.K_rest = np.asarray(Khess)                                   # (E, 15, 15) tangent at rest
        self.omega_global, self.omega_global_iterations, converged = global_omega_max(self.K_rest, M, tri)
        if numerics.stability_bound == "global" and not converged:
            log.warning("the power iteration for the global stability bound did not converge in %d "
                        "iterations; using the element bound", self.omega_global_iterations)
        use_global = numerics.stability_bound == "global" and converged
        self.omega_step = min(self.omega_global, self.omega_max) if use_global else self.omega_max
        self.dt_stable = 2.0 / self.omega_step

        # ---- device arrays in simulation precision
        f = lambda a: jnp.asarray(a, dtype=self.dtype)
        self.tri_j = jnp.asarray(tri, dtype=jnp.int32)
        self.Gx_j, self.Gy_j, self.area_j = f(mesh.Gx), f(mesh.Gy), f(area)
        self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j = f(self.z0x), f(self.z0y), f(self.Cm_e), f(self.Ke)
        self.rhoh_area_j = f(rho * self.h_e * area)
        self.h_node_j = f(self.h_node)
        self.normal_j = f(self.normal_node)
        self.M_j, self.Minv_j, self.Rm_j = f(M), f(1.0 / M), f(Rm)
        self._sigma1 = float(material.sigma1)
        self._N = N
        self._f_int = jax.grad(lambda q: -self.total_energy(q))

    # ---- energies and forces (plain JAX functions; the integrator jits the whole step)
    def energy_parts_per_element(self, q):
        """(membrane, bending) energy of every element for nodal state q (N, 5)."""
        return jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
            q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j,
            self.Cm_j, self.Ke_j, self.nl)

    def linear_energy_parts(self, q):
        """(membrane, bending) strain energy of q without the von Karman terms; both quadratic forms."""
        Em, Eb = jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
            q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j, 0.0)
        return Em.sum(), Eb.sum()

    def linear_energy(self, q):
        """Strain energy of q with the quadratic (von Karman) strain terms dropped: a quadratic
        form, so that sqrt(kinetic_energy(v) + linear_energy(q)) is a norm on the state."""
        Em, Eb = jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
            q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j, 0.0)
        return Em.sum() + Eb.sum()

    def total_energy(self, q):
        Em, Eb = self.energy_parts_per_element(q)
        return Em.sum() + Eb.sum()

    def internal_force(self, q):
        """f_int = -dU/dq, shape (N, 5)."""
        return self._f_int(q)

    def kinetic_energy(self, v):
        return 0.5 * (self.M_j * v * v).sum()

    def damping_force_w(self, vw):
        """The sigma1 term as a nodal force on w: -2 sigma1 rho h Laplacian(w_dot) integrated
        with the CST gradients, i.e. 2 sigma1 sum_e (rho h A)_e G^T (G w_dot_e)."""
        ve = vw[self.tri_j]
        vx, vy = (self.Gx_j * ve).sum(1), (self.Gy_j * ve).sum(1)
        c = self.rhoh_area_j[:, None] * (self.Gx_j * vx[:, None] + self.Gy_j * vy[:, None])
        return 2 * self._sigma1 * jax.ops.segment_sum(c.ravel(), self.tri_j.ravel(), num_segments=self._N)

    def project_rigid(self, v):
        """Remove the M-orthogonal projection onto the six rigid-body modes."""
        coef = jnp.einsum("knd,nd->k", self.Rm_j, self.M_j * v)
        return v - jnp.einsum("k,knd->nd", coef, self.Rm_j)

    # ---- diagnostics
    def element_checks(self) -> ElementChecks:
        """Element self-tests (float64) and the rigid-mode test of the assembled tangent."""
        mesh = self.mesh
        x, y = mesh.xy[mesh.tri, 0], mesh.xy[mesh.tri, 1]
        E = mesh.n_elements
        U_tilt = np.stack([y, np.ones_like(y), np.zeros_like(y)], 2).reshape(E, 9)
        res = np.abs(np.einsum("eij,ej->ei", self.Ke, U_tilt)).max() / np.abs(self.Ke).max()
        U_cur = np.stack([0.5 * x ** 2, np.zeros_like(x), -x], 2).reshape(E, 9)
        Eb = 0.5 * np.einsum("ei,eij,ej->e", U_cur, self.Ke, U_cur)
        err = np.abs(Eb / (0.5 * self.D_e * mesh.area) - 1).max()
        n_zero = (self._lam < 1e-9 * self._lam[:, -1:]).sum(1)

        q0 = jnp.zeros((mesh.n_nodes, N_DOF), self.dtype)
        K_apply = lambda d: jax.jvp(self._f_int, (q0,), (d,))[1]
        rng = np.random.default_rng(0)
        worst = 0.0
        for k in range(6):
            rnd = rng.standard_normal((mesh.n_nodes, N_DOF))
            rnd *= np.linalg.norm(self.Rm[k]) / np.linalg.norm(rnd)
            worst = max(worst, float(jnp.abs(K_apply(self.Rm_j[k])).max()
                                     / jnp.abs(K_apply(jnp.asarray(rnd, self.dtype))).max()))
        return ElementChecks(float(res), float(err), int(n_zero.min()), int(n_zero.max()), worst)