# Source code for ufl_legacy.algorithms.apply_function_pullbacks

    # -*- coding: utf-8 -*-
    """Algorithm for replacing gradients in an expression with reference gradients and coordinate mappings."""

    # Copyright (C) 2008-2016 Martin Sandve Alnæs
    #
    # This file is part of UFL (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    #
    # Modified by Lizao Li <lzlarryli@gmail.com>, 2016

    from itertools import chain, accumulate, repeat

    from ufl_legacy.log import error

    from ufl_legacy.core.multiindex import indices
    from ufl_legacy.corealg.multifunction import MultiFunction, memoized_handler
    from ufl_legacy.algorithms.map_integrands import map_integrand_dags

    from ufl_legacy.classes import (ReferenceValue,
                                    Jacobian, JacobianInverse, JacobianDeterminant)

    from ufl_legacy.tensors import as_tensor, as_vector
    from ufl_legacy.utils.sequences import product
    import numpy

    def sub_elements_with_mappings(element):
        "Return an ordered list of the largest subelements that have a defined mapping."
        if element.mapping() != "undefined":
            return [element]
        elements = []
        for subelm in element.sub_elements():
            if subelm.mapping() != "undefined":
                elements.append(subelm)
            else:
                elements.extend(sub_elements_with_mappings(subelm))
        return elements

    def create_nested_lists(shape):
        if len(shape) == 0:
            return [None]
        elif len(shape) == 1:
            return [None] * shape[0]
        else:
            return [create_nested_lists(shape[1:]) for i in range(shape[0])]

    def reshape_to_nested_list(components, shape):
        if len(shape) == 0:
            assert len(components) == 1
            return [components[0]]
        elif len(shape) == 1:
            assert len(components) == shape[0]
            return components
        else:
            n = product(shape[1:])
            return [reshape_to_nested_list(components[n * i:n * (i + 1)], shape[1:]) for i in range(shape[0])]

    def apply_known_single_pullback(r, element):
        """Apply pullback with given mapping.

        :arg r: Expression wrapped in ReferenceValue
        :arg element: The element defining the mapping
        """
        # Need to pass in r rather than the physical space thing, because
        # the latter may be a ListTensor or similar, rather than a
        # Coefficient/Argument (in the case of mixed elements, see below
        # in apply_single_function_pullbacks), to which we cannot apply ReferenceValue
        mapping = element.mapping()
        domain = r.ufl_domain()
        if mapping == "physical":
            return r
        elif mapping == "identity":
            return r
        elif mapping == "contravariant Piola":
            J = Jacobian(domain)
            detJ = JacobianDeterminant(J)
            transform = (1.0 / detJ) * J
            # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
            *k, i, j = indices(len(r.ufl_shape) + 1)
            kj = (*k, j)
            f = as_tensor(transform[i, j] * r[kj], (*k, i))
            return f
        elif mapping == "covariant Piola":
            K = JacobianInverse(domain)
            # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
            *k, i, j = indices(len(r.ufl_shape) + 1)
            kj = (*k, j)
            f = as_tensor(K[j, i] * r[kj], (*k, i))
            return f
        elif mapping == "L2 Piola":
            detJ = JacobianDeterminant(domain)
            return r / detJ
        elif mapping == "double contravariant Piola":
            J = Jacobian(domain)
            detJ = JacobianDeterminant(J)
            transform = (1.0 / detJ) * J
            # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
            *k, i, j, m, n = indices(len(r.ufl_shape) + 2)
            kmn = (*k, m, n)
            f = as_tensor((1.0 / detJ)**2 * J[i, m] * r[kmn] * J[j, n], (*k, i, j))
            return f
        elif mapping == "double covariant Piola":
            K = JacobianInverse(domain)
            # Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
            *k, i, j, m, n = indices(len(r.ufl_shape) + 2)
            kmn = (*k, m, n)
            f = as_tensor(K[m, i] * r[kmn] * K[n, j], (*k, i, j))
            return f
        else:
            error("Should never be reached!")

    def apply_single_function_pullbacks(r, element):
        """Apply an appropriate pullback to something in physical space

        :arg r: An expression wrapped in ReferenceValue.
        :arg element: The element this expression lives in.
        :returns: a pulled back expression."""
        mapping = element.mapping()
        if r.ufl_shape != element.reference_value_shape():
            error("Expecting reference space expression with shape '%s', got '%s'" % (element.reference_value_shape(), r.ufl_shape))
        if mapping in {"physical", "identity",
                       "contravariant Piola", "covariant Piola",
                       "double contravariant Piola", "double covariant Piola",
                       "L2 Piola"}:
            # Base case in recursion through elements. If the element
            # advertises a mapping we know how to handle, do that
            # directly.
            f = apply_known_single_pullback(r, element)
            if f.ufl_shape != element.value_shape():
                error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape))
            return f
        elif mapping in {"symmetries", "undefined"}:
            # Need to pull back each unique piece of the reference space thing
            gsh = element.value_shape()
            rsh = r.ufl_shape
            if mapping == "symmetries":
                subelem = element.sub_elements()[0]
                fcm = element.flattened_sub_element_mapping()
                offsets = (product(subelem.reference_value_shape()) * i for i in fcm)
                elements = repeat(subelem)
            else:
                elements = sub_elements_with_mappings(element)
                # Python >= 3.8 has an initial keyword argument to
                # accumulate, but 3.7 does not.
                offsets = chain([0],
                                accumulate(product(e.reference_value_shape())
                                           for e in elements))
            rflat = as_vector([r[idx] for idx in numpy.ndindex(rsh)])
            g_components = []
            # For each unique piece in reference space, apply the appropriate pullback
            for offset, subelem in zip(offsets, elements):
                sub_rsh = subelem.reference_value_shape()
                rm = product(sub_rsh)
                rsub = [rflat[offset + i] for i in range(rm)]
                rsub = as_tensor(numpy.asarray(rsub).reshape(sub_rsh))
                rmapped = apply_single_function_pullbacks(rsub, subelem)
                # Flatten into the pulled back expression for the whole thing
                g_components.extend([rmapped[idx]
                                     for idx in numpy.ndindex(rmapped.ufl_shape)])
            # And reshape appropriately
            f = as_tensor(numpy.asarray(g_components).reshape(gsh))
            if f.ufl_shape != element.value_shape():
                error("Expecting pulled back expression with shape '%s', got '%s'" % (element.value_shape(), f.ufl_shape))
            return f
        else:
            error("Unhandled mapping type '%s'" % mapping)

    class FunctionPullbackApplier(MultiFunction):
        def __init__(self):
            MultiFunction.__init__(self)

        expr = MultiFunction.reuse_if_untouched

        def terminal(self, t):
            return t

        @memoized_handler
        def form_argument(self, o):
            # Represent 0-derivatives of form arguments on reference
            # element
            f = apply_single_function_pullbacks(ReferenceValue(o), o.ufl_element())
            assert f.ufl_shape == o.ufl_shape
            return f

    def apply_function_pullbacks(expr):
        """Change representation of coefficients and arguments in expression
        by applying Piola mappings where applicable and representing all
        form arguments in reference value.

        @param expr:
            An Expr.
        """
        return map_integrand_dags(FunctionPullbackApplier(), expr)