"""Linear elasticity, CG-FEM: the Timoshenko-Goodier cantilever, solved exactly.

    -div sigma(u) = 0    in (0, L) x (-D/2, D/2)
    sigma(u) = 2 mu eps(u) + lambda tr(eps(u)) I,   eps(u) = (grad u + grad u^T)/2

A beam clamped on its west end and loaded on its east end by a parabolic
shear of resultant ``P``, in plane stress. The classical reference of the
adaptive-FEM literature (Timoshenko & Goodier, *Theory of Elasticity*, §21;
Augarde & Deeks 2008), and what makes it the right first elasticity case
here: the 2D solution is known in closed form, and it is a POLYNOMIAL --
cubic in ``x``, cubic in ``y`` -- so Q1 and Q2 must converge at their
nominal rates (2 and 3 in L2), and a Q3 space would contain it exactly.

The stress is never written by hand. ``u_exact`` is the only formula, and
``sigma_exact`` is obtained from it by autodiff and the same constitutive
law the weak form uses (as the Stokes example differentiates its source).
The script first checks ``div sigma_exact = 0`` and that the top and bottom
faces are traction-free, so a wrong constant in ``u_exact`` would be caught
before any solve.

Boundary conditions are keyed by side name, as everywhere: the clamped end
takes the exact displacement (the standard choice for this case, a plain
``u = 0`` is not compatible with the exact stress), the loaded end takes the
exact traction ``sigma . n``, the two long faces are free.

The displacement space is a Lagrange basis with one scalar component per
direction -- the velocity space of Stokes -- since ``u`` lives in ``H^1``:
there is nothing to conform to in ``H(div)`` or ``H(curl)`` here.
"""

# %%
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

