# Path: FIAT/reference_element.py
# class UFCInterval(UFCSimplex):
#     """This is the reference interval with vertices (0.0,) and (1.0,)."""
# 
#     def __init__(self):
#         verts = ((0.0,), (1.0,))
#         edges = {0: (0, 1)}
#         topology = {0: {0: (0,), 1: (1,)},
#                     1: edges}
#         super(UFCInterval, self).__init__(LINE, verts, topology)
# 
# class UFCTriangle(UFCSimplex):
#     """This is the reference triangle with vertices (0.0,0.0),
#     (1.0,0.0), and (0.0,1.0)."""
# 
#     def __init__(self):
#         verts = ((0.0, 0.0), (1.0, 0.0), (0.0, 1.0))
#         edges = {0: (1, 2), 1: (0, 2), 2: (0, 1)}
#         faces = {0: (0, 1, 2)}
#         topology = {0: {0: (0,), 1: (1,), 2: (2,)},
#                     1: edges, 2: faces}
#         super(UFCTriangle, self).__init__(TRIANGLE, verts, topology)
# 
#     def compute_normal(self, i):
#         "UFC consistent normal"
#         t = self.compute_tangents(1, i)[0]
#         n = numpy.array((t[1], -t[0]))
#         return n / numpy.linalg.norm(n)
# 
# Path: FIAT/finite_element.py
# def entity_support_dofs(elem, entity_dim):
#     """Return the map of entity id to the degrees of freedom for which the
#     corresponding basis functions take non-zero values
# 
#     :arg elem: FIAT finite element
#     :arg entity_dim: Dimension of the cell subentity.
#     """
#     if not hasattr(elem, "_entity_support_dofs"):
#         elem._entity_support_dofs = {}
#     cache = elem._entity_support_dofs
#     try:
#         return cache[entity_dim]
#     except KeyError:
#         pass
# 
#     ref_el = elem.get_reference_element()
#     dim = ref_el.get_spatial_dimension()
# 
#     entity_cell = ref_el.construct_subelement(entity_dim)
#     quad = create_quadrature(entity_cell, max(2*elem.degree(), 1))
#     weights = quad.get_weights()
# 
#     eps = 1.e-8  # Is this a safe value?
# 
#     result = {}
#     for f in elem.entity_dofs()[entity_dim].keys():
#         entity_transform = ref_el.get_entity_transform(entity_dim, f)
#         points = list(map(entity_transform, quad.get_points()))
# 
#         # Integrate the square of the basis functions on the facet.
#         vals = numpy.double(elem.tabulate(0, points)[(0,) * dim])
#         # Ints contains the square of the basis functions
#         # integrated over the facet.
#         if elem.value_shape():
#             # Vector-valued functions.
#             ints = numpy.dot(numpy.einsum("...ij,...ij->...j", vals, vals), weights)
#         else:
#             ints = numpy.dot(vals**2, weights)
# 
#         result[f] = [dof for dof, i in enumerate(ints) if i > eps]
# 
#     cache[entity_dim] = result
#     return result


import pytest
import FIAT
    import os
from FIAT.reference_element import UFCInterval, UFCTriangle
from FIAT.finite_element import entity_support_dofs
    W0 = FIAT.Hdiv(FIAT.TensorProductElement(W0_h, W0_v))

    W1_v = FIAT.Lagrange(UFCInterval(), 1)
    W1 = FIAT.Hdiv(FIAT.TensorProductElement(W1_h, W1_v))

    elem = FIAT.EnrichedElement(W0, W1)
    assert {0: [0, 1, 2], 1: [0, 1, 3]} == entity_support_dofs(elem, (1, 0))
    assert {0: [0, 2, 3], 1: [1, 2, 3]} == entity_support_dofs(elem, (0, 1))


@pytest.mark.parametrize(('base', 'extr', 'horiz_expected', 'vert_expected'),
                         [(("Discontinuous Lagrange", 0), ("Discontinuous Lagrange", 0),
                           {0: [0], 1: [0]},
                           {0: [0], 1: [0], 2: [0]}),
                          (("Discontinuous Lagrange", 1), ("Discontinuous Lagrange", 1),
                           {0: [0, 2, 4], 1: [1, 3, 5]},
                           {0: [2, 3, 4, 5], 1: [0, 1, 4, 5], 2: [0, 1, 2, 3]}),
                          (("Lagrange", 1), ("Lagrange", 1),
                           {0: [0, 2, 4], 1: [1, 3, 5]},
                           {0: [2, 3, 4, 5], 1: [0, 1, 4, 5], 2: [0, 1, 2, 3]}),
                          (("Discontinuous Lagrange", 0), ("Lagrange", 1),
                           {0: [0], 1: [1]},
                           {0: [0, 1], 1: [0, 1], 2: [0, 1]}),
                          (("Lagrange", 1), ("Discontinuous Lagrange", 0),
                           {0: [0, 1, 2], 1: [0, 1, 2]},
                           {0: [1, 2], 1: [0, 2], 2: [0, 1]})])
def test_prism(base, extr, horiz_expected, vert_expected):
    elem_A = FIAT.supported_elements[base[0]](UFCTriangle(), base[1])
    elem_B = FIAT.supported_elements[extr[0]](UFCInterval(), extr[1])
    elem = FIAT.TensorProductElement(elem_A, elem_B)

