# ==============================================
# Defining conforming Crouzeix--Raviart elements
# ==============================================
#
# In this demo, we show how Basix's custom element functionality can be used to
# create a conforming Crouzeix--Raviart element.
#
# First, we import Basix and Numpy.

import basix
import numpy as np
from basix import CellType, MapType, PolynomialType, LatticeType

from mpl_toolkits import mplot3d  # noqa: F401
import matplotlib.pyplot as plt

# Conforming CR element on a triangle
# ===================================
#
# We begin by implementing this element on a triangle. The following function
# implements this element for an arbitrary degree. Details of the definition of
# this element can be found at
# https://defelement.com/elements/conforming-crouzeix-raviart.html.
#
# As the input to this function, we use the degree of the element as shown on
# DefElement. For most degrees, the highest degree polynomial in this elements
# polynomial space is actually one degree higher, so we pass `degree + 1` into
# `create_custom_element`.


def create_ccr_triangle(degree):

    if degree == 1:
        wcoeffs = np.eye(3)

        x = [[], [], [], []]
        x[0].append(np.array([[0.0, 0.0]]))
        x[0].append(np.array([[1.0, 0.0]]))
        x[0].append(np.array([[0.0, 1.0]]))

        for _ in range(3):
            x[1].append(np.zeros((0, 2)))
        x[2].append(np.zeros((0, 2)))

        M = [[], [], [], []]
        for _ in range(3):
            M[0].append(np.array([[[1.]]]))

        for _ in range(3):
            M[1].append(np.zeros((0, 1, 0)))
        M[2].append(np.zeros((0, 1, 0)))

        return basix.create_custom_element(
            CellType.triangle, [], wcoeffs, x, M, MapType.identity, False, 1, 1)

    npoly = (degree + 2) * (degree + 3) // 2
    ndofs = degree * (degree + 5) // 2
    wcoeffs = np.zeros((ndofs, npoly))

    dof_n = 0
    for i in range((degree + 1) * (degree + 2) // 2):
        wcoeffs[dof_n, dof_n] = 1
        dof_n += 1

    pts, wts = basix.make_quadrature(CellType.triangle, 2 * (degree + 1))
    poly = basix.tabulate_polynomials(PolynomialType.legendre, CellType.triangle, degree + 1, pts)
    for i in range(1, degree):
        x = pts[:, 0]
        y = pts[:, 1]
        f = x ** i * y ** (degree - i) * (x + y)

        for j in range(npoly):
            wcoeffs[dof_n, j] = sum(f * poly[:, j] * wts)
        dof_n += 1

    geometry = basix.geometry(CellType.triangle)
    topology = basix.topology(CellType.triangle)
    x = [[], [], [], []]
    M = [[], [], [], []]
    for v in topology[0]:
        x[0].append(np.array(geometry[v]))
        M[0].append(np.array([[[1.]]]))
    pts = basix.create_lattice(CellType.interval, degree, LatticeType.equispaced, False)
    mat = np.zeros((len(pts), 1, len(pts)))
    mat[:, 0, :] = np.eye(len(pts))
    for e in topology[1]:
        edge_pts = []
        v0 = geometry[e[0]]
        v1 = geometry[e[1]]
        for p in pts:
            edge_pts.append(v0 + p * (v1 - v0))
        x[1].append(np.array(edge_pts))
        M[1].append(mat)
    pts = basix.create_lattice(CellType.triangle, degree + 1, LatticeType.equispaced, False)
    x[2].append(pts)
    mat = np.zeros((len(pts), 1, len(pts)))
    mat[:, 0, :] = np.eye(len(pts))
    M[2].append(mat)

    return basix.create_custom_element(
        CellType.triangle, [], wcoeffs, x, M, MapType.identity, False, degree, degree + 1)


# We can then create a degree 2 conforming CR element.

e = create_ccr_triangle(2)

# We now visualise the basis functions of the element we have created.

pts = basix.create_lattice(CellType.triangle, 30, LatticeType.equispaced, True)

x = pts[:, 0]
y = pts[:, 1]
z = e.tabulate(0, pts)[0]

fig = plt.figure(figsize=(8, 8))

for n in range(7):
    if n == 6:
        ax = plt.subplot(3, 3, n + 2, projection='3d')
    else:
        ax = plt.subplot(3, 3, n + 1, projection='3d')
    ax.plot([0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0], "k-")
    ax.scatter(x, y, z[:, n, 0])

plt.savefig("ccr_triangle_2.png")

# .. image:: ccr_triangle_2.png
#  :width: 100%
#  :alt: The basis functions of a degree 2 conforming CR element

# We also visualise the basis functions of a degree 3 conforming CR element.

e = create_ccr_triangle(3)

pts = basix.create_lattice(CellType.triangle, 30, LatticeType.equispaced, True)

x = pts[:, 0]
y = pts[:, 1]
z = e.tabulate(0, pts)[0]

fig = plt.figure(figsize=(11, 8))

for n in range(12):
    ax = plt.subplot(3, 4, n + 1, projection='3d')
    ax.plot([0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0], "k-")
    ax.scatter(x, y, z[:, n, 0])

plt.savefig("ccr_triangle_3.png")

# .. image:: ccr_triangle_3.png
#  :width: 100%
#  :alt: The basis functions of a degree 3 conforming CR element