from scimba_jax.linear_approximation.basis.analytic_bases import (
    local_lagrange_basis,
    local_lagrange_basis_by_logical,
)
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.error_analysis import l2_error
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.classical_weakform.linear_elasticity_weak_form import (  # noqa: E501
    LinearElasticityWeakForm,
    lame_parameters,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet, Neumann
from scimba_jax.plots.plots_galerkin import sample_solution

DIM = 2

# Geometry and material. A slender beam (L/D = 8) so that the bending is
# visible, E large enough that the deflection stays small compared to L --
# this is LINEAR elasticity, the deformed plot is scaled anyway.
LENGTH, DEPTH = 8.0, 1.0
YOUNG, POISSON = 1.0e3, 0.3
LOAD = 1.0  # resultant of the end shear, upwards
INERTIA = DEPTH**3 / 12.0
MU, LAM = lame_parameters(YOUNG, POISSON, plane_stress=True)


def u_exact(x):
    """Timoshenko-Goodier cantilever, plane stress, end shear ``P`` upwards.

    Cubic in both variables: Q3 reproduces it exactly.
    """
    X, Y = x[0], x[1]
    scale = LOAD / (6.0 * YOUNG * INERTIA)
    u_x = (
        -scale
        * Y
        * ((6.0 * LENGTH - 3.0 * X) * X + (2.0 + POISSON) * (Y * Y - DEPTH**2 / 4.0))
    )
    u_y = scale * (
        3.0 * POISSON * Y * Y * (LENGTH - X)
        + (4.0 + 5.0 * POISSON) * DEPTH**2 * X / 4.0
        + (3.0 * LENGTH - X) * X * X
    )
    return jnp.array([u_x, u_y])


def sigma_exact(x):
    """``sigma(u_exact)``, differentiated rather than written out."""
    jac = jax.jacfwd(u_exact)(x)
    eps = 0.5 * (jac + jac.T)
    return 2.0 * MU * eps + LAM * jnp.trace(eps) * jnp.eye(DIM)


def traction(normal):
    """The exact traction ``sigma . n`` on a side of constant outward normal."""
    n = jnp.asarray(normal)
    return lambda x: sigma_exact(x) @ n


class Cantilever(AbstractPhysicalWeakModel):
    """Clamped west (exact displacement), loaded east, free north/south."""

    def __init__(self):
        super().__init__(dim=DIM)
        self.weak_forms = {
            "interior": LinearElasticityWeakForm(dim=DIM, mu=MU, lam=LAM),
        }
        self.boundary_conditions = {
            "west": Dirichlet(u_exact),
            "east": Neumann(traction((1.0, 0.0))),
            # Traction-free: the exact traction is zero there, and saying so
            # keeps the two long faces natural rather than constrained.
            "south": Neumann(lambda x: jnp.zeros(DIM)),
            "north": Neumann(lambda x: jnp.zeros(DIM)),
        }


# %% Self-check of the reference: equilibrium and free faces, by autodiff.
def _div_sigma(x):
    return jnp.trace(jax.jacfwd(sigma_exact)(x), axis1=1, axis2=2)


probe = jnp.array([[0.3, 0.2], [5.0, -0.4], [7.9, 0.49]])
assert float(jnp.abs(jax.vmap(_div_sigma)(probe)).max()) < 1e-10
faces = jnp.array([[1.0, DEPTH / 2], [6.0, DEPTH / 2], [1.0, -DEPTH / 2]])
assert float(jnp.abs(jax.vmap(traction((0.0, 1.0)))(faces)).max()) < 1e-10
tip = float(u_exact(jnp.array([LENGTH, 0.0]))[1])
print(
    f"reference checked: div sigma = 0, free faces.  Tip deflection {tip:.4f}"
    f" (Euler-Bernoulli P L^3 / 3EI = {LOAD * LENGTH**3 / (3 * YOUNG * INERTIA):.4f},"
    " the rest is shear)"
)


# %% Mesh, space, scheme.
def make_scheme(order, n_y):
    """Q``order`` on an ``8 n_y x n_y`` grid of the beam (cells stay square)."""
    n_x = int(round(n_y * LENGTH / DEPTH))
    # The unit square scaled to the beam: an affine map, so the Cartesian
    # mesh keeps its closed-form cell inverse and the boundary sides keep
    # their names.
    to_beam = InvertibleFunction(
        lambda x: jnp.array([LENGTH * x[0], DEPTH * (x[1] - 0.5)]),
        lambda y: jnp.array([y[0] / LENGTH, y[1] / DEPTH + 0.5]),
    )
    mesh = Mesh(
        dim=DIM,
        n_cells=(n_x, n_y),
        # Exact for products of two Q`order` gradients.
        ref_quad=UnitSquareTensorized(dim=DIM, order=2 * order + 1),
        mapping=Mapping(mappings=[to_beam]),
        is_identity_mapping=False,
    )
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** DIM,
        out_dim=DIM,
        mesh=mesh,
        local_basis=lambda y, i, m: local_lagrange_basis(
            y, i, m, order=order, out_dim=DIM
        ),
        local_basis_by_logical=lambda y, i, m: local_lagrange_basis_by_logical(
            y, i, m, order=order, out_dim=DIM
        ),
        basis_type="field",
    )
    return EllipticFEscheme(Cantilever(), VariablesFE(basis=basis, nb_variables=DIM))


def solve(order, n_y, matrix_free=False):
    scheme = make_scheme(order, n_y)
    # Linear, and symmetric positive definite (Korn): one Newton step, and
    # the matrix-free path runs plain CG on it.
    options = {"max_iter_linear": 20000} if matrix_free else {}
    scheme = EllipticFEscheme.solve(
        scheme, matrix_free=matrix_free, max_iter=1, **options
    )
    return scheme, float(l2_error(scheme, u_exact, relative=True))


# %% Convergence: Q1 and Q2 at their nominal rates.
print(f"\n{'element':>8}{'cells':>10}{'rel. L2':>12}{'rate':>7}")
print("-" * 37)
solved = {}
for order in (1, 2):
    previous = None
    for n_y in (1, 2, 4):
        scheme, err = solve(order, n_y)
        solved[(order, n_y)] = scheme
        rate = "" if previous is None else f"{np.log2(previous / err):7.2f}"
        previous = err
        print(f"{'Q' + str(order):>8}{f'{8 * n_y}x{n_y}':>10}{err:>12.3e}{rate:>7}")
    print()
print("Q1 and Q2 converge at their nominal L2 rates, 2 and 3.")

# %% Both solver paths see the same boundary terms (assembled J vs matrix-free
# CG are written separately): they must agree to the CG tolerance, and CG
# only applies because the Dirichlet lift keeps the operator symmetric.
_, err_blocks = solve(2, 2)
_, err_mf = solve(2, 2, matrix_free=True)
print(f"\nQ2, 16x2:  assembled {err_blocks:.6e}   matrix-free CG {err_mf:.6e}")


