# Source code for ufl.differentiation

    """Differential operators."""

    # 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

    from ufl.argument import Argument, Coargument
    from ufl.checks import is_cellwise_constant
    from ufl.coefficient import Coefficient
    from ufl.constantvalue import Zero
    from ufl.core.base_form_operator import BaseFormOperator
    from ufl.core.expr import Expr
    from ufl.core.operator import Operator
    from ufl.core.terminal import Terminal
    from ufl.core.ufl_type import ufl_type
    from ufl.domain import extract_unique_domain, find_geometric_dimension
    from ufl.exprcontainers import ExprList, ExprMapping
    from ufl.form import BaseForm
    from ufl.precedence import parstr
    from ufl.variable import Variable

    # --- Basic differentiation objects ---

    @ufl_type(is_abstract=True, is_differential=True)
    class Derivative(Operator):
        """Base class for all derivative types."""

        __slots__ = ()

        def __init__(self, operands):
            """Initalise."""
            Operator.__init__(self, operands)

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class CoefficientDerivative(Derivative):
        """Derivative of form integrand w.r.t. the degrees of freedom in a discrete Coefficient."""

        __slots__ = ()

        def __new__(cls, integrand, coefficients, arguments, coefficient_derivatives):
            """Create a new CoefficientDerivative."""
            if not isinstance(coefficients, ExprList):
                raise ValueError("Expecting ExprList instance with Coefficients.")
            if not isinstance(arguments, ExprList):
                raise ValueError("Expecting ExprList instance with Arguments.")
            if not isinstance(coefficient_derivatives, ExprMapping):
                raise ValueError("Expecting ExprMapping for coefficient derivatives.")
            if isinstance(integrand, Zero):
                return integrand
            return Derivative.__new__(cls)

        def __init__(self, integrand, coefficients, arguments, coefficient_derivatives):
            """Initalise."""
            if not isinstance(coefficient_derivatives, ExprMapping):
                coefficient_derivatives = ExprMapping(coefficient_derivatives)
            Derivative.__init__(self, (integrand, coefficients, arguments, coefficient_derivatives))

        def __str__(self):
            """Format as a string."""
            return (
                f"d/dfj {{ {self.ufl_operands[0]} }}, with fh={self.ufl_operands[1]}, dfh/dfj = "
                f"{self.ufl_operands[2]}, and coefficient derivatives {self.ufl_operands[3]}"
            )

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class CoordinateDerivative(CoefficientDerivative):
        """Derivative of the integrand of a form w.r.t. the SpatialCoordinates."""

        __slots__ = ()

        def __str__(self):
            """Format as a string."""
            return (
                f"d/dfj {{ {self.ufl_operands[0]} }}, with fh={self.ufl_operands[1]}, dfh/dfj = "
                f"{self.ufl_operands[2]}, and coordinate derivatives {self.ufl_operands[3]}"
            )

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class BaseFormDerivative(CoefficientDerivative, BaseForm):
        """Derivative of a base form w.r.t the degrees of freedom in a discrete Coefficient."""

        _ufl_noslots_ = True
        _ufl_required_methods_: tuple[str, ...] = (
            CoefficientDerivative._ufl_required_methods_ + BaseForm._ufl_required_methods_
        )

        def __init__(self, base_form, coefficients, arguments, coefficient_derivatives):
            """Initalise."""
            CoefficientDerivative.__init__(
                self, base_form, coefficients, arguments, coefficient_derivatives
            )
            BaseForm.__init__(self)

        def _analyze_form_arguments(self):
            """Collect the arguments of the corresponding BaseForm."""
            from ufl.algorithms.analysis import extract_coefficients, extract_type

            base_form, _, arguments, _ = self.ufl_operands

            def arg_type(x):
                if isinstance(x, BaseForm):
                    return Coargument
                return Argument

            # Each derivative arguments can either be a:
            # - `ufl.BaseForm`: if it contains a `ufl.Coargument`
            # - or a `ufl.Expr`: if it contains a `ufl.Argument`
            # When a `Coargument` is encountered, it is treated as an
            # argument (i.e. as V* -> V* and not V* x V -> R) and should
            # result in one single argument (in the dual space).
            base_form_args = base_form.arguments() + tuple(
                arg for a in arguments.ufl_operands for arg in extract_type(a, arg_type(a))
            )
            # BaseFormDerivative's arguments don't necessarily contain BaseArgument objects only
            # -> e.g. `derivative(u ** 2, u, u)` with `u` a Coefficient.
            base_form_coeffs = base_form.coefficients() + tuple(
                arg for a in arguments.ufl_operands for arg in extract_coefficients(a)
            )
            # Reconstruct arguments for correct numbering
            self._arguments = tuple(
                type(arg)(arg.ufl_function_space(), arg.number(), arg.part()) for arg in base_form_args
            )
            self._coefficients = base_form_coeffs

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class BaseFormCoordinateDerivative(BaseFormDerivative, CoordinateDerivative):
        """Derivative of a base form w.r.t. the SpatialCoordinates."""

        _ufl_noslots_ = True

        def __init__(self, base_form, coefficients, arguments, coefficient_derivatives):
            """Initalise."""
            BaseFormDerivative.__init__(
                self, base_form, coefficients, arguments, coefficient_derivatives
            )

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class BaseFormOperatorDerivative(BaseFormDerivative, BaseFormOperator):
        """Derivative of a base form operator w.r.t the degrees of freedom in a discrete Coefficient."""

        _ufl_noslots_ = True

        # BaseFormOperatorDerivative is only needed because of a different
        # differentiation procedure for BaseformOperator objects.
        def __init__(self, base_form, coefficients, arguments, coefficient_derivatives):
            """Initalise."""
            BaseFormDerivative.__init__(
                self, base_form, coefficients, arguments, coefficient_derivatives
            )
            self._argument_slots = base_form._argument_slots

        # Enforce Operator reconstruction as Operator is a parent class of
        # both: BaseFormDerivative and BaseFormOperator.
        # Therefore the latter overwrites Operator reconstruction and we would have:
        #   -> BaseFormOperatorDerivative._ufl_expr_reconstruct_ =
        #   BaseFormOperator._ufl_expr_reconstruct_
        def _ufl_expr_reconstruct_(
            self, *operands, function_space=None, derivatives=None, argument_slots=None
        ):
            if (
                (function_space is not None)
                or (derivatives is not None)
                or (argument_slots is not None)
            ):
                raise ValueError("_ufl_expr_reconstruct_ invoked with unused argument(s).")

            return Operator._ufl_expr_reconstruct_(self, *operands)

        # Set __repr__
        __repr__ = Operator.__repr__

        def argument_slots(self, outer_form=False):
            """Return a tuple of expressions containing argument and coefficient based expressions."""
            from ufl.algorithms.analysis import extract_arguments

            base_form, _, arguments, _ = self.ufl_operands
            argument_slots = base_form.argument_slots(outer_form) + tuple(
                arg for a in arguments for arg in extract_arguments(a)
            )
            return argument_slots

    @ufl_type(num_ops=4, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
    class BaseFormOperatorCoordinateDerivative(BaseFormOperatorDerivative, CoordinateDerivative):
        """Derivative of a base form operator w.r.t. the SpatialCoordinates."""

        _ufl_noslots_ = True

        def __init__(self, base_form, coefficients, arguments, coefficient_derivatives):
            """Initalise."""
            BaseFormOperatorDerivative.__init__(
                self, base_form, coefficients, arguments, coefficient_derivatives
            )

    @ufl_type(num_ops=2)
    class VariableDerivative(Derivative):
        """Variable Derivative."""

        __slots__ = (
            "ufl_free_indices",
            "ufl_index_dimensions",
            "ufl_shape",
        )

        def __new__(cls, f, v):
            """Create a new VariableDerivative."""
            # Checks
            if not isinstance(f, Expr):
                raise ValueError("Expecting an Expr in VariableDerivative.")
            if not isinstance(v, Variable | Coefficient):
                raise ValueError("Expecting a Variable in VariableDerivative.")
            if v.ufl_free_indices:
                raise ValueError("Differentiation variable cannot have free indices.")

            # Simplification
            # Return zero if expression is trivially independent of variable
            if f._ufl_is_terminal_ and f != v:
                return Zero(f.ufl_shape + v.ufl_shape, f.ufl_free_indices, f.ufl_index_dimensions)

            # Construction
            return Derivative.__new__(cls)

        def __init__(self, f, v):
            """Initalise."""
            Derivative.__init__(self, (f, v))
            self.ufl_free_indices = f.ufl_free_indices
            self.ufl_index_dimensions = f.ufl_index_dimensions
            self.ufl_shape = f.ufl_shape + v.ufl_shape

        def __str__(self):
            """Format as a string."""
            if isinstance(self.ufl_operands[0], Terminal):
                return f"d{self.ufl_operands[0]}/d[{self.ufl_operands[1]}]"
            return f"d/d[{self.ufl_operands[1]}] {parstr(self.ufl_operands[0], self)}"

    # --- Compound differentiation objects ---

    @ufl_type(is_abstract=True)
    class CompoundDerivative(Derivative):
        """Base class for all compound derivative types."""

        __slots__ = ()

        def __init__(self, operands):
            """Initalise."""
            Derivative.__init__(self, operands)

    @ufl_type(num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True)
    class Grad(CompoundDerivative):
        """Grad."""

        __slots__ = ("_dim",)

        def __new__(cls, f):
            """Create a new Grad."""
            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                dim = find_geometric_dimension(f)
                return Zero(f.ufl_shape + (dim,), f.ufl_free_indices, f.ufl_index_dimensions)
            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))
            self._dim = find_geometric_dimension(f)

        def _ufl_expr_reconstruct_(self, op):
            """Return a new object of the same type with new operands."""
            if is_cellwise_constant(op):
                if op.ufl_shape != self.ufl_operands[0].ufl_shape:
                    raise ValueError("Operand shape mismatch in Grad reconstruct.")
                if self.ufl_operands[0].ufl_free_indices != op.ufl_free_indices:
                    raise ValueError("Free index mismatch in Grad reconstruct.")
                return Zero(self.ufl_shape, self.ufl_free_indices, self.ufl_index_dimensions)
            return self._ufl_class_(op)

        def evaluate(self, x, mapping, component, index_values, derivatives=()):
            """Get child from mapping and return the component asked for."""
            component, i = component[:-1], component[-1]
            derivatives = derivatives + (i,)
            result = self.ufl_operands[0].evaluate(
                x, mapping, component, index_values, derivatives=derivatives
            )
            return result

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return self.ufl_operands[0].ufl_shape + (self._dim,)

        def __str__(self):
            """Format as a string."""
            return f"grad({self.ufl_operands[0]})"

    @ufl_type(
        num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True, is_in_reference_frame=True
    )
    class ReferenceGrad(CompoundDerivative):
        """Reference grad."""

        __slots__ = ("_dim",)

        def __new__(cls, f):
            """Create a new ReferenceGrad."""
            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                # TODO: Use max topological dimension if there are multiple topological dimensions.
                dim = extract_unique_domain(f, expand_mesh_sequence=False).topological_dimension()
                return Zero(f.ufl_shape + (dim,), f.ufl_free_indices, f.ufl_index_dimensions)
            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))
            # TODO: Use max topological dimension if there are multiple topological dimensions.
            self._dim = extract_unique_domain(f, expand_mesh_sequence=False).topological_dimension()

        def _ufl_expr_reconstruct_(self, op):
            """Return a new object of the same type with new operands."""
            if is_cellwise_constant(op):
                if op.ufl_shape != self.ufl_operands[0].ufl_shape:
                    raise ValueError("Operand shape mismatch in ReferenceGrad reconstruct.")
                if self.ufl_operands[0].ufl_free_indices != op.ufl_free_indices:
                    raise ValueError("Free index mismatch in ReferenceGrad reconstruct.")
                return Zero(self.ufl_shape, self.ufl_free_indices, self.ufl_index_dimensions)
            return self._ufl_class_(op)

        def evaluate(self, x, mapping, component, index_values, derivatives=()):
            """Get child from mapping and return the component asked for."""
            component, i = component[:-1], component[-1]
            derivatives = derivatives + (i,)
            result = self.ufl_operands[0].evaluate(
                x, mapping, component, index_values, derivatives=derivatives
            )
            return result

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return self.ufl_operands[0].ufl_shape + (self._dim,)

        def __str__(self):
            """Format as a string."""
            return f"reference_grad({self.ufl_operands[0]})"

    @ufl_type(num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True)
    class Div(CompoundDerivative):
        """Div."""

        __slots__ = ()

        def __new__(cls, f):
            """Create a new Div."""
            if f.ufl_free_indices:
                raise ValueError("Free indices in the divergence argument is not allowed.")

            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                return Zero(f.ufl_shape[:-1])  # No free indices asserted above

            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return self.ufl_operands[0].ufl_shape[:-1]

        def __str__(self):
            """Format as a string."""
            return f"div({self.ufl_operands[0]})"

    @ufl_type(
        num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True, is_in_reference_frame=True
    )
    class ReferenceDiv(CompoundDerivative):
        """Reference divergence."""

        __slots__ = ()

        def __new__(cls, f):
            """Create a new ReferenceDiv."""
            if f.ufl_free_indices:
                raise ValueError("Free indices in the divergence argument is not allowed.")

            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                return Zero(f.ufl_shape[:-1])  # No free indices asserted above

            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return self.ufl_operands[0].ufl_shape[:-1]

        def __str__(self):
            """Format as a string."""
            return f"reference_div({self.ufl_operands[0]})"

    @ufl_type(num_ops=1, inherit_indices_from_operand=0)
    class NablaGrad(CompoundDerivative):
        """Nabla grad."""

        __slots__ = ("_dim",)

        def __new__(cls, f):
            """Create a new NablaGrad."""
            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                dim = find_geometric_dimension(f)
                return Zero((dim,) + f.ufl_shape, f.ufl_free_indices, f.ufl_index_dimensions)
            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))
            self._dim = find_geometric_dimension(f)

        def _ufl_expr_reconstruct_(self, op):
            """Return a new object of the same type with new operands."""
            if is_cellwise_constant(op):
                if op.ufl_shape != self.ufl_operands[0].ufl_shape:
                    raise ValueError("Operand shape mismatch in NablaGrad reconstruct.")
                if self.ufl_operands[0].ufl_free_indices != op.ufl_free_indices:
                    raise ValueError("Free index mismatch in NablaGrad reconstruct.")
                return Zero(self.ufl_shape, self.ufl_free_indices, self.ufl_index_dimensions)
            return self._ufl_class_(op)

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return (self._dim,) + self.ufl_operands[0].ufl_shape

        def __str__(self):
            """Format as a string."""
            return f"nabla_grad({self.ufl_operands[0]})"

    @ufl_type(num_ops=1, inherit_indices_from_operand=0)
    class NablaDiv(CompoundDerivative):
        """Nabla div."""

        __slots__ = ()

        def __new__(cls, f):
            """Create a new NablaDiv."""
            if f.ufl_free_indices:
                raise ValueError("Free indices in the divergence argument is not allowed.")

            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                return Zero(f.ufl_shape[1:])  # No free indices asserted above

            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))

        @property
        def ufl_shape(self):
            """Get the UFL shape."""
            return self.ufl_operands[0].ufl_shape[1:]

        def __str__(self):
            """Format as a string."""
            return f"nabla_div({self.ufl_operands[0]})"

    _curl_shapes = {(): (2,), (2,): (), (3,): (3,)}

    @ufl_type(num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True)
    class Curl(CompoundDerivative):
        """Compound derivative."""

        __slots__ = ("ufl_shape",)

        def __new__(cls, f):
            """Create a new CompoundDerivative."""
            # Validate input
            sh = f.ufl_shape
            if f.ufl_shape not in ((), (2,), (3,)):
                raise ValueError("Expecting a scalar, 2D vector or 3D vector.")
            if f.ufl_free_indices:
                raise ValueError("Free indices in the curl argument is not allowed.")

            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                sh = {(): (2,), (2,): (), (3,): (3,)}[sh]
                return Zero(sh)  # No free indices asserted above
            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))
            self.ufl_shape = _curl_shapes[f.ufl_shape]

        def __str__(self):
            """Format as a string."""
            return f"curl({self.ufl_operands[0]})"

    @ufl_type(
        num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True, is_in_reference_frame=True
    )
    class ReferenceCurl(CompoundDerivative):
        """Reference curl."""

        __slots__ = ("ufl_shape",)

        def __new__(cls, f):
            """Create a new ReferenceCurl."""
            # Validate input
            sh = f.ufl_shape
            if f.ufl_shape not in ((), (2,), (3,)):
                raise ValueError("Expecting a scalar, 2D vector or 3D vector.")
            if f.ufl_free_indices:
                raise ValueError("Free indices in the curl argument is not allowed.")

            # Return zero if expression is trivially constant
            if is_cellwise_constant(f):
                sh = {(): (2,), (2,): (), (3,): (3,)}[sh]
                return Zero(sh)  # No free indices asserted above
            return CompoundDerivative.__new__(cls)

        def __init__(self, f):
            """Initalise."""
            CompoundDerivative.__init__(self, (f,))
            self.ufl_shape = _curl_shapes[f.ufl_shape]

        def __str__(self):
            """Format as a string."""
            return f"reference_curl({self.ufl_operands[0]})"