r"""
Reaction forces — global equilibrium and root moment from a tip load
====================================================================

Every linear static solve produces two displacement fields: the
``(N,)`` solved-for ``displacement`` at the free DOFs, and the
``(N,)`` ``reaction`` field nonzero only at the constrained DOFs —
the force the supports must carry to hold the prescribed boundary
displacements.  Sum the reactions and you reconstruct the support
loads engineers actually need: bolt loads, weld loads, foundation
reactions.

Two textbook checks every static solve must pass (Hibbeler 2017 §9.5,
Cook et al. 2002 §11.6):

.. math::

    \mathbf{F}_\text{ext, free} + \mathbf{F}_\text{reaction} = \mathbf{0},
    \qquad
    \sum (\mathbf{r} \times \mathbf{F})_\text{free}
      + \sum (\mathbf{r} \times \mathbf{F})_\text{reaction}
      = \mathbf{0}.

A tip-loaded cantilever is the cleanest demo because every term has a
simple closed form:

* Total clamped-face reaction in :math:`y` equals :math:`-P`
  (shear balance).
* Net :math:`x`-reaction is zero, but the *individual* node-by-node
  :math:`R_x` values are nonzero — they form an internal couple whose
  arm × force product gives the root bending moment :math:`M_z = -PL`.

The second point is what surprises new users: the clamp face is in
*compression on the bottom, tension on the top*, and that couple is
exactly what carries the bending moment back into the support.

Implementation
--------------

A slender HEX8-EAS cantilever (L = 2 m, 0.05 × 0.05 m, 40 × 3 × 3
mesh) clamped at :math:`x = 0` and loaded with :math:`P = -1000` N
in :math:`y` distributed uniformly across the tip face nodes.
:attr:`StaticResult.reaction` is the DOF-indexed reaction-force
vector — index it via the model's DOF map to read the components per
node.

References
----------

* Hibbeler, R. C. (2017) *Mechanics of Materials*, 10th ed., Pearson,
  §9.5 — internal stress resultants from cross-section equilibrium.
* Cook, R. D., Malkus, D. S., Plesha, M. E., Witt, R. J. (2002)
  *Concepts and Applications of Finite Element Analysis*, 4th ed.,
  Wiley, §11.6 — reaction-force recovery in displacement-based FEM.
* Bathe, K.-J. (2014) *Finite Element Procedures*, 2nd ed.,
  Prentice Hall, §4.2.2 — partitioned static system and the
  recovered reactions.
"""

from __future__ import annotations

import numpy as np
import pyvista as pv

import femorph_solver
from femorph_solver import ELEMENTS

# %%
# Build the cantilever
# --------------------

E = 2.0e11
NU = 0.30
RHO = 7850.0
L = 2.0
WIDTH = 0.05
HEIGHT = 0.05
P_TOTAL = -1000.0  # N — total tip load in -y

NX, NY, NZ = 40, 3, 3
xs = np.linspace(0.0, L, NX + 1)
ys = np.linspace(0.0, WIDTH, NY + 1)
zs = np.linspace(0.0, HEIGHT, NZ + 1)
grid = pv.StructuredGrid(*np.meshgrid(xs, ys, zs, indexing="ij")).cast_to_unstructured_grid()

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

pts = np.asarray(m.grid.points)
node_nums = np.asarray(m.grid.point_data["ansys_node_num"], dtype=np.int64)
clamped_pt_idx = np.where(pts[:, 0] < 1e-9)[0]
tip_pt_idx = np.where(pts[:, 0] > L - 1e-9)[0]

m.fix(nodes=node_nums[clamped_pt_idx].tolist(), dof="ALL")

# Distribute the total tip load uniformly across the tip-face nodes.
load_per_node = P_TOTAL / float(tip_pt_idx.size)
for k in tip_pt_idx:
    m.apply_force(int(node_nums[k]), fy=load_per_node)

print("Tip-loaded cantilever — reaction-force recovery")
print(f"  L = {L} m, cross = {WIDTH} × {HEIGHT} m, {NX}×{NY}×{NZ} HEX8-EAS")
print(f"  P_tip = {P_TOTAL:.1f} N in -y, distributed across {tip_pt_idx.size} face nodes")

# %%
# Static solve
# ------------

res = m.solve_static()
print(f"\n  {res!r}")

# %%
# Sum reactions at the clamped face
# ---------------------------------
#
# ``StaticResult.reaction`` is ``(n_dof,)`` and nonzero only at the
# constrained DOFs.  ``Model.dof_map()`` gives ``(node_number, dof_idx)``
# per row so we can fold the flat vector into per-node force vectors.

dof_map = m.dof_map()
reaction = np.asarray(res.reaction, dtype=np.float64)
clamped_node_set = set(int(n) for n in node_nums[clamped_pt_idx])

n_points = m.grid.n_points
node_to_pt_idx = {int(node_nums[i]): i for i in range(n_points)}

# Gather (R_x, R_y, R_z) per clamped node.
clamp_nodes = sorted(clamped_node_set)
clamp_reactions = np.zeros((len(clamp_nodes), 3), dtype=np.float64)
for row, (node, dof_idx) in enumerate(dof_map.tolist()):
    if int(node) in clamped_node_set and dof_idx < 3:
        i = clamp_nodes.index(int(node))
        clamp_reactions[i, int(dof_idx)] += reaction[row]

