# Source code for ufl.split_functions

    """Algorithm for splitting a Coefficient or Argument into subfunctions."""
    # 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 Anders Logg, 2008
    # Modified by Pablo Brubeck, 2025

    import numpy as np

    from ufl.domain import extract_unique_domain
    from ufl.functionspace import FunctionSpace
    from ufl.indexed import Indexed
    from ufl.permutation import compute_indices
    from ufl.pullback import SymmetricPullback
    from ufl.tensors import ListTensor, as_tensor
    from ufl.utils.indexflattening import flatten_multiindex, shape_to_strides
    from ufl.utils.sequences import product

    def split(v):
        """Split a coefficient or argument.

        If v is a Coefficient or Argument in a mixed space, returns a tuple
        with the function components corresponding to the subelements.
        """
        # Default range is all of v
        begin = 0
        end = None

        if isinstance(v, Indexed):
            # Special case: split previous output of split again
            # Consistent with simple element, just return function in a tuple
            return (v,)

        elif isinstance(v, ListTensor):
            # Special case: split previous output of split again
            ops = tuple(v[i] for i in np.ndindex(v.ufl_shape))

            if all(isinstance(comp, Indexed) for comp in ops):
                args = [comp.ufl_operands[0] for comp in ops]
                if all(args[0] == args[i] for i in range(1, len(args))):
                    # Get innermost terminal here and its element
                    v = args[0]
                    # Get relevant range of v components
                    (begin,) = ops[0].ufl_operands[1]
                    (end,) = ops[-1].ufl_operands[1]
                    begin = int(begin)
                    end = int(end) + 1
                else:
                    raise ValueError(f"Don't know how to split {v}.")
            else:
                raise ValueError(f"Don't know how to split {v}.")

        # Special case: simple element, just return function in a tuple
        element = v.ufl_element()
        if element.num_sub_elements == 0:
            assert end is None
            return (v,)

        if len(v.ufl_shape) != 1:
            raise ValueError(
                "Don't know how to split tensor valued mixed functions without flattened index space."
            )

        domain = extract_unique_domain(v, expand_mesh_sequence=False)

        # Compute value size and set default range end
        value_size = v.ufl_function_space().value_size
        if end is None:
            end = value_size
        else:
            # Recursively dive into mixedelement in to subelement
            # corresponding to beginning of range
            j = begin
            while True:
                domains = domain.iterable_like(element)
                for d, e in zip(domains, element.sub_elements):
                    if j < FunctionSpace(d, e).value_size:
                        domain = d
                        element = e
                        break
                    j -= FunctionSpace(d, e).value_size
                # Then break when we find the subelement that covers the whole range
                if FunctionSpace(domain, element).value_size == (end - begin):
                    break

        # Build expressions representing the subfunction of v for each subelement
        offset = begin
        sub_functions = []

        # Deal with symmetry
        if isinstance(element.pullback, SymmetricPullback):
            symmetry = element.pullback._symmetry
            indices = (comp for idx, comp in sorted(symmetry.items()))
        else:
            indices = range(len(element.sub_elements))

        domains = tuple(domain.iterable_like(element))
        for i in indices:
            d = domains[i]
            e = element.sub_elements[i]
            # Get shape, size, indices, and v components
            # corresponding to subelement value
            shape = FunctionSpace(d, e).value_shape
            strides = shape_to_strides(shape)
            rank = len(shape)
            sub_size = product(shape)
            subindices = [flatten_multiindex(c, strides) for c in compute_indices(shape)]
            components = [v[k + offset] for k in subindices]

            # Shape components into same shape as subelement
            if rank == 0:
                (subv,) = components
            else:
                subv = as_tensor(np.reshape(components, shape))

            offset += sub_size
            sub_functions.append(subv)

        if end != offset:
            raise ValueError(
                "Function splitting failed to extract components for whole intended range."
            )

        return tuple(sub_functions)