# Source code for ufl.algorithms.replace_derivative_nodes

    """Algorithm for replacing derivative nodes in a BaseForm or Expr."""

    import ufl
    from ufl.algorithms.analysis import extract_arguments
    from ufl.algorithms.map_integrands import map_integrand_dags
    from ufl.constantvalue import as_ufl
    from ufl.corealg.multifunction import MultiFunction
    from ufl.tensors import ListTensor

    class DerivativeNodeReplacer(MultiFunction):
        """Replace derivative nodes with new derivative nodes."""

        def __init__(self, mapping, **derivative_kwargs):
            """Initialise."""
            super().__init__()
            self.mapping = mapping
            self.der_kwargs = derivative_kwargs

        ufl_type = MultiFunction.reuse_if_untouched

        def coefficient_derivative(self, cd, o, coefficients, arguments, coefficient_derivatives):
            """Apply to coefficient_derivative."""
            der_kwargs = self.der_kwargs
            new_coefficients = tuple(
                self.mapping[c] if c in self.mapping.keys() else c for c in coefficients.ufl_operands
            )

            # Ensure type compatibility for arguments!
            if "argument" not in der_kwargs.keys():
                # Argument's number/part can be retrieved from the former coefficient derivative.
                arguments = arguments.ufl_operands
                new_arguments = ()
                for c, a in zip(new_coefficients, arguments):
                    if isinstance(a, ListTensor):
                        (a,) = extract_arguments(a)
                    new_arguments += (type(a)(c.ufl_function_space(), a.number(), a.part()),)
                der_kwargs.update({"argument": new_arguments})

            return ufl.derivative(o, new_coefficients, **der_kwargs)

    def replace_derivative_nodes(expr, mapping, **derivative_kwargs):
        """Replaces derivative nodes.

        Replaces the variable with respect to which the derivative is taken.
        This is called during apply_derivatives to treat delayed derivatives.

        Example: Let u be a Coefficient, N an ExternalOperator independent
            of u (i.e. N's operands don't depend on u),
            and let uhat and Nhat be Arguments.

            F = u ** 2 * N * dx
            dFdu = derivative(F, u, uhat)
            dFdN = replace_derivative_nodes(dFdu, {u: N}, argument=Nhat)

            Then, by subsequently expanding the derivatives we have:

            dFdu -> 2 * u * uhat * N * dx
            dFdN -> u ** 2 * Nhat * dx

        Args:
            expr:
                An Expr or BaseForm.
            mapping:
                A dict with from:to replacements to perform.
            derivative_kwargs:
                A dict containing the keyword arguments for derivative (i.e. `argument` and
                `coefficient_derivatives`).
        """
        mapping2 = dict((k, as_ufl(v)) for (k, v) in mapping.items())
        return map_integrand_dags(DerivativeNodeReplacer(mapping2, **derivative_kwargs), expr)