import matplotlib.pyplot as plt
import meshio
import numpy as np
import torch
from scipy.interpolate import RegularGridInterpolator

from torchfem import Shell
from torchfem.data import get_data
from torchfem.materials import IsotropicElasticityPlaneStress
from torchfem.mesh import rect_quad, rect_tri

torch.set_default_dtype(torch.float64)

q = 1.0e-4  # Distributed load per unit area
E = 1.7472e7  # Young's modulus
nu = 0.3  # Poisson's ratio
t = 0.0001  # Thickness of the shell

# Extract reference solution obtained from Abaqus with 48x48 S8R elements
reference = meshio.read(get_data("clamped_plate_uniform_S8R.vtk"))
M = 49
x = reference.points[: M**2, 0]
y = reference.points[: M**2, 1]
w = reference.point_data["U"][: M**2, 2]

# Interpolate reference solution to nodes
interp = RegularGridInterpolator((np.unique(x), np.unique(y)), w.reshape(M, M))

N = [4, 8, 16, 32]
variants = ("up", "down", "center", "quad")
mesh_sizes = {variant: [] for variant in variants}
errors = {variant: [] for variant in variants}

for variant in variants:
    for n in N:
        # Material parameters
        material = IsotropicElasticityPlaneStress(E=E, nu=nu)

        # Define nodes and element of the plate
        if variant == "quad":
            nodes, elements = rect_quad(n + 1, n + 1)
        else:
            nodes, elements = rect_tri(n + 1, n + 1, variant=variant)
        nodes = torch.hstack([nodes, torch.zeros((len(nodes), 1))])

        # Define Shell model
        plate = Shell(nodes, elements, material, thickness=t)

        # Apply a uniform pressure in the negative z-direction
        surface = torch.ones(plate.n_nod, dtype=torch.bool)
        plate.forces[:, 0:3] = plate.integrate_surface_load(surface, -q)

        # Apply displacement boundary conditions
        plate.constraints[nodes[:, 0] < 0.01, :] = True
        plate.constraints[nodes[:, 1] < 0.01, :] = True
        plate.constraints[nodes[:, 0] > 0.99, 0] = True
        plate.constraints[nodes[:, 0] > 0.99, 4] = True
        plate.constraints[nodes[:, 1] > 0.99, 1] = True
        plate.constraints[nodes[:, 1] > 0.99, 3] = True

        # Solve
        u, f, _, _, _ = plate.solve(method="direct")

        # Compute error with respect to reference solutio
        w_i = interp(nodes[:, :2])
        error = np.linalg.norm(w_i - u[:, 2].numpy()) / np.linalg.norm(w_i)

        # Store error and mesh size
        errors[variant].append(error)
        mesh_sizes[variant].append(1 / n)

fig, ax = plt.subplots(1, len(variants), figsize=(12, 3))
for i, variant in enumerate(variants):
    logh = np.log10(mesh_sizes[variant])
    loge = np.log10(errors[variant])
    ax[i].plot(logh, loge, marker="o")
    ax[i].plot(logh, 2 * logh, "k")
    ax[i].grid(True)
    ax[i].set_xlabel("log(h)")
    ax[i].set_ylabel("log(E)")
    ax[i].set_title(f"Variant: {variant}")

plt.tight_layout()
plt.show()

m = 1.0  # Distributed moment per length
E = 1.7472e7  # Young's modulus
nu = 0.0  # Poisson's ratio
t = 0.0001  # Thickness of the shell

# Material parameters
material = IsotropicElasticityPlaneStress(E=E, nu=nu)