# %% The picture: deformed beam coloured by the von Mises stress.
def sample_fields(scheme, n_side=6):
    """``u_h`` and ``sigma(u_h)`` on the reference grid of every cell.

    Same layout as ``sample_solution`` (which is used for the points and the
    triangles, and returns one component only). The stress is the weak
    form's own ``stress`` applied to ``local_variables`` -- the operator
    acts on the expansion as a ``ParamFunction``, never on sampled values,
    which is what makes a derivative of ``u_h`` expressible at all.
    """
    mesh = scheme.variables.mesh
    weak_form = scheme.pde.weak_forms["interior"]
    grid = jnp.stack(
        jnp.meshgrid(
            jnp.linspace(0.0, 1.0, n_side),
            jnp.linspace(0.0, 1.0, n_side),
            indexing="ij",
        ),
        axis=-1,
    ).reshape(-1, 2)
    dofs = scheme.variables.dofsl

    def one_cell(cell):
        # Logical cell, then the beam: `local_variables` takes physical points.
        points = mesh.mapping.local_mapping(mesh._unit_hypercube_to_cell(cell, grid))
        u_h = scheme.local_variables(cell, dofs)
        sigma_h = weak_form.stress(u_h)
        return (
            jax.vmap(lambda x: u_h(scheme, x))(points),
            jax.vmap(lambda x: sigma_h(scheme, x))(points),
        )

    u_h, sigma_h = jax.jit(jax.vmap(one_cell))(jnp.arange(mesh.n_cells_total))
    return np.asarray(u_h).reshape(-1, DIM), np.asarray(sigma_h).reshape(-1, DIM, DIM)


def von_mises(sigma):
    """Plane-stress von Mises equivalent, ``sqrt(s_xx^2 - s_xx s_yy + s_yy^2 + 3 s_xy^2)``."""
    sxx, syy, sxy = sigma[:, 0, 0], sigma[:, 1, 1], sigma[:, 0, 1]
    return np.sqrt(sxx**2 - sxx * syy + syy**2 + 3.0 * sxy**2)


scheme = solved[(2, 4)]
points, triangles, _ = sample_solution(scheme, n_side=6)
displacement, sigma_h = sample_fields(scheme, n_side=6)
mises = von_mises(sigma_h)
mises_exact = von_mises(np.asarray(jax.vmap(sigma_exact)(jnp.asarray(points))))

AMPLIFY = 0.5  # deflection ~2 on a beam of length 8: half of it is readable
deformed = points + AMPLIFY * displacement

# What the error panel shows, measured on the 32x4 Q2 run: the stress is the
# gradient of a C0 field, so its error sits on the mesh lines (|d tau_xy| 0.08
# on the edges, 0.03 mid-cell) -- superconvergent inside, not across. It is
# the SAME on the five horizontal lines; the neutral axis only stands out
# because von Mises there is pure shear (sqrt(3)|tau| ~ 2.6, against ~12 one
# line up): the same absolute error, 5% instead of 1%, and a sensitivity
# d(VM)/d(tau) of sqrt(3) where sigma_xx no longer dominates.
fig, axes = plt.subplots(2, 1, figsize=(11, 6), constrained_layout=True)
error = np.abs(mises - mises_exact)
for ax, field, title in (
    (
        axes[0],
        mises,
        f"Q2, 32x4 cells: von Mises stress on the deformed beam (x{AMPLIFY})",
    ),
    (axes[1], error, f"|von Mises error|, max {error.max():.2e} (own colour range)"),
):
    tpc = ax.tripcolor(
        deformed[:, 0], deformed[:, 1], triangles, field, shading="gouraud"
    )
    ax.plot(
        [0, LENGTH, LENGTH, 0, 0],
        [-DEPTH / 2, -DEPTH / 2, DEPTH / 2, DEPTH / 2, -DEPTH / 2],
        "k--",
        lw=0.8,
        label="undeformed",
    )
    ax.set_aspect("equal")
    ax.set_title(title)
    fig.colorbar(tpc, ax=ax, shrink=0.8)
axes[0].legend(loc="upper left")
plt.show()