%matplotlib widget
%config InlineBackend.figure_format = 'svg'

import matplotlib.pyplot as plt

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

from firedrake import *
set_log_level(ERROR)

mesh = ExtrudedMesh(UnitSquareMesh(10, 10, quadrilateral=True), 10)

p = 5
V = FunctionSpace(mesh, "CG", p)
u = TrialFunction(V)
v = TestFunction(V)
a = dot(grad(u), grad(v)) *dx  # Laplace operator

from tsfc import compile_form

kernel_vanilla, = compile_form(a, parameters={"mode": "vanilla"})

print("Local assembly FLOPs with vanilla mode is {0:.3g}".format(kernel_vanilla.flop_count))

kernel_spectral, = compile_form(a)
print("Local assembly FLOPs with spectral mode is {0:.3g}".format(kernel_spectral.flop_count))

import FIAT, finat

def gauss_lobatto_legendre_line_rule(degree):
    fiat_make_rule = FIAT.quadrature.GaussLobattoLegendreQuadratureLineRule
    fiat_rule = fiat_make_rule(FIAT.ufc_simplex(1), degree + 1)
    finat_ps = finat.point_set.GaussLobattoLegendrePointSet
    points = finat_ps(fiat_rule.get_points())
    weights = fiat_rule.get_weights()
    return finat.quadrature.QuadratureRule(points, weights)

def gauss_lobatto_legendre_cube_rule(dimension, degree):
    make_tensor_rule = finat.quadrature.TensorProductQuadratureRule
    result = gauss_lobatto_legendre_line_rule(degree)
    for _ in range(1, dimension):
        line_rule = gauss_lobatto_legendre_line_rule(degree)
        result = make_tensor_rule([result, line_rule])
    return result

element = FiniteElement('CG', mesh.ufl_cell(), degree=p, variant='spectral')
V = FunctionSpace(mesh, element)
u = TrialFunction(V)
v = TestFunction(V)

gll_quadrature_rule = gauss_lobatto_legendre_cube_rule(dimension=3, degree=p)
a_gll = dot(grad(u), grad(v)) *dx(scheme=gll_quadrature_rule)

kernel_gll, = compile_form(a_gll)
print("Local assembly FLOPs with GLL quadrature is {0:.3g}".format(kernel_gll.flop_count))

from collections import defaultdict
import matplotlib.pyplot as plt
import numpy

flops = defaultdict(list)
ps = range(1, 33)  # polynomial degrees
modes = {
    'gll': {'mode': 'spectral', 'variant': 'spectral', 'rule': gauss_lobatto_legendre_cube_rule},
    'spectral': {'mode': 'spectral', 'variant': None, 'rule': lambda *args: None},
    'vanilla': {'mode': 'vanilla', 'variant': None, 'rule': lambda *args: None}
}

for p in ps:
    for mode in modes:
        element = FiniteElement('CG', mesh.ufl_cell(), degree=p, variant=modes[mode]['variant'])
        V = FunctionSpace(mesh, element)
        u = TrialFunction(V)
        v = TestFunction(V)
        a = dot(grad(u), grad(v))*dx(scheme=modes[mode]['rule'](3, p))
        kernel, = compile_form(a, parameters={"mode": modes[mode]['mode']})
        flops[mode].append(kernel.flop_count)

fig, ax = plt.subplots(1, 1)
ax.set_xscale('log')
ax.set_yscale('log')
for mode in modes:
    ax.plot(ps, flops[mode], label=mode)
x = numpy.linspace(1, 32, 100)
for p, style, offset in zip([5,7,9], ['-.','--',':'], [10, 3, 5]):
    ax.plot(x, numpy.power(x, p)*offset, label=r"$p^{0}$".format(p), color='grey', linestyle=style)
ax.legend(loc='upper left');

# This might take some time to run
flops_curl = defaultdict(list)
ps_curl = range(1, 17)
for p in ps_curl:
    for mode in modes:
        element = FiniteElement('NCE', mesh.ufl_cell(), degree=p, variant=modes[mode]['variant'])
        V = FunctionSpace(mesh, element)
        u = TrialFunction(V)
        v = TestFunction(V)
        a = dot(curl(u), curl(v))*dx(scheme=modes[mode]['rule'](3, p))
        kernel, = compile_form(a, parameters={"mode": modes[mode]['mode']})
        flops_curl[mode].append(kernel.flop_count)

fig, ax = plt.subplots(1, 1)
ax.set_xscale('log')
ax.set_yscale('log')
for mode in modes:
    ax.plot(ps_curl, flops_curl[mode], label=mode)
x = numpy.linspace(1, 16, 100)
for p, style, offset in zip([5,7,9], ['-.','--',':'], [800,40,60]):
    ax.plot(x, numpy.power(x, p)*offset, label=r"$p^{0}$".format(p), color='grey', linestyle=style)
ax.legend(loc='upper left');