r"""
Propped cantilever under uniformly-distributed load
===================================================

Statically-indeterminate first-order beam: clamped at
:math:`x = 0`, simply supported at :math:`x = L`, with a
uniform downward distributed load :math:`q` over the full
span.  The closed-form Euler-Bernoulli solution (Roark &
Young 1989, Table 8 case 7; Timoshenko 1955 §40) gives

.. math::

    R_\mathrm{prop} = \tfrac{3}{8} q L,
    \qquad
    R_\mathrm{clamp} = \tfrac{5}{8} q L,
    \qquad
    M_\mathrm{clamp} = \tfrac{1}{8} q L^{2},

and the maximum (downward) deflection occurs at
:math:`x^{*} \approx 0.5785\, L` with magnitude

.. math::

    \delta_\mathrm{max}
        \;=\; \frac{q\, L^{4}}{185\, E I},
    \qquad
    \delta(L/2) = \frac{q L^{4}}{192\, E I}.

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

Standard 3D BEAM2 (Hermite-cubic Bernoulli kinematics, Cook
Table 16.3-1) line of :math:`N_e = 20` elements along the
span.  Both end-section restraints use the native :meth:`Model.fix`
API:

* node 1 (clamp): all 6 DOFs pinned.
* node :math:`N_e + 1` (prop): UY = 0; out-of-plane DOFs
  (UZ, ROTX, ROTY, ROTZ) suppressed to keep the response purely
  in-plane.

The UDL is applied as the consistent nodal-force vector for a
uniform load on a Hermite beam (Cook §16.5,
:math:`\{f_e\} = q L_e \{1/2,\, L_e/12,\, 1/2,\, -L_e/12\}^{T}`),
which is what an APDL ``F,FY,…`` over the line equivalents to.

References
----------

* Timoshenko, S. P. (1955) *Strength of Materials, Part I*,
  3rd ed., Van Nostrand, §40.
* Roark, R. J. and Young, W. C. (1989) *Roark's Formulas for
  Stress and Strain*, 6th ed., McGraw-Hill, Table 8 case 7.
* Cook, R. D., Malkus, D. S., Plesha, M. E., Witt, R. J. (2002)
  *Concepts and Applications of Finite Element Analysis*, 4th
  ed., Wiley, §16.5 (consistent loads on beam elements).
"""

from __future__ import annotations

import numpy as np
import pyvista as pv
from vtkmodules.util.vtkConstants import VTK_LINE

import femorph_solver
from femorph_solver import ELEMENTS

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

E = 2.0e11  # Pa (steel)
NU = 0.30
RHO = 7850.0  # kg/m^3
b = 0.050  # square cross-section side [m]
A = b * b
I = b**4 / 12.0  # second moment of area  # noqa: E741
J = 2.0 * I  # thin-square Saint-Venant J (no torsion in this BC set)
L = 1.0  # span [m]
q = 1.0e3  # uniform line load [N/m] (downward → -y)

N_ELEM = 20

# Closed-form quantities -----------------------------------------------

R_clamp_published = 5.0 / 8.0 * q * L
R_prop_published = 3.0 / 8.0 * q * L
M_clamp_published = q * L**2 / 8.0
delta_mid_published = q * L**4 / (192.0 * E * I)
delta_max_published = q * L**4 / (185.0 * E * I)
x_star = 0.5785  # location of max deflection / L

print(f"q = {q:.1f} N/m, L = {L:.2f} m, EI = {E * I:.3e} N m^2")
print(f"R_clamp  = 5/8 q L          = {R_clamp_published:.4f} N")
print(f"R_prop   = 3/8 q L          = {R_prop_published:.4f} N")
print(f"M_clamp  = 1/8 q L^2        = {M_clamp_published:.4f} N m")
print(f"δ(L/2)   = q L^4 / (192 EI) = {delta_mid_published:.4e} m")
print(f"δ_max    = q L^4 / (185 EI) = {delta_max_published:.4e} m  (at x ≈ {x_star} L)")

# %%
# Build a 20-element BEAM2 line
# -----------------------------

points = np.array(
    [[i * L / N_ELEM, 0.0, 0.0] for i in range(N_ELEM + 1)],
    dtype=np.float64,
)
cells_list: list[int] = []
for i in range(N_ELEM):
    cells_list.extend([2, i, i + 1])
cells = np.asarray(cells_list, dtype=np.int64)
cell_types = np.full(N_ELEM, VTK_LINE, dtype=np.uint8)
grid = pv.UnstructuredGrid(cells, cell_types, points)

m = femorph_solver.Model.from_grid(grid)
m.assign(
    ELEMENTS.BEAM2,
    material={"EX": E, "PRXY": NU, "DENS": RHO},
    real=(A, I, I, J),
)

# %%
# Boundary conditions
# -------------------
#
# Node 1 — clamp (all 6 DOFs pinned).
# Node N_e + 1 — pin only UY (the prop) and suppress
# out-of-plane DOFs so the response is in the x-y plane.

m.fix(nodes=[1], dof="ALL")
prop = N_ELEM + 1
m.fix(nodes=[prop], dof="UY")
# Suppress out-of-plane / torsional DOFs so the response stays
# in the x-y plane.  ROTZ is **free** at the prop — pinning it
# would turn the simple support into a clamp and the problem
# would be a clamped-clamped beam (50/50 reaction split) instead
# of a propped cantilever (5/8 - 3/8 split).
m.fix(nodes=[prop], dof="UZ")
m.fix(nodes=[prop], dof="ROTX")
m.fix(nodes=[prop], dof="ROTY")

