# Source code for ufl.algorithms.remove_component_tensors

    """Remove component tensors.

    This module contains classes and functions to remove component tensors.
    """
    # Copyright (C) 2025 Pablo Brubeck
    #
    # This file is part of UFL (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later

    from collections import defaultdict

    from ufl.algorithms.map_integrands import map_integrand_dags
    from ufl.classes import ComponentTensor, Index, MultiIndex, Zero
    from ufl.corealg.map_dag import map_expr_dag
    from ufl.corealg.multifunction import MultiFunction
    from ufl.index_combination_utils import unique_sorted_indices

    class IndexReplacer(MultiFunction):
        """Replace Indices."""

        def __init__(self, fimap: dict):
            """Initialise.

            Args:
               fimap: map for index replacements.

            """
            MultiFunction.__init__(self)
            self.fimap = fimap

        expr = MultiFunction.reuse_if_untouched

        def zero(self, o):
            """Handle Zero."""
            indices = tuple(map(Index, o.ufl_free_indices))
            if not any(i in self.fimap for i in indices):
                # Reuse if untouched
                return o

            fi = []
            for i, d in zip(indices, o.ufl_index_dimensions):
                j = self.fimap.get(i, i)
                if isinstance(j, Index):
                    fi.append((j.count(), d))

            fi = unique_sorted_indices(sorted(fi))
            free_indices, index_dimensions = zip(*fi)

            return Zero(
                shape=o.ufl_shape,
                free_indices=free_indices,
                index_dimensions=index_dimensions,
            )

        def multi_index(self, o):
            """Handle MultiIndex."""
            if not any(i in self.fimap for i in o):
                # Reuse if untouched
                return o

            indices = tuple(self.fimap.get(i, i) for i in o)
            return MultiIndex(indices)

    class IndexRemover(MultiFunction):
        """Remove Indexed."""

        def __init__(self):
            """Initialise."""
            MultiFunction.__init__(self)
            self.rules = {}
            # caches for reuse in the dispatched transformers
            self.vcaches = defaultdict(dict)
            self.rcaches = defaultdict(dict)

        expr = MultiFunction.reuse_if_untouched

        def indexed(self, o, o1, i1):
            """Simplify Indexed."""
            if isinstance(o1, ComponentTensor):
                # Simplify Indexed ComponentTensor
                o2, i2 = o1.ufl_operands
                # Replace outer indices
                rkey = (i2, i1)
                rule = self.rules.get(rkey)
                if rule is None:
                    # NOTE: Replace with `fimap = dict(zip(i2, i1, strict=True))` when
                    # Python>=3.10
                    assert len(i2) == len(i1)
                    fimap = dict(zip(i2, i1))
                    rule = IndexReplacer(fimap)
                    self.rules[rkey] = rule

                key = (IndexReplacer, *rkey)
                return map_expr_dag(rule, o2, vcache=self.vcaches[key], rcache=self.rcaches[key])

            elif o.ufl_operands[0] is o1:
                # Reuse if untouched
                return o
            else:
                return o._ufl_expr_reconstruct_(o1, i1)

    def remove_component_tensors(o):
        """Remove component tensors."""
        rule = IndexRemover()
        return map_integrand_dags(rule, o)