r"""
Shear-locking demonstration — HEX8 integration variants
========================================================

The textbook trap of low-order solid elements: a fully-
integrated 8-node trilinear hex (HEX8 with full 2 × 2 × 2
Gauss, **no** B-bar / enhanced strain) **cannot** reproduce
linear-bending kinematics, because the constant-strain
Gauss point at the element centre overstates the bending
stiffness (Cook §6.6, §6.13).  The result is **shear locking**:
even on a slender cantilever where the Euler-Bernoulli
analytical answer is exact, the FE displacement is too small —
sometimes by orders of magnitude — and the missing strain
energy goes into spurious transverse shear at the Gauss point.

Two practitioner cures, both shipped with femorph-solver's
:doc:`/reference/elements/hex8`:

* **B-bar volumetric stabilisation** (Hughes 1980, default
  ``ELEMENTS.HEX8``): averages the volumetric component of the
  strain field across the element, suppressing the related
  *volumetric* locking at :math:`\nu \to 0.5` but not the
  shear locking on its own.

* **Enhanced assumed strain** (EAS — Simo–Rifai 1990,
  ``ELEMENTS.HEX8(integration="enhanced_strain")``): adds nine
  static-condensed internal modes that fix bending-mode shear
  locking exactly.  The natural choice for thin-bending solid
  geometries.

This example solves the same cantilever-tip-load problem with
**three integration variants** side by side and prints the tip-
deflection ratio against the Euler-Bernoulli closed form
:math:`\delta = P L^{3} / (3 E I)`:

* ``ELEMENTS.HEX8(integration="plain_gauss")`` — full 2×2×2
  Gauss, no stabilisation (the textbook locking case).
* ``ELEMENTS.HEX8`` — full Gauss + B-bar (default).
* ``ELEMENTS.HEX8(integration="enhanced_strain")`` — EAS.

References
----------

* Cook, R. D., Malkus, D. S., Plesha, M. E., Witt, R. J.
  (2002) *Concepts and Applications of Finite Element
  Analysis*, 4th ed., Wiley, §6.6 (volumetric locking),
  §6.13 (shear locking + EAS), §17.2 (low-order shell
  cure).
* Hughes, T. J. R. (1980) "Generalization of selective
  integration procedures to anisotropic and nonlinear media,"
  *International Journal for Numerical Methods in
  Engineering* 15 (9), 1413–1418 (B-bar).
* Simo, J. C. and Rifai, M. S. (1990) "A class of mixed
  assumed strain methods and the method of incompatible
  modes," *International Journal for Numerical Methods in
  Engineering* 29 (8), 1595–1638 (EAS).
* Bathe, K.-J. (2014) *Finite Element Procedures*, 2nd ed.,
  §5.3.4 (incompatible modes / EAS).
* Belytschko, T., Liu, W. K., Moran, B. and Elkhodary,
  K. (2014) *Nonlinear Finite Elements for Continua and
  Structures*, 2nd ed., Wiley, §8.4 (locking modes).
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv

import femorph_solver
from femorph_solver import ELEMENTS

# %%
# Problem data
# ------------

E = 2.0e11
NU = 0.3
RHO = 7850.0
b = 0.01  # cross-section side; L / b = 100 → very slender → severe locking
A = b * b
I = b**4 / 12.0  # noqa: E741
L = 1.0
P_TOTAL = -100.0  # N (small load → linear regime)

NX, NY, NZ = 20, 1, 1  # 20 elements axial, 1 through each thickness

delta_eb = P_TOTAL * L**3 / (3.0 * E * I)
print(f"L / b = {L / b:.0f} (slender)")
print(f"δ_EB = P L^3 / (3 EI) = {delta_eb * 1e6:.3f} µm")


def run_with_integration(label: str, kwargs: dict) -> tuple[float, float]:
    """Solve the same cantilever with the given HEX8 spec; return
    (UY tip mid-surface, ratio_to_EB)."""
    xs = np.linspace(0.0, L, NX + 1)
    ys = np.linspace(0.0, b, NY + 1)
    zs = np.linspace(0.0, b, NZ + 1)
    grid = pv.StructuredGrid(*np.meshgrid(xs, ys, zs, indexing="ij")).cast_to_unstructured_grid()

    m = femorph_solver.Model.from_grid(grid)
    if kwargs:
        m.assign(ELEMENTS.HEX8(**kwargs), material={"EX": E, "PRXY": NU, "DENS": RHO})
    else:
        m.assign(ELEMENTS.HEX8, material={"EX": E, "PRXY": NU, "DENS": RHO})

    pts = np.asarray(m.grid.points)
    node_nums = np.asarray(m.grid.point_data["ansys_node_num"])
    fixed = node_nums[pts[:, 0] < 1e-9]
    m.fix(nodes=fixed.tolist(), dof="ALL")
    tip_nodes = node_nums[pts[:, 0] > L - 1e-9]
    fz_per = P_TOTAL / len(tip_nodes)
    for nn in tip_nodes:
        m.apply_force(int(nn), fy=fz_per)

    res = m.solve_static()
    g = femorph_solver.io.static_result_to_grid(m, res)
    pts_g = np.asarray(g.points)
    tip_mask = pts_g[:, 0] > L - 1e-9
    uy_tip = float(g.point_data["displacement"][tip_mask, 1].mean())
    return uy_tip, uy_tip / delta_eb


# %%
# Three integration variants — side by side
# -----------------------------------------

variants = [
    ("plain_gauss", {"integration": "plain_gauss"}, "#d62728"),
    ("default (B-bar)", {}, "#1f77b4"),
    ("enhanced_strain (EAS)", {"integration": "enhanced_strain"}, "#2ca02c"),
]
results: list[tuple[str, float, float, str]] = []

print()
print(f"{'integration':>26}  {'UY_tip [µm]':>12}  {'UY / δ_EB':>10}")
print(f"{'-' * 26:>26}  {'-' * 12:>12}  {'-' * 10:>10}")
for label, kwargs, colour in variants:
    uy, ratio = run_with_integration(label, kwargs)
    results.append((label, uy, ratio, colour))
    print(f"{label:>26}  {uy * 1e6:12.4f}  {ratio:10.6f}")

# %%
# Verify the textbook ordering
# ----------------------------
#
# The locked plain-Gauss solution must underpredict the
# Euler-Bernoulli reference on a slender beam (textbook
# expectation: locks to ~10–30 % of the analytical answer at
# L/b = 100, NX = 20 with one element through the thickness).
# B-bar is **not** a shear-locking cure — it cures volumetric
# locking only — so it should give the same locked deflection
# as plain Gauss in this benchmark.  Only EAS recovers the
# bending mode.

ratios = {label: ratio for label, _, ratio, _ in results}
assert ratios["plain_gauss"] < 0.5, "plain_gauss should lock to << 1.0"
assert ratios["default (B-bar)"] < 0.5, "B-bar alone should also lock here"
assert ratios["enhanced_strain (EAS)"] > 0.85, "EAS should recover ≥ 85 % of EB"

# %%
# Render the deflection-ratio comparison
# --------------------------------------

labels = [r[0] for r in results]
ratios_values = [r[2] for r in results]
colours = [r[3] for r in results]

fig, ax = plt.subplots(1, 1, figsize=(7.0, 4.0))
bars = ax.bar(labels, ratios_values, color=colours)
ax.axhline(1.0, linestyle="--", color="black", lw=1.5, label="Euler-Bernoulli (target)")
ax.set_ylabel(r"$\delta^{h}_{\mathrm{tip}} / \delta_{\mathrm{EB}}$")
ax.set_title(
    f"Cantilever HEX8 shear-locking — L / b = {L / b:.0f}, N_x = {NX}, 1 element through h",
    fontsize=10,
)
ax.set_ylim(0.0, 1.10)
for bar, ratio in zip(bars, ratios_values):
    ax.annotate(
        f"{ratio:.3f}",
        xy=(bar.get_x() + bar.get_width() / 2, ratio),
        xytext=(0, 4),
        textcoords="offset points",
        ha="center",
        fontsize=9,
    )
ax.legend(loc="upper left", fontsize=9)
ax.grid(True, axis="y", ls=":", alpha=0.5)
fig.tight_layout()
fig.show()

# %%
# Closing note
# ------------

print()
print("Take-aways:")
print(
    f"  • plain_gauss    → {ratios['plain_gauss']:.1%} of EB  (locked; spurious shear-strain energy)"
)
print(
    f"  • B-bar default  → {ratios['default (B-bar)']:.1%} of EB  (B-bar cures *volumetric* locking, not shear)"
)
print(
    f"  • EAS variant    → {ratios['enhanced_strain (EAS)']:.1%} of EB  (the right cure for thin-bending solid kernels)"
)