r"""
QUAD4 shell reference geometry — Mindlin–Reissner shape functions
=================================================================

The 4-node flat Mindlin–Reissner shell maps to the natural
coordinates :math:`(\xi, \eta) \in [-1, 1]^2` and carries six
DOFs per node — three translations
:math:`(u_x, u_y, u_z)` plus three rotations
:math:`(\theta_x, \theta_y, \theta_z)` — for 24 DOFs per element.

In-plane shape functions (bilinear Lagrange):

.. math::

    N_i(\xi, \eta) = \tfrac{1}{4}(1 + \xi_i\xi)(1 + \eta_i\eta),
    \qquad i = 1, \ldots, 4.

The first-order shear-deformation theory (Mindlin 1951;
Reissner 1945) assumes the kinematic split

.. math::

    u_x = u_x^0 + z\, \theta_y, \quad
    u_y = u_y^0 - z\, \theta_x, \quad
    u_z = u_z^0,

so the bending-membrane coupling is governed by independent
rotation DOFs and the through-thickness shear strain is constant
across :math:`z`.

Integration:

* **Membrane + bending** — 2×2 Gauss-Legendre (4 points), exact
  for the bilinear :math:`B^T D B` integrand on the reference
  square.
* **Transverse shear** — 1×1 selective-reduced Gauss (1 point at
  the centre), per Malkus & Hughes (CMAME 1978), to suppress
  shear locking on thin shells.

Drilling-DOF stabilisation: a per-node penalty
:math:`\alpha\, G\, t\, A` (Allman 1984; Hughes-Brezzi 1989)
removes the local rank deficiency when :math:`\theta_z` is free.

References
----------
* Mindlin, R. D. (1951) "Influence of rotatory inertia and shear
  on flexural motions of isotropic elastic plates," *J. Applied
  Mechanics* 18, 31–38.
* Reissner, E. (1945) "The effect of transverse shear deformation
  on the bending of elastic plates," *J. Applied Mechanics* 12,
  A-69–A-77.
* Malkus, D. S. and Hughes, T. J. R. (1978) "Mixed finite element
  methods — Reduced and selective integration techniques: A
  unification of concepts," *CMAME* 15 (1), 63–81.
* Allman, D. J. (1984) "A compatible triangular element including
  vertex rotations for plane elasticity analysis," *Computers &
  Structures* 19 (1-2), 1–8.

Implementation: :class:`femorph_solver.elements.quad4_shell.Quad4Shell`.
"""

from __future__ import annotations

import numpy as np
import pyvista as pv

# %%
# Reference quad with both quadrature rules
# -----------------------------------------

corners = np.array([[-1.0, -1.0, 0.0], [+1.0, -1.0, 0.0], [+1.0, +1.0, 0.0], [-1.0, +1.0, 0.0]])
cells = np.hstack([[4], np.arange(4, dtype=np.int64)])
cell_types = np.array([pv.CellType.QUAD], dtype=np.uint8)
ref_quad = pv.UnstructuredGrid(cells, cell_types, corners)

# 2x2 Gauss for membrane + bending
g = 1.0 / np.sqrt(3.0)
gauss_4 = np.array([[a, b, 0.0] for a in (-g, +g) for b in (-g, +g)])

# 1x1 selective-reduced for transverse shear
gauss_1 = np.array([[0.0, 0.0, 0.0]])

# %%
# Render

plotter = pv.Plotter(off_screen=True, window_size=(640, 480))
plotter.add_mesh(
    ref_quad, style="wireframe", color="black", line_width=2.5, label="reference QUAD4 shell"
)
plotter.add_points(
    corners, render_points_as_spheres=True, point_size=18, color="black", label="corner nodes (4)"
)
plotter.add_points(
    gauss_4,
    render_points_as_spheres=True,
    point_size=14,
    color="#d62728",
    label="2×2 membrane / bending",
)
plotter.add_points(
    gauss_1,
    render_points_as_spheres=True,
    point_size=22,
    color="#9467bd",
    label="1×1 transverse shear",
)
plotter.view_xy()
plotter.camera.zoom(1.3)
plotter.add_legend(face=None, size=(0.27, 0.14), bcolor="white")
plotter.show()

# %%
# Sanity — N_i at all 5 quadrature locations sums to 1

xi_i = np.array([-1.0, +1.0, +1.0, -1.0])
eta_i = np.array([-1.0, -1.0, +1.0, +1.0])
all_pts = np.vstack([gauss_4, gauss_1])
sums = []
for q in all_pts:
    sums.append(sum(0.25 * (1 + xi_i[i] * q[0]) * (1 + eta_i[i] * q[1]) for i in range(4)))
sums = np.array(sums)
np.testing.assert_allclose(sums, 1.0, atol=1e-14)
print(f"5 quadrature locations; max |Σ N_i - 1| = {np.max(np.abs(sums - 1)):.2e}")
print("OK — partition of unity verified at every Gauss point.")