# Path: FIAT/reference_element.py
# class UFCInterval(UFCSimplex):
#     """This is the reference interval with vertices (0.0,) and (1.0,)."""
# 
#     def __init__(self):
#         verts = ((0.0,), (1.0,))
#         edges = {0: (0, 1)}
#         topology = {0: {0: (0,), 1: (1,)},
#                     1: edges}
#         super(UFCInterval, self).__init__(LINE, verts, topology)
# 
# class UFCTriangle(UFCSimplex):
#     """This is the reference triangle with vertices (0.0,0.0),
#     (1.0,0.0), and (0.0,1.0)."""
# 
#     def __init__(self):
#         verts = ((0.0, 0.0), (1.0, 0.0), (0.0, 1.0))
#         edges = {0: (1, 2), 1: (0, 2), 2: (0, 1)}
#         faces = {0: (0, 1, 2)}
#         topology = {0: {0: (0,), 1: (1,), 2: (2,)},
#                     1: edges, 2: faces}
#         super(UFCTriangle, self).__init__(TRIANGLE, verts, topology)
# 
#     def compute_normal(self, i):
#         "UFC consistent normal"
#         t = self.compute_tangents(1, i)[0]
#         n = numpy.array((t[1], -t[0]))
#         return n / numpy.linalg.norm(n)
# 
# Path: FIAT/finite_element.py
# def entity_support_dofs(elem, entity_dim):
#     """Return the map of entity id to the degrees of freedom for which the
#     corresponding basis functions take non-zero values
# 
#     :arg elem: FIAT finite element
#     :arg entity_dim: Dimension of the cell subentity.
#     """
#     if not hasattr(elem, "_entity_support_dofs"):
#         elem._entity_support_dofs = {}
#     cache = elem._entity_support_dofs
#     try:
#         return cache[entity_dim]
#     except KeyError:
#         pass
# 
#     ref_el = elem.get_reference_element()
#     dim = ref_el.get_spatial_dimension()
# 
#     entity_cell = ref_el.construct_subelement(entity_dim)
#     quad = create_quadrature(entity_cell, max(2*elem.degree(), 1))
#     weights = quad.get_weights()
# 
#     eps = 1.e-8  # Is this a safe value?
# 
#     result = {}
#     for f in elem.entity_dofs()[entity_dim].keys():
#         entity_transform = ref_el.get_entity_transform(entity_dim, f)
#         points = list(map(entity_transform, quad.get_points()))
# 
#         # Integrate the square of the basis functions on the facet.
#         vals = numpy.double(elem.tabulate(0, points)[(0,) * dim])
#         # Ints contains the square of the basis functions
#         # integrated over the facet.
#         if elem.value_shape():
#             # Vector-valued functions.
#             ints = numpy.dot(numpy.einsum("...ij,...ij->...j", vals, vals), weights)
#         else:
#             ints = numpy.dot(vals**2, weights)
# 
#         result[f] = [dof for dof, i in enumerate(ints) if i > eps]
# 
#     cache[entity_dim] = result
#     return result


# Path: test/unit/test_facet_support_dofs.py
import pytest
import FIAT
    import os
from FIAT.reference_element import UFCInterval, UFCTriangle
from FIAT.finite_element import entity_support_dofs

    W0 = FIAT.Hdiv(FIAT.TensorProductElement(W0_h, W0_v))

    W1_v = FIAT.Lagrange(UFCInterval(), 1)
    W1 = FIAT.Hdiv(FIAT.TensorProductElement(W1_h, W1_v))

    elem = FIAT.EnrichedElement(W0, W1)
    assert {0: [0, 1, 2], 1: [0, 1, 3]} == entity_support_dofs(elem, (1, 0))
    assert {0: [0, 2, 3], 1: [1, 2, 3]} == entity_support_dofs(elem, (0, 1))


@pytest.mark.parametrize(('base', 'extr', 'horiz_expected', 'vert_expected'),
                         [(("Discontinuous Lagrange", 0), ("Discontinuous Lagrange", 0),
                           {0: [0], 1: [0]},
                           {0: [0], 1: [0], 2: [0]}),
                          (("Discontinuous Lagrange", 1), ("Discontinuous Lagrange", 1),
                           {0: [0, 2, 4], 1: [1, 3, 5]},
                           {0: [2, 3, 4, 5], 1: [0, 1, 4, 5], 2: [0, 1, 2, 3]}),
                          (("Lagrange", 1), ("Lagrange", 1),
                           {0: [0, 2, 4], 1: [1, 3, 5]},
                           {0: [2, 3, 4, 5], 1: [0, 1, 4, 5], 2: [0, 1, 2, 3]}),
                          (("Discontinuous Lagrange", 0), ("Lagrange", 1),
                           {0: [0], 1: [1]},
                           {0: [0, 1], 1: [0, 1], 2: [0, 1]}),
                          (("Lagrange", 1), ("Discontinuous Lagrange", 0),
                           {0: [0, 1, 2], 1: [0, 1, 2]},
                           {0: [1, 2], 1: [0, 2], 2: [0, 1]})])
def test_prism(base, extr, horiz_expected, vert_expected):
    elem_A = FIAT.supported_elements[base[0]](UFCTriangle(), base[1])
    elem_B = FIAT.supported_elements[extr[0]](UFCInterval(), extr[1])
    elem = FIAT.TensorProductElement(elem_A, elem_B)

    assert horiz_expected == entity_support_dofs(elem, (2, 0))