# %%
# Consistent UDL on a Hermite beam
# --------------------------------
#
# For a single Bernoulli element of length :math:`L_e` under
# uniform line load :math:`q`, the consistent nodal forces /
# moments at its two endpoints are
#
# .. math::
#
#     \{f_e\} = q\,L_e
#     \begin{Bmatrix} 1/2 \\ L_e/12 \\ 1/2 \\ -L_e/12 \end{Bmatrix},
#
# i.e. half the segment load goes to each endpoint as a
# transverse force, and a fixing moment of magnitude
# :math:`q L_e^{2}/12` accumulates at each end with opposite
# signs (Cook §16.5).  We assemble these onto the model nodes.

L_e = L / N_ELEM
fy_per_node = np.zeros(N_ELEM + 1)
mz_per_node = np.zeros(N_ELEM + 1)
for e in range(N_ELEM):
    f_half = -q * L_e / 2.0  # half of -q L_e per endpoint (downward)
    m_left = -q * L_e**2 / 12.0  # = -q L_e^2 / 12 from {-qL/2,-qL^2/12,-qL/2,+qL^2/12}
    m_right = +q * L_e**2 / 12.0
    fy_per_node[e] += f_half
    fy_per_node[e + 1] += f_half
    mz_per_node[e] += m_left
    mz_per_node[e + 1] += m_right

for i in range(N_ELEM + 1):
    if abs(fy_per_node[i]) > 0:
        m.apply_force(i + 1, fy=float(fy_per_node[i]))
    if abs(mz_per_node[i]) > 0:
        m.apply_force(i + 1, mz=float(mz_per_node[i]))

# %%
# Static solve + reaction extraction
# ----------------------------------

res = m.solve_static()
dof = m.dof_map()


def _react(node: int, dof_idx: int) -> float:
    rows = np.where((dof[:, 0] == node) & (dof[:, 1] == dof_idx))[0]
    return float(res.reaction[rows[0]]) if len(rows) else 0.0


R_clamp_y = _react(1, 1)  # UY reaction at clamp
M_clamp_z = _react(1, 5)  # ROTZ reaction at clamp
R_prop_y = _react(prop, 1)  # UY reaction at prop

print()
print("reactions  (femorph-solver) → (analytical)")
print(f"  R_clamp_UY  = {R_clamp_y:+.4f}  →  {+R_clamp_published:+.4f} N")
print(f"  R_prop_UY   = {R_prop_y:+.4f}  →  {+R_prop_published:+.4f} N")
print(f"  M_clamp_RZ  = {M_clamp_z:+.4f}  →  {+M_clamp_published:+.4f} N m")

# %%
# Equilibrium check + tolerance assertions
# ----------------------------------------
#
# Sum of vertical reactions must balance the applied UDL
# (5/8 + 3/8 = 1) and the moment about the clamp from the
# downward load + prop reaction must be zero.

assert np.isclose(R_clamp_y + R_prop_y, q * L, rtol=1e-6), "vertical equilibrium failed"
assert np.isclose(R_clamp_y, R_clamp_published, rtol=1e-6), "clamp UY reaction off"
assert np.isclose(R_prop_y, R_prop_published, rtol=1e-6), "prop UY reaction off"
assert np.isclose(M_clamp_z, M_clamp_published, rtol=1e-6), "clamp moment off"

# %%
# Mid-span deflection
# -------------------
#
# Hermite cubic bases recover Euler-Bernoulli kinematics
# exactly for prismatic beams under nodal loads, so the
# mid-span UY equals the analytical :math:`\delta(L/2) = q L^4
# / (192 E I)` to machine precision on a uniform mesh.

mid_node = N_ELEM // 2 + 1
mid_uy_val = float(res.displacement[np.where((dof[:, 0] == mid_node) & (dof[:, 1] == 1))[0][0]])
err_mid = (mid_uy_val - (-delta_mid_published)) / (-delta_mid_published)
print()
print(f"δ(L/2) computed   = {mid_uy_val:+.4e} m")
print(f"δ(L/2) published  = {-delta_mid_published:+.4e} m")
print(f"relative error    = {err_mid * 100:+.4f} %")
assert abs(err_mid) < 1.0e-4, "mid-span deflection error too large for Bernoulli BEAM2"

# %%
# Render the deflected line
# -------------------------

pts = np.asarray(grid.points)
disp_y = np.zeros(N_ELEM + 1)
for i in range(N_ELEM + 1):
    rows = np.where((dof[:, 0] == i + 1) & (dof[:, 1] == 1))[0]
    if len(rows):
        disp_y[i] = float(res.displacement[rows[0]])
warped = grid.copy()
warped.points = pts + np.column_stack([np.zeros(N_ELEM + 1), 200.0 * disp_y, np.zeros(N_ELEM + 1)])
warped["uy"] = disp_y

plotter = pv.Plotter(off_screen=True, window_size=(720, 280))
plotter.add_mesh(grid, color="grey", line_width=2, opacity=0.5)
plotter.add_mesh(warped, scalars="uy", line_width=5, cmap="viridis")
plotter.add_points(
    np.array([[0.0, 0.0, 0.0], [L, 0.0, 0.0]]),
    render_points_as_spheres=True,
    point_size=18,
    color="black",
)
plotter.view_xy()
plotter.camera.zoom(1.05)
plotter.show()