# Code in this cell makes plots appear an appropriate size and resolution in the browser window
%matplotlib widget
%config InlineBackend.figure_format = 'svg'

import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = (11, 4)

from firedrake import *
length = 1
width = 0.2
mesh = RectangleMesh(40, 20, length, width)

V = VectorFunctionSpace(mesh, "Lagrange", 1)

bc = DirichletBC(V, Constant([0, 0]), 1)

rho = Constant(0.01)
g = Constant(1)
f = as_vector([0, -rho*g])
mu = Constant(1)
lambda_ = Constant(0.25)
Id = Identity(mesh.geometric_dimension()) # 2x2 Identity tensor

def epsilon(u):
    return 0.5*(grad(u) + grad(u).T)

def sigma(u):
    return lambda_*div(u)*Id + 2*mu*epsilon(u)

u = TrialFunction(V)
v = TestFunction(V)
a = inner(sigma(u), epsilon(v))*dx
L = dot(f, v)*dx

uh = Function(V)
solve(a == L, uh, bcs=bc, solver_parameters={"ksp_monitor": None})

displaced_coordinates = interpolate(SpatialCoordinate(mesh) + uh, V)
displaced_mesh = Mesh(displaced_coordinates)

# NBVAL_IGNORE_OUTPUT
from firedrake.pyplot import triplot

fig, axes = plt.subplots()
triplot(displaced_mesh, axes=axes)
axes.set_aspect("equal");



def solve_elasticity(nx, ny, options=None, **kwargs):
    length = 1
    width = 0.2
    mesh = RectangleMesh(nx, ny, length, width)
    V = VectorFunctionSpace(mesh, "Lagrange", 1)
    rho = Constant(0.01)
    g = Constant(1)
    f = as_vector([0, -rho*g])
    mu = Constant(1)
    lambda_ = Constant(0.25)
    Id = Identity(mesh.geometric_dimension()) # 2x2 Identity tensor
    
    bc = DirichletBC(V, Constant([0, 0]), 1)
    u = TrialFunction(V)
    v = TestFunction(V)
    a = inner(sigma(u), epsilon(v))*dx
    L = dot(f, v)*dx

    uh = Function(V)
    solve(a == L, uh, bcs=bc, solver_parameters=options, **kwargs)
    return uh

# NBVAL_RAISES_EXCEPTION
uh = solve_elasticity(100, 100, options={"ksp_max_it": 100, "ksp_type": "cg", "pc_type": "none"})

uh = solve_elasticity(200, 200, options={"ksp_type": "cg", 
                                         "ksp_max_it": 100, 
                                         "pc_type": "gamg",
                                         "pc_gamg_aggressive_square_graph": None,
                                         "pc_gamg_mis_k_minimum_degree_ordering": True,
                                         "mg_levels_pc_type": "sor",
                                         "mat_type": "aij",
                                         "ksp_converged_reason": None})

def solve_elasticity(nx, ny, options=None, **kwargs):
    length = 1
    width = 0.2
    mesh = RectangleMesh(nx, ny, length, width)
    V = VectorFunctionSpace(mesh, "CG", 1)
    rho = Constant(0.01)
    g = Constant(1)
    f = as_vector([0, -rho*g])
    mu = Constant(1)
    lambda_ = Constant(0.25)
    Id = Identity(mesh.geometric_dimension()) # 2x2 Identity tensor
    def epsilon(u):
        return 0.5*(grad(u) + grad(u).T)

    def sigma(u):
        return lambda_*div(u)*Id + 2*mu*epsilon(u)    
    bc = DirichletBC(V, Constant([0, 0]), 1)
    u = TrialFunction(V)
    v = TestFunction(V)
    a = inner(sigma(u), epsilon(v))*dx
    L = dot(f, v)*dx

    # create rigid body modes
    x, y = SpatialCoordinate(mesh)
    b0 = Function(V)
    b1 = Function(V)
    b2 = Function(V)
    b0.interpolate(Constant([1, 0]))
    b1.interpolate(Constant([0, 1]))
    b2.interpolate(as_vector([-y, x]))
    nullmodes = VectorSpaceBasis([b0, b1, b2])
    # Make sure they're orthonormal.
    nullmodes.orthonormalize()
    uh = Function(V)
    solve(a == L, uh, bcs=bc, solver_parameters=options, near_nullspace=nullmodes)
    return uh

uh = solve_elasticity(200, 200, options={"ksp_type": "cg", 
                                         "ksp_max_it": 100, 
                                         "pc_type": "gamg",
                                         "pc_gamg_aggressive_square_graph": None,
                                         "pc_gamg_mis_k_minimum_degree_ordering": True,
                                         "mat_type": "aij",
                                         "ksp_converged_reason": None})