r_total = clamp_reactions.sum(axis=0)
print()
print(f"  ΣR_x  = {r_total[0]:+.3f} N   (expected  0)")
print(f"  ΣR_y  = {r_total[1]:+.3f} N   (expected {-P_TOTAL:+.3f})")
print(f"  ΣR_z  = {r_total[2]:+.3f} N   (expected  0)")

residual = np.abs(r_total + np.array([0.0, P_TOTAL, 0.0]))
print(f"\n  global force residual  ‖ΣR + ΣF_applied‖ = {np.linalg.norm(residual):.3e} N")
assert np.linalg.norm(residual) < 1e-6 * abs(P_TOTAL), "global force balance violated"

# %%
# Bending-moment recovery from the root couple
# --------------------------------------------
#
# Even though ΣR_x = 0, the per-node R_x values are *not* zero — they
# form a couple about the neutral axis:
#
#     M_z(root) = Σ over clamp nodes of  y · R_x  +  (something)
#
# For a transverse load in y on a tip-loaded cantilever the dominant
# couple is the ``y * R_x`` integral.  Equivalently, the root moment
# is :math:`-P \cdot L` (textbook), so we expect
# :math:`M_z \approx 2000\,\text{N·m}` for our configuration.

clamp_pts = pts[[node_to_pt_idx[n] for n in clamp_nodes]]
y = clamp_pts[:, 1] - 0.5 * WIDTH  # offset to neutral axis

# Σ (r × R)_z = Σ (x · R_y − y · R_x).  At the clamp, x = 0, so
# M_z(at clamp) = − Σ y · R_x.  Add the contribution from R_y across
# x = 0 (which is zero by construction).
m_z_clamp = -float(np.sum(y * clamp_reactions[:, 0]))
m_z_pub = -P_TOTAL * L
print()
print(f"  M_z(at clamp) recovered from ΣyR_x  = {m_z_clamp:+.3f} N·m")
print(f"  M_z(at clamp) closed form  −P·L     = {m_z_pub:+.3f} N·m")
print(
    f"  relative error                       = {abs(m_z_clamp - m_z_pub) / abs(m_z_pub) * 100:.3f}%"
)

# %%
# Per-node table for the four extreme corners
# -------------------------------------------

print()
print(f"  {'node':>6}  {'y':>8}  {'z':>8}  {'R_x [N]':>12}  {'R_y [N]':>12}  {'R_z [N]':>12}")
print("  " + "-" * 64)
# Sort by descending |R_x| to highlight the bending couple.
order = np.argsort(-np.abs(clamp_reactions[:, 0]))
for idx in order[:8]:
    node = clamp_nodes[idx]
    p = clamp_pts[idx]
    r = clamp_reactions[idx]
    print(
        f"  {node:>6}  {p[1]:>8.4f}  {p[2]:>8.4f}  {r[0]:>+12.3f}  {r[1]:>+12.3f}  {r[2]:>+12.3f}"
    )

# %%
# Render the reaction-force vectors at the clamped face
# -----------------------------------------------------

reaction_grid = pv.PolyData(clamp_pts)
reaction_grid["reaction"] = clamp_reactions
reaction_grid["|R|"] = np.linalg.norm(clamp_reactions, axis=1)

# Scale glyphs so the longest reaction arrow reaches 25 % of beam length.
max_reaction = float(reaction_grid["|R|"].max())
arrow_scale = 0.25 * L / max_reaction if max_reaction > 0 else 1.0

reaction_grid.set_active_vectors("reaction")
arrows = reaction_grid.glyph(orient="reaction", scale="|R|", factor=arrow_scale)

plotter = pv.Plotter(off_screen=True, window_size=(900, 540))
plotter.add_mesh(m.grid, style="wireframe", color="grey", opacity=0.4, line_width=1)
plotter.add_mesh(
    arrows,
    scalars="|R|",
    cmap="plasma",
    show_scalar_bar=True,
    scalar_bar_args={"title": "|R|  [N]"},
)
plotter.add_text(
    f"clamp-face reactions  —  ΣR_y = {r_total[1]:.1f} N", position="upper_left", font_size=11
)
plotter.view_xy()
plotter.camera.zoom(1.1)
plotter.show()

# %%
# Take-aways
# ----------

print()
print("Take-aways:")
print(
    "  • StaticResult.reaction is a DOF-indexed (N,) array — nonzero only at "
    "constrained DOFs.  Index it via Model.dof_map() to get per-node force "
    "vectors at the supports."
)
print(
    "  • Global force balance: ΣF_applied + ΣF_reaction = 0 to machine precision.  "
    "Use this as a sanity check on every solve before trusting downstream "
    "stress / strain recovery."
)
print(
    "  • Internal moments are carried by NODE-TO-NODE couples at the "
    "support: ΣR_x = 0 net, but R_x varies linearly with the lever-arm "
    "coordinate.  Σ y·R_x reconstructs the root bending moment to within "
    "FEM discretisation error."
)
print(
    "  • Reaction-force visualisation at the clamp is the engineer's first "
    "check that loads are flowing through the structure as intended — "
    "asymmetries here usually point at a misapplied BC or a missing fix."
)