#!/usr/bin/env python
# coding: utf-8

# In[12]:


import torch

from torchfem import Planar
from torchfem.materials import IsotropicPlaneStress
from torchfem.io import import_mesh
from torchfem.examples import get_example_file


# ## A simple cantilever beam in 2D 

# In[13]:


# Dimensions
Nx = 20
Ny = 10
Lx = 2.0
Ly = 1.0

# Material model (plane stress)
material = IsotropicPlaneStress(E=1000.0, nu=0.3)

# Create nodes
n1 = torch.linspace(0.0, Lx, Nx + 1)
n2 = torch.linspace(0.0, Ly, Ny + 1)
n1, n2 = torch.stack(torch.meshgrid(n1, n2, indexing="xy"))
nodes = torch.stack([n1.ravel(), n2.ravel()], dim=1)

# Create elements connecting nodes
elements = []
for j in range(Ny):
    for i in range(Nx):
        # Quad elements
        n0 = i + j * (Nx + 1)
        elements.append([n0, n0 + 1, n0 + Nx + 2, n0 + Nx + 1])
elements = torch.tensor(elements)

# Load at tip
forces = torch.zeros_like(nodes)
forces[(int((Ny + 1) / 2) + 1) * (Nx + 1) - 1, 1] = -1.0

# Constrained displacement at left end
displacements = torch.zeros_like(nodes)
constraints = torch.zeros_like(nodes, dtype=bool)
for i in range(Ny + 1):
    constraints[i * (Nx + 1), :] = True

# Thickness
thickness = 0.1 * torch.ones(len(elements))

cantilever = Planar(
    nodes, elements, forces, displacements, constraints, thickness, material.C()
)
cantilever.plot()


# In[14]:


# Solve
u, f = cantilever.solve()
# Plot
cantilever.plot(u, node_property=torch.norm(u, dim=1))


# ## A simple tensile specimen in 2D 

# In[15]:


# Material model (plane stress)
material = IsotropicPlaneStress(E=72000.0, nu=0.2)

# Import mesh
specimen = import_mesh(get_example_file("tensile.vtu"), C=material.C())

# Fixed end BCs
left = specimen.nodes[:, 0] < 0.1
specimen.constraints[left, 0] = True
specimen.constraints[left & (torch.abs(specimen.nodes[:, 1]) < 0.1), 1] = True

# Loaded end BCs
right = specimen.nodes[:, 0] >= 149.9
specimen.constraints[right, 0] = True
specimen.displacements[right, 0] = 5.0


# In[16]:


u, f = specimen.solve()


# In[17]:


specimen.plot(
    u, axes=True, bcs=False, node_property=u[:, 0], linewidth=0.1, figsize=(18, 3)
)


# In[ ]:


# Stress
sigma = specimen.compute_stress(u)
mises = torch.sqrt(
    sigma[:, 0] ** 2
    - sigma[:, 0] * sigma[:, 1]
    + sigma[:, 1] ** 2
    + 3 * sigma[:, 2] ** 2
)
specimen.plot(
    u, axes=True, bcs=False, element_property=mises, linewidth=0.1, figsize=(18, 3)
)