# Source code for FIAT.tensor_product

    # Copyright (C) 2008 Robert C. Kirby (Texas Tech University)
    # Copyright (C) 2013 Andrew T. T. McRae
    # Modified by Thomas H. Gibson, 2016
    #
    # This file is part of FIAT (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later

    import numpy
    from FIAT.finite_element import FiniteElement
    from FIAT.reference_element import TensorProductCell, UFCQuadrilateral, UFCHexahedron, flatten_entities, compute_unflattening_map
    from FIAT.dual_set import DualSet
    from FIAT.polynomial_set import mis
    from FIAT import dual_set
    from FIAT import functional

    def _first_point(node):
        return tuple(node.get_point_dict().keys())[0]

    def _first_point_pair(node):
        return tuple(node.get_point_dict().items())[0]

    class TensorProductElement(FiniteElement):
        """Class implementing a finite element that is the tensor product
        of two existing finite elements."""

        def __init__(self, A, B):
            # set up simple things
            order = min(A.get_order(), B.get_order())
            if A.get_formdegree() is None or B.get_formdegree() is None:
                formdegree = None
            else:
                formdegree = A.get_formdegree() + B.get_formdegree()

            # set up reference element
            ref_el = TensorProductCell(A.get_reference_element(),
                                       B.get_reference_element())

            if A.mapping()[0] != "affine" and B.mapping()[0] == "affine":
                mapping = A.mapping()[0]
            elif B.mapping()[0] != "affine" and A.mapping()[0] == "affine":
                mapping = B.mapping()[0]
            elif A.mapping()[0] == "affine" and B.mapping()[0] == "affine":
                mapping = "affine"
            else:
                raise ValueError("check tensor product mappings - at least one must be affine")

            # set up entity_ids
            Adofs = A.entity_dofs()
            Bdofs = B.entity_dofs()
            Bsdim = B.space_dimension()
            entity_ids = {}

            for curAdim in Adofs:
                for curBdim in Bdofs:
                    entity_ids[(curAdim, curBdim)] = {}
                    dim_cur = 0
                    for entityA in Adofs[curAdim]:
                        for entityB in Bdofs[curBdim]:
                            entity_ids[(curAdim, curBdim)][dim_cur] = \
                                [x*Bsdim + y for x in Adofs[curAdim][entityA]
                                    for y in Bdofs[curBdim][entityB]]
                            dim_cur += 1

            # set up dual basis
            Anodes = A.dual_basis()
            Bnodes = B.dual_basis()

            # build the dual set by inspecting the current dual
            # sets item by item.
            # Currently supported cases:
            # PointEval x PointEval = PointEval [scalar x scalar = scalar]
            # PointScaledNormalEval x PointEval = PointScaledNormalEval [vector x scalar = vector]
            # ComponentPointEvaluation x PointEval [vector x scalar = vector]
            nodes = []
            for Anode in Anodes:
                if isinstance(Anode, functional.PointEvaluation):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: PointEval x PointEval
                            # the PointEval functional just requires the
                            # coordinates. these are currently stored as
                            # the key of a one-item dictionary. we retrieve
                            # these by calling get_point_dict(), and
                            # use the concatenation to make a new PointEval
                            nodes.append(functional.PointEvaluation(ref_el, _first_point(Anode) + _first_point(Bnode)))
                        elif isinstance(Bnode, functional.IntegralMoment):
                            # dummy functional for product with integral moments
                            nodes.append(functional.Functional(None, None, None,
                                                               {}, "Undefined"))
                        elif isinstance(Bnode, functional.PointDerivative):
                            # dummy functional for product with point derivative
                            nodes.append(functional.Functional(None, None, None,
                                                               {}, "Undefined"))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.PointScaledNormalEvaluation):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: PointScaledNormalEval x PointEval
                            # this could be wrong if the second shape
                            # has spatial dimension >1, since we are not
                            # explicitly scaling by facet size
                            if len(_first_point(Bnode)) > 1:
                                # TODO: support this case one day
                                raise NotImplementedError("PointScaledNormalEval x PointEval is not yet supported if the second shape has dimension > 1")
                            # We cannot make a new functional.PSNEval in
                            # the natural way, since it tries to compute
                            # the normal vector by itself.
                            # Instead, we create things manually, and
                            # call Functional() with these arguments
                            sd = ref_el.get_spatial_dimension()
                            # The pt_dict is a one-item dictionary containing
                            # the details of the functional.
                            # The key is the spatial coordinate, which
                            # is just a concatenation of the two parts.
                            # The value is a list of tuples, representing
                            # the normal vector (scaled by the volume of
                            # the facet) at that point.
                            # Each tuple looks like (foo, (i,)); the i'th
                            # component of the scaled normal is foo.

                            # The following line is only valid when the second
                            # shape has spatial dimension 1 (enforced above)
                            Apoint, Avalue = _first_point_pair(Anode)
                            pt_dict = {Apoint + _first_point(Bnode): Avalue + [(0.0, (len(Apoint),))]}

                            # The following line should be used in the
                            # general case
                            # pt_dict = {Anode.get_point_dict().keys()[0] + Bnode.get_point_dict().keys()[0]: Anode.get_point_dict().values()[0] + [(0.0, (ii,)) for ii in range(len(Anode.get_point_dict().keys()[0]), len(Anode.get_point_dict().keys()[0]) + len(Bnode.get_point_dict().keys()[0]))]}

                            # THE FOLLOWING IS PROBABLY CORRECT BUT UNTESTED
                            shp = (sd,)
                            nodes.append(functional.Functional(ref_el, shp, pt_dict, {}, "PointScaledNormalEval"))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.PointEdgeTangentEvaluation):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: PointEdgeTangentEval x PointEval
                            # this is very similar to the case above, so comments omitted
                            if len(_first_point(Bnode)) > 1:
                                raise NotImplementedError("PointEdgeTangentEval x PointEval is not yet supported if the second shape has dimension > 1")
                            sd = ref_el.get_spatial_dimension()
                            Apoint, Avalue = _first_point_pair(Anode)
                            pt_dict = {Apoint + _first_point(Bnode): Avalue + [(0.0, (len(Apoint),))]}

                            # THE FOLLOWING IS PROBABLY CORRECT BUT UNTESTED
                            shp = (sd,)
                            nodes.append(functional.Functional(ref_el, shp, pt_dict, {}, "PointEdgeTangent"))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.ComponentPointEvaluation):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: ComponentPointEval x PointEval
                            # the CptPointEval functional requires the component
                            # and the coordinates. very similar to PE x PE case.
                            sd = ref_el.get_spatial_dimension()
                            nodes.append(functional.ComponentPointEvaluation(ref_el, Anode.comp, (sd,), _first_point(Anode) + _first_point(Bnode)))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.FrobeniusIntegralMoment):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: FroIntMom x PointEval
                            sd = ref_el.get_spatial_dimension()
                            pt_dict = {}
                            pt_old = Anode.get_point_dict()
                            for pt in pt_old:
                                pt_dict[pt+_first_point(Bnode)] = pt_old[pt] + [(0.0, sd-1)]
                            # THE FOLLOWING IS PROBABLY CORRECT BUT UNTESTED
                            shp = (sd,)
                            nodes.append(functional.Functional(ref_el, shp, pt_dict, {}, "FrobeniusIntegralMoment"))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.IntegralMoment):
                    for Bnode in Bnodes:
                        if isinstance(Bnode, functional.PointEvaluation):
                            # case: IntMom x PointEval
                            sd = ref_el.get_spatial_dimension()
                            pt_dict = {}
                            pt_old = Anode.get_point_dict()
                            for pt in pt_old:
                                pt_dict[pt+_first_point(Bnode)] = pt_old[pt]
                            # THE FOLLOWING IS PROBABLY CORRECT BUT UNTESTED
                            shp = (sd,)
                            nodes.append(functional.Functional(ref_el, shp, pt_dict, {}, "IntegralMoment"))
                        else:
                            raise NotImplementedError("unsupported functional type")

                elif isinstance(Anode, functional.Functional):
                    # this should catch everything else
                    for Bnode in Bnodes:
                        nodes.append(functional.Functional(None, None, None, {}, "Undefined"))
                else:
                    raise NotImplementedError("unsupported functional type")

            dual = dual_set.DualSet(nodes, ref_el, entity_ids)

            super(TensorProductElement, self).__init__(ref_el, dual, order, formdegree, mapping)
            # Set up constituent elements
            self.A = A
            self.B = B

            # degree for quadrature rule
            self.polydegree = max(A.degree(), B.degree())

        def degree(self):
            """Return the degree of the (embedding) polynomial space."""
            return self.polydegree

        def get_nodal_basis(self):
            """Return the nodal basis, encoded as a PolynomialSet object,
            for the finite element."""
            raise NotImplementedError("get_nodal_basis not implemented")

        def get_coeffs(self):
            """Return the expansion coefficients for the basis of the
            finite element."""
            raise NotImplementedError("get_coeffs not implemented")

        def tabulate(self, order, points, entity=None):
            """Return tabulated values of derivatives up to given order of
            basis functions at given points."""
            if entity is None:
                entity = (self.ref_el.get_dimension(), 0)
            entity_dim, entity_id = entity

            shape = tuple(len(c.get_topology()[d])
                          for c, d in zip(self.ref_el.cells, entity_dim))
            idA, idB = numpy.unravel_index(entity_id, shape)

            # Factor the entity argument to get entities of the component elements
            entityA_dim, entityB_dim = entity_dim
            entityA = (entityA_dim, idA)
            entityB = (entityB_dim, idB)

            pointsAdim, pointsBdim = [c.get_spatial_dimension()
                                      for c in self.ref_el.construct_subelement(entity_dim).cells]
            pointsA = [point[:pointsAdim] for point in points]
            pointsB = [point[pointsAdim:pointsAdim + pointsBdim] for point in points]

            Asdim = self.A.ref_el.get_spatial_dimension()
            Bsdim = self.B.ref_el.get_spatial_dimension()
            # Note that for entities other than cells, the following
            # tabulations are already appropriately zero-padded so no
            # additional zero padding is required.
            Atab = self.A.tabulate(order, pointsA, entityA)
            Btab = self.B.tabulate(order, pointsB, entityB)
            npoints = len(points)

            # allow 2 scalar-valued FE spaces, or 1 scalar-valued,
            # 1 vector-valued. Combining 2 vector-valued spaces
            # into a tensor-valued space via an outer-product
            # seems to be a sensible general option, but I don't
            # know how to handle the nestedness of the arrays
            # if someone then tries to make a new "tensor finite
            # element" where one component is already a
            # tensor-valued space!
            A_valuedim = len(self.A.value_shape())  # scalar: 0, vector: 1
            B_valuedim = len(self.B.value_shape())  # scalar: 0, vector: 1
            if A_valuedim + B_valuedim > 1:
                raise NotImplementedError("tabulate does not support two vector-valued inputs")
            result = {}
            for i in range(order + 1):
                alphas = mis(Asdim+Bsdim, i)  # thanks, Rob!
                for alpha in alphas:
                    if A_valuedim == 0 and B_valuedim == 0:
                        # for each point, get outer product of (A's basis
                        # functions f1, f2, ... evaluated at that point)
                        # with (B's basis functions g1, g2, ... evaluated
                        # at that point). This gives temp[point][f_i][g_j].
                        # Flatten this, so bfs are
                        # in the order f1g1, f1g2, ..., f2g1, f2g2, ...
                        # which is compatible with the entity_dofs order.
                        # We now have temp[point][full basis function]
                        # Transpose this to get temp[bf][point],
                        # and we are done.
                        temp = numpy.array([numpy.outer(
                                           Atab[alpha[0:Asdim]][..., j],
                                           Btab[alpha[Asdim:Asdim+Bsdim]][..., j])
                            .ravel() for j in range(npoints)])
                        result[alpha] = temp.transpose()
                    elif A_valuedim == 1 and B_valuedim == 0:
                        # similar to above, except A's basis functions
                        # are now vector-valued. numpy.outer flattens the
                        # array, so it's like taking the OP of
                        # f1_x, f1_y, f2_x, f2_y, ... with g1, g2, ...
                        # this gives us
                        # temp[point][f1x, f1y, f2x, f2y, ...][g_j].
                        # reshape once to get temp[point][f_i][x/y][g_j]
                        # transpose to get temp[point][x/y][f_i][g_j]
                        # reshape to flatten the last two indices, this
                        # gives us temp[point][x/y][full bf_i]
                        # finally, transpose the first and last indices
                        # to get temp[bf_i][x/y][point], and we are done.
                        temp = numpy.array([numpy.outer(
                                           Atab[alpha[0:Asdim]][..., j],
                                           Btab[alpha[Asdim:Asdim+Bsdim]][..., j])
                            for j in range(npoints)])
                        assert temp.shape[1] % 2 == 0
                        temp2 = temp.reshape((temp.shape[0],
                                              temp.shape[1]//2,
                                              2,
                                              temp.shape[2]))\
                            .transpose(0, 2, 1, 3)\
                            .reshape((temp.shape[0], 2, -1))\
                            .transpose(2, 1, 0)
                        result[alpha] = temp2
                    elif A_valuedim == 0 and B_valuedim == 1:
                        # as above, with B's functions now vector-valued.
                        # we now do... [numpy.outer ... for ...] gives
                        # temp[point][f_i][g1x,g1y,g2x,g2y,...].
                        # reshape to temp[point][f_i][g_j][x/y]
                        # flatten middle: temp[point][full bf_i][x/y]
                        # transpose to temp[bf_i][x/y][point]
                        temp = numpy.array([numpy.outer(
                            Atab[alpha[0:Asdim]][..., j],
                            Btab[alpha[Asdim:Asdim+Bsdim]][..., j])
                            for j in range(len(Atab[alpha[0:Asdim]][0]))])
                        assert temp.shape[2] % 2 == 0
                        temp2 = temp.reshape((temp.shape[0], temp.shape[1],
                                              temp.shape[2]//2, 2))\
                            .reshape((temp.shape[0], -1, 2))\
                            .transpose(1, 2, 0)
                        result[alpha] = temp2
            return result

        def value_shape(self):
            """Return the value shape of the finite element functions."""
            if len(self.A.value_shape()) == 0 and len(self.B.value_shape()) == 0:
                return ()
            elif len(self.A.value_shape()) == 1 and len(self.B.value_shape()) == 0:
                return (self.A.value_shape()[0],)
            elif len(self.A.value_shape()) == 0 and len(self.B.value_shape()) == 1:
                return (self.B.value_shape()[0],)
            else:
                raise NotImplementedError("value_shape not implemented")

        def dmats(self):
            """Return dmats: expansion coefficients for basis function
            derivatives."""
            raise NotImplementedError("dmats not implemented")

        def get_num_members(self, arg):
            """Return number of members of the expansion set."""
            raise NotImplementedError("get_num_members not implemented")

        def is_nodal(self):
            # This element is nodal iff all factor elements are nodal.
            return all([self.A.is_nodal(), self.B.is_nodal()])

    class FlattenedDimensions(FiniteElement):
        """A wrapper class that flattens entity dimensions of a FIAT element defined
        on a TensorProductCell to one with quadrilateral/hexahedron entities.
        TensorProductCell has dimension defined as a tuple of factor element dimensions
        (i, j) in 2D  and (i, j, k) in 3D.
        Flattened dimension is a sum of the tuple elements."""

        def __init__(self, element):

            nodes = element.dual.nodes
            dim = element.ref_el.get_spatial_dimension()

            if dim == 2:
                ref_el = UFCQuadrilateral()
            elif dim == 3:
                ref_el = UFCHexahedron()
            else:
                raise ValueError("Illegal element dimension %s" % dim)

            entity_ids = element.dual.entity_ids

            flat_entity_ids = flatten_entities(entity_ids)
            dual = DualSet(nodes, ref_el, flat_entity_ids)
            super(FlattenedDimensions, self).__init__(ref_el, dual, element.get_order(), element.get_formdegree(), element._mapping)
            self.element = element

            # Construct unflattening map for passing correct values to tabulate()
            self.unflattening_map = compute_unflattening_map(self.element.ref_el.get_topology())

        def degree(self):
            """Return the degree of the (embedding) polynomial space."""
            return self.element.degree()

        def tabulate(self, order, points, entity=None):
            """Return tabulated values of derivatives up to given order of
            basis functions at given points."""
            if entity is None:
                entity = (self.get_reference_element().get_spatial_dimension(), 0)

            # Entity is provided in flattened form (d, i)
            # Appropriate product entity is taken from the unflattening_map dict
            entity_dim, entity_id = entity
            product_entity = self.unflattening_map[(entity_dim, entity_id)]

            return self.element.tabulate(order, points, product_entity)

        def value_shape(self):
            """Return the value shape of the finite element functions."""
            return self.element.value_shape()

        def get_nodal_basis(self):
            """Return the nodal basis, encoded as a PolynomialSet object,
            for the finite element."""
            raise self.element.get_nodal_basis()

        def get_coeffs(self):
            """Return the expansion coefficients for the basis of the
            finite element."""
            raise self.element.get_coeffs()

        def dmats(self):
            """Return dmats: expansion coefficients for basis function
            derivatives."""
            raise self.element.dmats()

        def get_num_members(self, arg):
            """Return number of members of the expansion set."""
            raise self.element.get_num_members(arg)

        def is_nodal(self):
            # This element is nodal iff unflattened element is nodal.
            return self.element.is_nodal()