# Define nodes and element of the plate
nodes = torch.tensor(
    [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
)
connectivities = {
    "Tria1": torch.tensor([[0, 1, 2], [1, 2, 3]]),
    "Quad1": torch.tensor([[0, 1, 3, 2]]),
}

for etype, elements in connectivities.items():
    # Define Shell model
    plate = Shell(nodes, elements, material, thickness=t)

    # Boundary conditions
    tip = nodes[:, 1] > 0.99
    plate.constraints[~tip, :] = True
    plate.forces[tip, 3] = m

    u, f, _, _, _ = plate.solve(method="direct")

    # Analytical solution with M = 2m, D = E t^3 / 12: w = M L^2 / (2D), rot_x = M L / D
    # Both element types reproduce this state of constant curvature exactly.
    assert torch.allclose(u[tip, 2], 686813.19 * torch.ones_like(u[tip, 2]))
    assert torch.allclose(u[tip, 3], 1373626.37 * torch.ones_like(u[tip, 3]))
    assert u[tip, 4].abs().max() < 1e-6 * u[tip, 3].abs().max()

# Plot the quadrilateral solution
plate.plot(node_property={"U": u[:, :3]})

m = 1.0  # Distributed moment per length
E = 1.7472e7  # Young's modulus
nu = 0.0  # Poisson's ratio

variants = ("up", "down", "center")

# Analytical strain energy is 12 m / (E t^2). The "up" and "down" patterns reproduce
# it, while "center" is 23% too high.
expected = {
    ("up", 0.001): 0.686814,
    ("up", 0.0001): 68.681320,
    ("down", 0.001): 0.686813,
    ("down", 0.0001): 68.681319,
    ("center", 0.001): 0.845311,
    ("center", 0.0001): 84.530856,
}

# Material parameters
material = IsotropicElasticityPlaneStress(E=E, nu=nu)

for variant in variants:
    for t in [0.001, 0.0001]:
        # Define nodes and element of the plate
        nodes, elements = rect_tri(2, 2, variant=variant)
        nodes = torch.hstack([nodes, torch.zeros((len(nodes), 1))])

        # Define Shell model
        plate = Shell(nodes, elements, material, thickness=t)

        # Boundary conditions
        left = nodes[:, 1] < 0.01
        right = nodes[:, 1] > 0.99
        back = nodes[:, 0] < 0.01
        front = nodes[:, 0] > 0.99
        plate.constraints[left, :] = True
        plate.constraints[back, :] = True
        plate.forces[right, 3] = m
        plate.forces[front, 4] = -m

        u, f, _, _, _ = plate.solve(method="direct")

        strain_energy = 0.5 * t * torch.einsum("...i, ...i->", u, f)
        assert torch.allclose(strain_energy, torch.tensor(expected[(variant, t)]))

R = 25.0  # Radius of the cylinder
L = 50.0  # Length of the cylinder
opening = 40.0  # Half opening angle in degrees
g = 90.0  # Gravity load per unit area
E = 4.32e8  # Young's modulus
nu = 0.0  # Poisson's ratio
t = 0.25  # Thickness of the shell
w_ref = 0.3024  # Reference deflection at the midpoint of the free edge

N = [4, 8, 16, 32]
variants = ("up", "down", "center", "quad")
deflections = {variant: [] for variant in variants}
material = IsotropicElasticityPlaneStress(E=E, nu=nu)


def build_roof(n, variant):
    """Quarter roof, meshed by wrapping a unit square onto the cylinder."""
    if variant == "quad":
        nodes, elements = rect_quad(n + 1, n + 1)
    else:
        nodes, elements = rect_tri(n + 1, n + 1, variant=variant)
    angle = torch.deg2rad(opening * nodes[:, 0])
    nodes = torch.stack(
        [R * torch.sin(angle), R * torch.cos(angle), L / 2 * nodes[:, 1]], dim=1
    )
    roof = Shell(nodes, elements, material, thickness=t)

    # Apply gravity in the negative y-direction
    surface = torch.ones(roof.n_nod, dtype=torch.bool)
    load_vec = torch.tensor([0.0, -g, 0.0])
    roof.forces[:, 0:3] = roof.integrate_surface_load(surface, load_vec)

    # Symmetry at midspan, rigid diaphragm at the end, symmetry at the crown
    x, z = nodes[:, 0], nodes[:, 2]
    roof.constraints[z < 1e-6, 2:5] = True
    roof.constraints[z > L / 2 - 1e-6, 0:2] = True
    roof.constraints[x < 1e-6, 0] = True
    roof.constraints[x < 1e-6, 4:6] = True
    return roof


for variant in variants:
    for n in N:
        roof = build_roof(n, variant)
        u, *_ = roof.solve(method="direct")

        # Deflection at the midpoint of the free edge
        x, z = roof.nodes[:, 0], roof.nodes[:, 2]
        edge = (z < 1e-6) & (x > x.max() - 1e-6)
        deflections[variant].append(-u[edge, 1].item() / w_ref)

# Every variant should approach the reference from below.
for variant, w in deflections.items():
    assert w == sorted(w) and 0.98 < w[-1] < 1.0

fig, ax = plt.subplots(figsize=(5, 3.5))
for variant in variants:
    ax.plot(N, deflections[variant], marker="o", label=variant)
ax.axhline(1.0, color="k", linewidth=0.8)
ax.set_xscale("log", base=2)
ax.set_xticks(N, [str(n) for n in N])
ax.grid(True)
ax.set_xlabel("Elements per side")
ax.set_ylabel("Normalized deflection")
ax.legend()

plt.tight_layout()
plt.show()

roof = build_roof(32, "quad")
u, *_ = roof.solve()
roof.plot(20.0 * u[:, 0:3], node_property={"u": u[:, 1]}, mirror=(True, False, True))

Standard tests for shells and their convergence behavior¶

Fully clamped plate problem¶

Cantilever bending plate problem¶

Two-sided clamped plate¶

Scordelis-Lo roof¶