def parse_tce_out(tce_out, range_cb, base_cb):
    """Parse a TCE output into a list of terms.

    A list of terms, and a dictionary of free symbols will be returned.
    """

    lines = []
    for line in tce_out.splitlines():
        stripped = line.strip()
        if len(stripped) > 0:
            lines.append(stripped)
        continue

    free_vars = collections.defaultdict(set)
    return list(itertools.chain.from_iterable(
        _parse_tce_line(line, range_cb, base_cb, free_vars)
        for line in lines
    )), free_vars NEG = 1 IDENT = 0 class Tensor:
    """The main tensor class.

    A tensor is an aggregate of terms distributed and managed by Spark.  Here
    most operations needed for tensors are defined.

    Normally, tensor instances are created from drudge methods or tensor
    operations.  Direct invocation of its constructor is seldom in user scripts.

    """

    __slots__ = [
        '_drudge',
        '_terms',
        '_local_terms',
        '_free_vars',
        '_expanded',
        '_repartitioned'
    ]

    #
    # Term creation
    #

    def __init__(self, drudge: 'Drudge', terms: RDD,
                 free_vars: typing.Set[Symbol] = None,
                 expanded=False, repartitioned=False):
        """Initialize the tensor.

        This function is not designed to be called by users directly.  Tensor
        creation should be carried out by factory function inside drudges and
        the operations defined here.

        The default values for the keyword arguments are always the safest
        choice, for better performance, manipulations are encouraged to have
        proper consideration of all the keyword arguments.
        """

        self._drudge = drudge
        self._terms = terms

        self._local_terms = None  # type: typing.List[Term]

        self._free_vars = free_vars
        self._expanded = expanded
        self._repartitioned = repartitioned

    # To be used by the apply method.
    _INIT_ARGS = {
        'free_vars': '_free_vars',
        'expanded': '_expanded',
        'repartitioned': '_repartitioned'
    }

    #
    # Basic information
    #

    @property
    def drudge(self):
        """The drudge created the tensor."""
        return self._drudge

    @property
    def terms(self):
        """The terms in the tensor, as an RDD object.

        Although for users, normally there is no need for direct manipulation of
        the terms, it is still exposed here for flexibility.
        """
        return self._terms

    @property
    def local_terms(self) -> typing.List[Term]:
        """Gather the terms locally into a list.

        The list returned by this is for read-only and should **never** be
        mutated.

        .. warning::

            This method will gather all terms into the memory of the driver.

        """
        if self._local_terms is None:
            self._local_terms = self._terms.collect()
        else:
            pass

        return self._local_terms

    @property
    def n_terms(self) -> int:
        """Get the number of terms.

        A zero number of terms signatures a zero tensor.  Accessing this
        property will make the tensor to be cached automatically.
        """
        if self._local_terms is not None:
            return len(self._local_terms)
        else:
            self.cache()  # We never get a tensor just to count its terms.
            return self._terms.count()

    def cache(self):
        """Cache the terms in the tensor.

        This method should be called when this tensor is an intermediate result
        that will be used multiple times.  The tensor itself will be returned
        for the ease of chaining.
        """

        self._terms.cache()
        return self

    def repartition(self, num_partitions=None, cache=False):
        """Repartition the terms across the Spark cluster.

        This function should be called when the terms need to be rebalanced
        among the workers.  Note that this incurs an Spark RDD shuffle operation
        and might be very expensive.  Its invocation and the number of
        partitions used need to be fine-tuned for different problems to achieve
        good performance.

        Parameters
        ----------

        num_partitions : int
            The number of partitions.  By default, the number is read from the
            drudge object.

        cache : bool
            If the result is going to be cached.

        """

        if not self._repartitioned:

            num_partitions = (
                self._drudge.num_partitions if num_partitions is None else
                num_partitions
            )
            if num_partitions is None:
                raise ValueError('No default number of partitions available')

            self._terms = self._terms.repartition(num_partitions)
            self._repartitioned = True

        if cache:
            self.cache()
        return self

    @property
    def is_scalar(self):
        """If the tensor is a scalar.

        A tensor is considered a scalar when none of its terms has a vector
        part.  This property will make the tensor automatically cached.
        """

        self.cache()

        # Work around a pyspark bug by doing the reduction locally.
        return all(
            self._terms.map(lambda x: x.is_scalar).collect()
        )

    @property
    def free_vars(self) -> typing.Set[Symbol]:
        """The free variables in the tensor.
        """
        if self._free_vars is None:
            self._free_vars = self._get_free_vars(self._terms)

        return self._free_vars

    @staticmethod
    def _get_free_vars(terms) -> typing.Set[Symbol]:
        """Get the free variables in the given terms."""

        # The terms are definitely going to be used for other purposes.
        terms.cache()

        return terms.map(
            lambda term: term.free_vars
        ).aggregate(set(), _union, _union)
        # TODO: investigate performance characteristic with treeAggregate.

    @property
    def expanded(self):
        """If the tensor is already expanded."""
        return self._expanded

    @property
    def repartitioned(self):
        """If the terms in the tensor is already repartitioned."""
        return self._repartitioned

    def has_base(self, base: typing.Union[IndexedBase, Symbol, Vec]) -> bool:
        """Find if the tensor has the given scalar or vector base.

        Parameters
        ----------

        base
            The base whose presence is to be queried.  When it is indexed base
            or a plain symbol, its presence in the amplitude part is tested.
            When it is a vector, its presence in the vector part is tested.

        """

        # A tensor is barely needed by a has_base decision only.
        self.cache()

        # Work around a possible pyspark bug in reduce.
        return any(
            self._terms.map(functools.partial(Term.has_base, base=base))
                .collect()
        )

    #
    # Printing support
    #

    def __repr__(self):
        """Get the machine string representation.

        In normal execution environment, only the memory address is displayed,
        since the tensor may or may not be evaluated yet.  In drudge scripts,
        the readable string representation is returned.
        """

        if self._drudge.inside_drs:
            return str(self)
        else:
            return super().__repr__()

    def __str__(self):
        """Get the string representation of the tensor.

        Note that this function will **gather** all terms into the driver.

        """
        if self.n_terms == 0:
            return '0'
        else:
            return '\n + '.join(str(i) for i in self.local_terms)

    def latex(self, **kwargs):
        r"""Get the latex form for the tensor.

        The actual printing is dispatched to the drudge object for the
        convenience of tuning the appearance.

        All keyword arguments are forwarded to the
        :py:meth:`Drudge.format_latex` method.

        """
        return self._drudge.format_latex(self, **kwargs)

    def display(self, if_return=True, **kwargs):
        """Display the tensor in interactive IPython notebook sessions.

        Parameters
        ----------

        if_return
            If the resulted equation be returned rather than directly displayed.
            It can be disabled for displaying equation in the middle of a
            Jupyter cell.

        kwargs
            All the rest of the keyword arguments are forwarded to the
            :py:meth:`Drudge.format_latex` method.

        """

        form = Math(self.latex(**kwargs))
        if if_return:
            return form
        else:
            display(form)
            return

    #
    # Pickling support
    #

    def __getstate__(self):
        """Get the current state of the tensor.

        Here we just have the local terms.  Other cached information are
        discarded.
        """
        return self.local_terms

    def __setstate__(self, state):
        """Set the state for the new tensor.

        This function reads the drudge to use from the module attribute, which
        is set in the :py:meth:`Drudge.pickle_env` method.
        """

        global current_drudge
        drudge = current_drudge
        if drudge is None:
            raise ValueError(
                'Tensor objects cannot be unpickled, '
                'need to be inside Drudge.pickle_env'
            )

        assert isinstance(drudge, Drudge)

        # Subclasses might be of __init__ function of a different signature.
        Tensor.__init__(self, drudge, drudge.ctx.parallelize(state))
        return

    #
    # Small manipulations
    #

    def apply(self, func, **kwargs):
        """Apply the given function to the RDD of terms.

        This function is analogous to the replace function of Python named
        tuples, the same value from self for the tensor initializer is going to
        be used when it is not given.  The terms get special treatment since it
        is the centre of tensor objects.  The drudge is kept the same always.

        Users generally do not need this method.  It is exposed here just for
        flexibility and convenience.

        .. warning::

            For developers:  Note that the resulted tensor will inherit all
            unspecified keyword arguments from self.  This method can give
            *unexpected results* if certain arguments are not correctly reset
            when they need to.  For instance, when expanded is not reset when
            the result is no longer guaranteed to be in expanded form, later
            expansions could be skipped when they actually need to be performed.

            So all functions using this methods need to be reviewed when new
            property are added to tensor class.  Direct invocation of the tensor
            constructor is a much safe alternative.

        """

        for k, v in self._INIT_ARGS.items():
            if k not in kwargs:
                kwargs[k] = getattr(self, v)

        return Tensor(self._drudge, func(self._terms), **kwargs)

    #
    # Here for a lot of methods, we have two versions, with one being public,
    # another being private with a leading underscore.  The private version
    # operates on given RDD of terms and returns another RDD of terms.  The
    # public version operates on the terms of the current tensor, and return
    # another tensor.
    #

    def reset_dumms(self, excl=None):
        """Reset the dummies.

        The dummies will be set to the canonical dummies according to the order
        in the summation list.  This method is especially useful on
        canonicalized tensors.

        Parameters
        ----------

        excl
            A set of symbols to be excluded in the dummy selection.  This option
            can be useful when some symbols already used as dummies are planned
            to be used for other purposes.

        """

        free_vars = self.free_vars
        if excl is None:
            excl = free_vars
        else:
            excl |= free_vars

        return self.apply(functools.partial(self._reset_dumms, excl=excl))

    def _reset_dumms(self, terms: RDD, excl) -> RDD:
        """Get terms with dummies reset.

        Note that this function does not automatically add the free variables in
        the terms to the excluded symbols.
        """

        dumms = self._drudge.dumms
        res_terms = terms.map(
            lambda term: term.reset_dumms(dumms=dumms.value, excl=excl)[0]
        )
        return res_terms

    def simplify_amps(self):
        """Simplify the amplitudes in the tensor.

        This method simplifies the amplitude in the terms of the tensor by using
        the facility from SymPy.  The zero terms will be filtered out as well.

        """

        # Some free variables might be canceled.
        return self.apply(
            self._simplify_amps, free_vars=None, repartitioned=False
        )

    @staticmethod
    def _simplify_amps(terms):
        """Get the terms with amplitude simplified by SymPy."""

        simplified_terms = terms.map(
            lambda term: term.map(lambda x: x.simplify(), skip_vecs=True)
        ).filter(_is_nonzero)

        return simplified_terms

    def simplify_deltas(self):
        """Simplify the deltas in the tensor.

        Kronecker deltas whose operands contains dummies will be attempted to be
        simplified.
        """

        return Tensor(
            self._drudge,
            self._simplify_deltas(self._terms, expanded=self._expanded),
            expanded=True,
        )

    def _simplify_deltas(self, terms, expanded):
        """Simplify the deltas in the terms.

        This function will leave the resulted terms in expanded form.
        """

        if not expanded:
            terms = self._expand(terms)

        resolvers = self._drudge.resolvers

        return terms.map(
            lambda x: x.simplify_deltas(resolvers.value)
        ).filter(_is_nonzero)

    def simplify_sums(self, simplifiers=True, excl_bases=True):
        """Simplify the summations within the amplitude in the tensor.

        This method attempts to perform simplification for concrete summations
        over dummies that does not appear on the vector part at all.

        First, bounded summations with dummies not actually involved in the term
        will be replaced by a multiplication with its size.  Then simplification
        is going to be performed on other summations within the amplitude part.

        Parameters
        ----------

        simplifiers

            The rules to simplify the internally summed factors of the
            amplitude.  It should be a Spark broadcast of a mapping from the
            number of summations indices able to be handled to an iterable of
            functions implementing the rules.

            The functions will be called with SymPy ``Sum`` objects as the first
            positional argument.  Then keyword arguments ``resolvers`` will be
            the resolvers for the ranges, and ``sums_dict`` will hold the
            original summation dictionary for the term.

            If the same expression or ``None`` is returned, it will be
            considered to be not simplifiable any more.  When it is
            simplifiable, the simplified SymPy expression should be returned.
            Note that, currently, the simplified form cannot be a Sum by itself.

            By default, it is set to True, which will use the
            :py:attr:`sum_simplifiers` attribute of the current drudge object.
            It can also be set to anything that evaluates to false to disable
            deep simplification of the summations.

            By default, the ``sum_simplifiers`` attribute holds a dictionary
            having only the SymPy ``eval_sum_symbolic`` function for summations
            with a single summation index.

        excl_bases
            If summations involved by indexed bases are also going to be
            excluded from attempts of simplifications.  This can be safely set
            to true unless customized rules exists.

        """

        return Tensor(self._drudge, self._simplify_sums(
            self._terms, simplifiers=simplifiers, excl_bases=excl_bases
        ))

    def _simplify_sums(
            self, terms: RDD, simplifiers=True, excl_bases=True
    ):
        """Simplify the summations in the given terms."""

        if simplifiers is True:
            simplifiers = self._drudge.sum_simplifiers.bcast

        # Make it a two-step process for future extensibility.
        terms = terms.map(lambda x: x.simplify_trivial_sums())

        if simplifiers:
            terms = terms.map(functools.partial(
                simplify_amp_sums_term,
                simplifiers=simplifiers, excl_bases=excl_bases,
                resolvers=self._drudge.resolvers
            ))

        return terms

    def expand(self):
        """Expand the terms in the tensor.

        By calling this method, terms in the tensor whose amplitude is the
        addition of multiple parts will be expanded into multiple terms.
        """
        if self._expanded:
            return self
        else:
            return self.apply(self._expand, expanded=True, repartitioned=False)

    @staticmethod
    def _expand(terms):
        """Get terms after they are fully expanded."""
        return terms.flatMap(lambda term: term.expand())

    def shallow_expand(self):
        """Expand terms with addition amplitudes into multiple terms.

        Different from the :py:meth:`expand` method, this method only expands
        a term when its amplitude is written as an addition on the
        top-level.  This is useful for cases where the inner structures of the
        amplitude expression tree is intended to be preserved.
        """
        return self.bind(self._shallow_expand)

    @staticmethod
    def _shallow_expand(term: Term) -> typing.Iterable[Term]:
        """Expand a term shallowly."""
        if isinstance(term.amp, Add):
            return [
                Term(sums=term.sums, amp=i, vecs=term.vecs)
                for i in term.amp.args
            ]
        else:
            return [term]

    def sort(self):
        """Sort the terms in the tensor.

        The terms will generally be sorted according to increasing complexity.
        """
        return self.apply(self._sort)

    @staticmethod
    def _sort(terms: RDD):
        """Sort the terms in the tensor."""
        return terms.sortBy(lambda term: term.sort_key)

    def merge(self, consts=None, gens=None):
        """Merge some terms.

        This function attempts to merge some terms into smaller number of terms.
        Two terms can be merged only when their summation list, vector part, and
        selected factors are *syntactically* the same.  So it is more useful
        when the canonicalization has been performed and the dummies reset.

        When simple merge is disabled, nothing is put into the selected factors,
        all terms with the same summations and vectors are tried to be merged.
        When simple merge is enabled, inside the amplitude, any factor
        containing one of the dummy variables and all indexed quantities are put
        into the selected factors.  By default, all other factors are considered
        to be symbols standing for some constants and will be merged.  If the
        ``consts`` argument is given, only the symbols given there will be
        considered to be constant symbols for coefficients, with all factors
        with any other symbol put into the selected factors.  If the ``gens``
        argument is given, only symbols given here are considered to be
        generators of the polynomial, with all others considered to be
        constants.

        """

        # All the traits could be invalidated by merging.
        return Tensor(
            self._drudge, self._merge(self._terms, consts, gens)
        )

    def _merge(self, terms, consts, gens):
        """Get the term when they are attempted to be merged."""
        if not self._drudge.simple_merge and consts is None and gens is None:
            specials = None
        else:
            specials = _DecomposeSpecials(consts, gens)

        return terms.map(
            functools.partial(_decompose_term, specials=specials)
        ).reduceByKey(operator.add).map(_recover_term)

    #
    # Canonicalization
    #

    def canon(self):
        """Canonicalize the terms in the tensor.

        This method will first expand the terms in the tensor.  Then the
        canonicalization algorithm is going to be applied to each of the terms.
        Note that this method does not rename the dummies.
        """
        return self.apply(
            functools.partial(self._canon, expanded=self._expanded),
            expanded=True, repartitioned=self._expanded and self._repartitioned
        )

    def _canon(self, terms, expanded):
        """Compute the canonicalized terms."""

        symms = self._drudge.symms

        vec_colour = self._drudge.vec_colour
        # Normally a static function, not broadcast variable.

        if not expanded:
            expanded_terms = self._expand(terms)
        else:
            expanded_terms = terms
        canoned = expanded_terms.map(
            lambda term: term.canon(symms=symms.value, vec_colour=vec_colour)
        )
        return canoned

    def normal_order(self):
        """Normal order the terms in the tensor.

        The actual work is dispatched to the drudge, who has domain specific
        knowledge about the noncommutativity of the vectors.

        """

        # Free variables, expanded, and repartitioned can all be invalidated.
        return Tensor(
            self._drudge, self._drudge.normal_order(self.terms)
        )

    #
    # The driver simplification.
    #

    def simplify(self):
        """Simplify the tensor.

        This is the master driver function for tensor simplification.  Inside
        drudge scripts, it also make eager evaluation and repartition the terms
        among the Spark workers, with the result cached.  This is for the ease
        of users unfamiliar with the Spark lazy execution model.

        """

        result = Tensor(
            self._drudge, self._simplify(self._terms), expanded=True
        )

        if self._drudge.inside_drs:
            result.repartition(cache=True)
            _ = result.n_terms

        return result

    def _simplify(self, terms):
        """Get the terms in the simplified form."""

        num_partitions = self._drudge.num_partitions

        repartitioned = self._repartitioned
        if not self._expanded:
            terms = self._expand(terms)
            repartitioned = False

        if not repartitioned and num_partitions is not None:
            terms = terms.repartition(num_partitions)

        # First we make the vector part normal-ordered.
        terms = self._drudge.normal_order(terms)
        if num_partitions is not None:
            terms = terms.repartition(num_partitions)

        # Simplify the amplitude part.
        terms = self._simplify_amps(terms)
        terms = self._simplify_sums(terms)
        terms = self._simplify_amps(terms)
        terms = self._simplify_deltas(terms, False)

        # Canonicalize the terms and see if they can be merged.
        terms = self._canon(terms, True)
        # In rare cases, normal order could make the result unexpanded.
        #
        # TODO: Find a design to skip repartition in most cases.

        free_vars = self._get_free_vars(terms)
        terms = self._reset_dumms(terms, excl=free_vars)
        terms = self._merge(terms, None, None)

        # Finally simplify the merged amplitude again.
        if self._drudge.full_simplify:
            terms = self._simplify_amps(terms)

        # Make the final expansion.
        terms = self._expand(terms)

        return terms

    #
    # Comparison operations
    #

    def __eq__(self, other):
        """Compare the equality of tensors.

        Note that this function only compares the syntactical equality of
        tensors.  Mathematically equal tensors might be compared to be unequal
        by this function when they are not simplified.

        Note that only comparison with zero is performed by counting the number
        of terms distributed.  Or this function gathers all terms in both
        tensors and can be very expensive.  So direct comparison of two tensors
        is mostly suitable for testing and debugging on small problems only.
        For large scale problems, it is advised to compare the simplified
        difference with zero.
        """

        n_terms = self.n_terms
        if isinstance(other, Tensor):
            return set(self.local_terms) == set(other.local_terms)
        elif other == 0:
            return n_terms == 0
        else:
            return self._drudge.sum(other)

    #
    # Mathematical operations
    #

    _op_priority = 20.0

    def __add__(self, other):
        """Add the two tensors together.

        The terms in the two tensors will be concatenated together, without any
        further processing.

        In addition to full tensors, tensor inputs can also be directly added.
        """
        return self._add(other)

    def __radd__(self, other):
        """Add tensor with something in front."""
        return self._add(other)

    def _add(self, other):
        """Add tensor with another thing."""
        if not isinstance(other, Tensor):
            other = self._drudge.sum(other)

        if self._free_vars is not None and other._free_vars is not None:
            free_vars = self.free_vars | other.free_vars
        else:
            free_vars = None

        return Tensor(
            self._drudge, self._terms.union(other.terms),
            free_vars=free_vars,
            expanded=self._expanded and other.expanded
        )

    def __sub__(self, other):
        """Subtract another tensor from this tensor.
        """
        return self._add(-other)

    def __rsub__(self, other):
        """Subtract the tensor from another quantity."""
        return (-self)._add(other)

    def __neg__(self):
        """Negate the current tensor.

        The result will be equivalent to multiplication with :math:`-1`.
        """
        return self.apply(
            lambda terms: terms.map(lambda x: x.scale(-1))
        )

    def __mul__(self, other) -> 'Tensor':
        """Multiply the tensor.

        This multiplication operation is done completely within the framework of
        free algebras.  The vectors are only concatenated without further
        processing.  The actual handling of the commutativity should be carried
        out at the normal ordering operation for different problems.

        In addition to full tensors, tensors can also be multiplied to user
        tensor input directly.
        """
        return self._mul(other)

    def __rmul__(self, other):
        """Multiply the tensor on the right.
        """
        return self._mul(other, right=True)

    def _mul(self, other, right=False):
        """Multiply the tensor with another."""
        prod, free_vars, expanded = self._cartesian_terms(other, right)

        dumms = self._drudge.dumms
        return Tensor(self._drudge, prod.map(
            lambda x: x[0].mul_term(x[1], dumms=dumms.value, excl=free_vars)
        ), free_vars=free_vars, expanded=expanded)

    def __or__(self, other):
        """Compute the commutator with another tensor.

        In the same way as multiplication, this can be used for both full
        tensors and local tensor input.
        """
        return self._comm(other)

    def __ror__(self, other):
        """Compute the commutator with another tensor on the right."""
        return self._comm(other, right=True)

    def _comm(self, other, right=False):
        """Compute the commutator."""

        if isinstance(other, Expr):
            msg = 'Taking commutator with commutative expression `{}`'.format(
                other
            )
            warnings.warn(msg)

        prod, free_vars, expanded = self._cartesian_terms(other, right)

        dumms = self._drudge.dumms
        return Tensor(self._drudge, prod.flatMap(
            lambda x: x[0].comm_term(x[1], dumms=dumms.value, excl=free_vars)
        ), free_vars=free_vars, expanded=expanded)

    def _cartesian_terms(self, other, right):
        """Cartesian the terms with the terms in another tensor.

        The other tensor will be attempted to be interpreted as a tensor when it
        is not given as one.  And the free variables used in both tensors will
        also be returned since it is going to be used frequently.
        """

        if isinstance(other, Tensor):

            if right:
                prod = other.terms.cartesian(self._terms)
            else:
                prod = self._terms.cartesian(other.terms)

            free_vars = self.free_vars | other.free_vars
            expanded = self._expanded and other._expanded

        else:
            # Special optimized version when the other terms are local.

            other_terms = parse_terms(other)
            if len(other_terms) > 1:
                prod = self._terms.flatMap(lambda term: [
                    (i, term) if right else (term, i)
                    for i in other_terms
                ])
            else:
                # Special optimization when we just have one term.
                other_term = other_terms[0]
                prod = self._terms.map(
                    lambda term:
                    (other_term, term) if right else (term, other_term)
                )

            free_vars = set.union(*[
                i.free_vars for i in other_terms
            ])
            free_vars |= self.free_vars
            expanded = False

        return prod, free_vars, expanded

    def __truediv__(self, other):
        """Divide tensor by a scalar quantity."""

        other = sympify(other)
        return self.apply(
            lambda terms: terms.map(lambda x: x.scale(1 / other)),
            free_vars=None
        )

    def __rtruediv__(self, other):
        """Make division over a tensor."""
        raise NotImplementedError('General tensors cannot be divided over.')

    #
    # Substitution
    #

    @staticmethod
    def _decr_vecs(vecs: typing.Tuple[Vec]):
        return tuple(
            Vec(label=(vec.label, _Decred(0)), indices=vec.indices)
            for vec in vecs
        )

    @staticmethod
    def _restore(term: Term, decr_vars=None):

        if decr_vars != None:
            restore_vars = {decr_var: var for var, decr_var in decr_vars.items()}
        else:
            suffix_index = -len(_DECR_SUFFIX)
            restore_vars = {}
            for var in term.free_vars:
                if isinstance(var, Symbol):
                    name = var.name
                    Ref = Symbol
                elif isinstance(var, IndexedBase):
                    name = var.label.name
                    Ref = IndexedBase
                else:
                    raise TypeError('Expecting Symbol or IndexedBase')

                if name.endswith(_DECR_SUFFIX):
                    restore_vars[var] = Ref(
                        name[:suffix_index], **var._assumptions
                    )

        func = lambda x: x.xreplace(restore_vars)

        new_vecs = tuple(
            Vec(label=vec.label[0], indices=vec.indices)
            if isinstance(vec.label, tuple) and len(vec.label) == 2
                and isinstance(vec.label[1], _Decred)
            else vec
            for vec in term.vecs
        )

        return term.map(func, vecs=new_vecs, skip_vecs=True)

    def subst(
            self, lhs, rhs, wilds=None, full_balance=False, excl=None,
            simult=True, keep_decorated=False
    ):
        """Substitute the all appearance of the defined tensor.

        When the given LHS is a plain SymPy symbol, all its appearances in the
        amplitude of the tensor will be replaced.  Or the LHS can also be
        indexed SymPy expression or indexed Vector, for which all of the
        appearances of the indexed base or vector base will be attempted to be
        matched against the indices on the LHS.  When a matching succeeds for
        all the indices, the RHS, with the substitution found in the matching
        performed, will be replace the indexed base in the amplitude, or the
        vector.  Note that for scalar LHS, the RHS must contain no vector.  For
        vector LHS, also supported is a product of vectors.  This will lead to
        pattern matching inside the vector part.

        A special case is when the LHS is a plain integral unity.  In this case,
        the scalar terms will be replaced by product with the RHS, as if the
        identity element has been replaced by the RHS.

        Parameters
        ----------

        lhs
            The left-hand side of the substitution.

        rhs
            The right-hand side of the substituion.

        wilds
            Since we do not commonly define tensors with wild symbols, an option
            ``wilds`` can be used to give a mapping translating plain symbols on
            the LHS and the RHS to the wild symbols that would like to be used.
            The default value of None could make all **plain** symbols in the
            indices of the LHS to be translated into a wild symbol with the same
            name and no exclusion. And empty dictionary can be used to disable
            all such automatic translation.  The default value of None should
            satisfy most needs.

        full_balance
            If load-balancing is going to be performed after each substitution
            step.  This is normally not needed.

        excl
            Symbols that needs to be excluded to be used as dummies.

        simult
            The core substitution algorithm works by making one substitution at
            a time until no more left-hand side can be matched successfully.  In
            this way, if the pattern in the LHS reappears on the RHS, infinite
            loop will be resulted.  When it is set to true, the
            newly-substituted part will no longer participate in the matching.

            By default, it is true.  Internally, it decorates the free variables
            on the RHS to special forms, so that the originally-present ones are
            different from the newly-substituted ones.  Note that its disabling
            makes sense only when the LHS is a complex pattern.

        Examples
        --------

        For instance, we can have a very simple tensor, the outer product of the
        same vector,

        .. doctest::

            >>> dr = Drudge(SparkContext())
            >>> r = Range('R')
            >>> a, b = dr.set_dumms(r, symbols('a b c d e f'))[:2]
            >>> dr.add_default_resolver(r)
            >>> x = IndexedBase('x')
            >>> v = Vec('v')
            >>> tensor = dr.einst(x[a] * x[b] * v[a] * v[b])
            >>> str(tensor)
            'sum_{a, b} x[a]*x[b] * v[a] * v[b]'

        We can replace the indexed base by the product of a matrix with another
        indexed base,

        .. doctest::

            >>> o = IndexedBase('o')
            >>> y = IndexedBase('y')
            >>> res = tensor.subst(x[a], dr.einst(o[a, b] * y[b]))
            >>> str(res)
            'sum_{a, b, c, d} y[c]*y[d]*o[a, c]*o[b, d] * v[a] * v[b]'

        We can also make substitution on the vectors,

        .. doctest::

            >>> w = Vec('w')
            >>> res = tensor.subst(v[a], dr.einst(o[a, b] * w[b]))
            >>> str(res)
            'sum_{a, b, c, d} x[a]*x[b]*o[a, c]*o[b, d] * w[c] * w[d]'

        After the substitution, we can always make a simplification, at least to
        make the naming of the dummies more aesthetically pleasing,

        .. doctest::

            >>> res = res.simplify()
            >>> str(res)
            'sum_{a, b, c, d} x[c]*x[d]*o[c, a]*o[d, b] * w[a] * w[b]'

        """

        assert simult or not keep_decorated

        # Special case of the unity LHS.
        if lhs == 1:
            scalar_part = self.filter(lambda x: len(x.vecs) == 0)
            other_part = self.filter(lambda x: len(x.vecs) > 0)
            res = other_part + scalar_part * rhs
            return res

        # Shallow parsing of the left-hand side.

        if isinstance(lhs, Indexed):
            if_scalar = True
            if_indexed = True
            bases = (lhs.base,)
        elif isinstance(lhs, Symbol):
            if_scalar = True
            if_indexed = False
            bases = (lhs,)
        elif isinstance(lhs, ATerms):
            if_scalar = False
            if_indexed = True

            if len(lhs.terms) != 1:
                raise ValueError('Invalid LHS to substitute', lhs.terms)
            term = lhs.terms[0]
            if len(term.sums) != 0 or term.amp != 1:
                raise ValueError('Invalid LHS to substitute', term)
            vecs = term.vecs
            bases = {vec.base for vec in vecs}

            # In this way, lhs is either an SymPy expression (Indexed or Symbol)
            # or a tuple.
            lhs = vecs
        else:
            raise TypeError(
                'Invalid LHS for substitution', lhs,
                'expecting vector, indexed, or symbol'
            )

        if not all(self.has_base(base) for base in bases):
            return self

        # Shallow processing of the right-hand side.

        # We need to gather, and later broadcast all the terms.  The rational is
        # that the RHS is usually small in real problems.
        if isinstance(rhs, Tensor):
            rhs_terms = rhs.local_terms
            free_vars = rhs.free_vars
        else:
            rhs_terms = parse_terms(rhs)
            free_vars = set.union(*[term.free_vars for term in rhs_terms])

        if if_scalar and not all(
                term.is_scalar for term in rhs_terms
        ):
            raise ValueError('Invalid RHS for substituting a scalar', rhs)

        # Handling of the wilds
        if wilds is None:
            wilds = {}
            if if_indexed:
                if if_scalar:
                    index_bunches = [lhs.indices]
                else:
                    index_bunches = [vec.indices for vec in lhs]
                for indices in index_bunches:
                    wilds.update(
                        (index, Wild(index.name))
                        for index in indices if isinstance(index, Symbol)
                    )

        if if_scalar:
            lhs = lhs.xreplace(wilds)
        else:
            lhs = tuple(
                vec.map(lambda x: x.xreplace(wilds)) for vec in lhs
            )

        # XXX: Why is it necessary to expand the terms?
        # Expansion is probably only necessary for complex matching patterns,
        # which has not been implemented for now.
        rhs_terms = [
            expanded_term.subst(wilds)
            for term in rhs_terms for expanded_term in term.expand()
        ]

        # "Decorate" the RHS to ensure simultaneous substitution
        decr_vars = {}
        if simult:
            # for symbols and indexedBases
            # (Note that free_vars have been defined before the handling of the 
            # wilds            for var in free_vars:
            for var in free_vars:
                if isinstance(var, Symbol):
                    decr_vars[var] = Symbol(
                        var.name + _DECR_SUFFIX, **var._assumptions
                    )
                elif isinstance(var, IndexedBase):
                    decr_vars[var] = IndexedBase(
                        var.label.name + _DECR_SUFFIX, **var._assumptions
                    )
                else:
                    raise TypeError('Expecting Symbol or IndexedBase')

            rhs_terms = [term.subst(decr_vars) for term in rhs_terms]

            # for vectors
            rhs_terms = [
                term.map(vecs=self._decr_vecs(term.vecs)) for term in rhs_terms
            ]

        # Invoke the core facility.
        res = self.expand()._subst(
            lhs, rhs_terms, full_balance=full_balance, excl=excl
        )

        if not simult or keep_decorated:
            return res
        else:
            restore = functools.partial(self._restore, decr_vars=decr_vars)
            return res.map(restore)

    def _subst(
            self, lhs: typing.Union[typing.Tuple[Vec], Indexed, Symbol],
            rhs_terms, full_balance, excl=None
    ):
        """Core substitution function.

        This function assumes the self and the substituting terms are already
        fully expanded.  And the LHS and the RHS of the substitution have been
        replaced with the wilds if it is needed.
        """

        rhs_free_vars_per_term = [term.free_vars for term in rhs_terms]
        if len(rhs_free_vars_per_term) > 0:
            rhs_free_vars = set.union(*rhs_free_vars_per_term)
        else:
            rhs_free_vars = set()

        free_vars_local = (self.free_vars | rhs_free_vars)
        if excl is not None:
            free_vars_local |= excl

        free_vars = self._drudge.ctx.broadcast(free_vars_local)
        dumms = self._drudge.dumms
        full_simplify = self._drudge.full_simplify

        # We keep the dummbegs dictionary for each term and substitute all
        # appearances of the lhs one-by-one.

        subs_states = self._terms.map(lambda x: x.reset_dumms(
            dumms=dumms.value, excl=free_vars.value
        ))

        rhs_terms = self._drudge.ctx.broadcast(rhs_terms)

        if isinstance(lhs, (Indexed, Symbol)):
            res = nest_bind(subs_states, lambda x: subst_factor_term(
                x[0], lhs, rhs_terms.value,
                dumms=dumms.value, dummbegs=x[1], excl=free_vars.value,
                full_simplify=full_simplify
            ), full_balance=full_balance)
        else:
            res = nest_bind(subs_states, lambda x: subst_vec_term(
                x[0], lhs, rhs_terms.value,
                dumms=dumms.value, dummbegs=x[1], excl=free_vars.value
            ), full_balance=full_balance)

        res_terms = res.map(operator.itemgetter(0))
        return Tensor(
            self._drudge, res_terms, free_vars=free_vars_local, expanded=True
        )

    def subst_all(
            self, defs, simplify=False, full_balance=False, excl=None,
            simult=True, simult_all=False
    ):
        """Substitute all given definitions.

        The definitions should be given as an iterable of either
        :py:class:`TensorDef` instances or pairs of left-hand side and
        right-hand side of the substitutions. Note that, by default, the
        substitutions are going to be performed **according to the given
        order** one-by-one; simultaneous substitutions can be turned on by
        setting ``simult_all`` to be true.
        """

        # assert simult or not simult_all

        if simplify and not simult_all:
            sequentially_simplify = True
        else:
            sequentially_simplify = False

        res = self
        for i in defs:
            if isinstance(i, TensorDef):
                lhs = i.lhs
                rhs = i.rhs
            elif isinstance(i, Sequence) and len(i) == 2:
                lhs, rhs = i
            else:
                raise TypeError(
                    'Invalid substitution', i,
                    'expecting definition or LHS/RHS pair'
                )

            res = res.subst(
                lhs, rhs, full_balance=full_balance, excl=excl, 
                simult=simult, keep_decorated=simult_all
            )
            if sequentially_simplify:
                res = res.simplify().repartition()

            # Mostly to make the evaluation eagerly.
            if res.n_terms == 0:
                return res

        if simult_all:
            res = res.map(self._restore)
            if simplify:
                res = res.simplify().repartition()

        return res

    def rewrite(self, vecs, new_amp):
        """Rewrite terms with the given vectors in terms of the new amplitude.

        This method will rewrite the terms whose vector part patches the given
        vectors in terms of the given new amplitude.  And all terms rewritten
        into the same form will be aggregated into a single term.

        Parameters
        ----------

        vecs
            A vector or a product of vectors.  They should be written in terms
            of SymPy wild symbols when they need to be matched against different
            actual vectors.

        new_amp
            The amplitude that the matched terms should have.  They are usually
            written in terms of the same wild symbols as the wilds in the
            vectors.

        Returns
        -------

        rewritten
            The tensor with the requested terms rewritten in term of the given
            amplitude.

        defs
            The actual definitions of the rewritten amplitude.  One for each
            rewritten term in the result.

        """

        vecs_terms = parse_terms(vecs)
        invalid_vecs = ValueError(
            'Invalid vectors to rewrite', vecs,
            'expecting just vectors'
        )
        if len(vecs_terms) != 1:
            raise invalid_vecs
        vecs_term = vecs_terms[0]
        if len(vecs_term.sums) > 0 or vecs_term.amp != 1:
            raise invalid_vecs
        vecs = vecs_term.vecs

        rewritten = self._terms.map(
            lambda term: rewrite_term(term, vecs, new_amp)
        ).cache()
        new_terms = [
            i for i in rewritten.countByKey().keys() if i is not None
        ]

        get_term = operator.itemgetter(1)
        untouched_terms = rewritten.filter(
            lambda x: x[0] is None
        ).map(get_term)
        new_defs = {}
        for i in new_terms:
            def_terms = rewritten.filter(lambda x: x[0] == i).map(get_term)
            # Eagerly evaluate to circumvent a Spark bug.
            def_terms.cache()
            def_terms.count()
            new_defs[i.amp] = Tensor(self._drudge, def_terms)
            continue

        return Tensor(self._drudge, untouched_terms.union(
            self._drudge.ctx.parallelize(new_terms)
        )), new_defs

    def expand_sums(
            self, range_: Range, expander: Sum_expander,
            exts=None, conv_accs=None
    ):
        """Expand some symbolic summations.

        The basic structure of drudge always treat summation dummies as
        primitive scalar quantities.  However, sometimes, we would want to use a
        single summation over an abstract domain to symbolically denote multiple
        simple summations.

        Rather than adding this to the base system with possibly significant
        complication of the code base, here we seek a pragmatic solution, where
        we can just use plain symbols to denote complex summations, and this
        function can be used to expand the summation into its primitive
        components when we need to.

        Parameters
        ----------

        range_
            The range all summations over which are to be expanded.

        expander
            A callable to give information how the summation over a dummy should
            be expanded.  It will be called with a dummy that is summed over the
            given range, and it is expected to return an iterable of triples,
            formed by the new dummy and new range, followed by previous form the
            the dummy in the expression, which will be replaced by the new
            dummy.  Note that after the substitution, all the original dummies
            over the expanded range should no longer be present in the
            expression, or an error will be raised.

        exts
            Free symbols that are also going to be decomposed in the same way as
            the summations over the given range.

        conv_accs
            When it is set to an iterable of accessors, all their access of the
            original dummies will be substituted by the same accessor applied to
            the first dummy from the expander.  This can be helpful for cases
            where we just strip one (some) component(s) from a symbolic bundle.

        """
        return self.map(functools.partial(
            expand_sums_term, range_=range_, expander=expander,
            exts=exts, conv_accs=conv_accs
        ))

    #
    # Analytic gradient
    #

    def diff(self, variable, real=False, wirtinger_conj=False):
        """Differentiate the tensor to get the analytic gradient.

        By this function, support is provided for evaluating the derivative with
        respect to either a plain symbol or a tensor component.  This is
        achieved by leveraging the core differentiation operation to SymPy.  So
        very wide range of expressions are supported.

        .. warning::

            For non-analytic complex functions, this function gives the
            Wittinger derivative with respect to the given variable only.  The
            other non-vanishing derivative with respect to the conjugate needs
            to be evaluated by another invocation with ``wittinger_conj`` set to
            true.

        .. warning::

            The differentiation algorithm currently does **not** take the
            symmetry of the tensor to be differentiated with respect to into
            account.  For differentiate with respect to symmetric tensor,
            further symmetrization of the result might be needed.


        Parameters
        ----------

        variable
            The variable to differentiate with respect to.  It should be either
            a plain SymPy symbol or a indexed quantity.  When it is an indexed
            quantity, the indices should be plain symbols with resolvable range.

        real : bool
            If the variable is going to be assumed to be real.  Real variables
            has conjugate equal to themselves.

        wirtinger_conj : bool
            If we evaluate the Wirtinger derivative with respect to the
            conjugate of the variable.

        """

        if real and wirtinger_conj:
            raise ValueError(
                'Wittinger conjugate derivative vanishes for real variables'
            )

        if isinstance(variable, Indexed):

            symms = self._drudge.symms.value
            if_symm = (
                    variable.base in symms or
                    (variable.base, len(variable.indices)) in symms
            )
            if if_symm:
                warnings.warn(
                    'Gradient wrt to symmetric tensor {} '.format(variable) +
                    'might need further symmetrization'
                )

            # We need a copy.
            excl = set(self.free_vars)

            for i in variable.indices:
                if not isinstance(i, Symbol):
                    raise ValueError(
                        'Invalid index', i, 'expecting plain symbol'
                    )
                if i in excl:
                    raise ValueError(
                        'Invalid index', i,
                        'clashing with existing free symbols'
                    )
                excl.add(i)
                continue

            terms = self._reset_dumms(self._terms, excl=excl)

        elif isinstance(variable, Symbol):
            terms = self._terms
        else:
            raise TypeError('Invalid variable to differentiate', variable)

        return Tensor(self._drudge, self._diff(
            terms, variable, real=real, wirtinger_conj=wirtinger_conj
        ), expanded=True)

    def _diff(self, terms, variable, real, wirtinger_conj):
        """Differentiate the terms."""

        diff = terms.map(
            lambda x: diff_term(x, variable, real, wirtinger_conj)
        )
        # We have got to simplify here to avoid confuse users.
        return self._simplify_deltas(diff, False)

    #
    # Advanced manipulations.
    #

    def filter(self, crit):
        """Filter out terms satisfying the given criterion.

        The criterion needs to be a callable accepting :py:class:`Term`
        objects.  In the result, only terms for which the given criterion
        function evaluates to True will be retained.
        """
        return self.apply(
            lambda terms: terms.filter(crit),
            free_vars=None, repartitioned=False
        )

    def map(self, func):
        """Map the given function to the terms in the tensor.

        The given function should take a :py:class:`Term` object and return a
        :py:class:`Term` object.  The resulted tensor has the result of the
        application of the given callable to each of the terms as its terms.
        """
        return Tensor(self._drudge, self._terms.map(func))

    def bind(self, func):
        """Map the given function to the terms and flatten.

        The given callable need to return an iterable of :py:class:`Term`
        objects for each of the terms in the current tensor.  It is going to
        be applied to each of the terms, all the terms from all the results
        is going to form the terms of the resulted tensor.  This is the
        ``bind`` operation for the monad of lists in Haskell, or ``flatMap``
        in Java/Scala/Spark collection API.

        :py:meth:`filter` and :py:meth:`map` can be understood as special
        case of this method.
        """
        return Tensor(self._drudge, self._terms.flatMap(func))

    def map2scalars(self, action, skip_vecs=False, skip_ranges=True):
        """Map the given action to the scalars in the tensor.

        The given action should return SymPy expressions for SymPy expressions,
        the amplitude for each terms and the indices to the vectors, in the
        tensor.  Note that this function is more supposed for free variables and
        does not change the summation dummies in the terms.

        Parameters
        ----------

        action
            The callable to be applied to the scalars inside the tensor.

        skip_vecs
            When it is set, the callable will no longer be mapped to the indices
            to the vectors.  It could be used to boost the performance when we
            know that the action need no application on the indices.

        skip_ranges
            When it is set, the callable will no longer be mapped to the bounds
            of the ranges.

        """

        return Tensor(self._drudge, self._terms.map(lambda x: x.map(
            action, skip_vecs=skip_vecs, skip_ranges=skip_ranges
        )))

    def map2amps(self, action):
        """Map the given action to the amplitudes in the tensor.

        This method is a plain wrapper for calling :py:meth:`map2scalars` with
        vectors skipped.
        """

        return self.map2scalars(action, skip_vecs=True)

    #
    # Operations from the drudge
    #

    def __getattr__(self, item):
        """Try to see if the item is a tensor method from the drudge.

        This enables individual drudges to dynamically add domain-specific
        operations on tensors.
        """

        try:
            meth = self._drudge.get_tensor_method(item)
        except KeyError:
            raise AttributeError('Invalid operation name on tensor', item)

        return functools.partial(meth, self) class TensorDef(Tensor):
    """Definition of a tensor.

    A tensor definition is basically a tensor with a name.  In additional to
    being a tensor, a tensor definition also has a left-hand side.  When the
    tensor is zero-order, the left-hand side is simply a symbol.  When it has
    external indices, the base and external indices for the it are both
    stored.  The base is a :py:class:`Vec` instance for tensors with vector
    part or it is an SymPy IndexedBase for scalar tensors.  For instance,

    .. math::

        \sum_j o_{i, j} f_j

    can be construed as a tensor.  By storing it as an :py:class:`Tensor`
    object, we can have mathematical manipulations on it.  With an explicit
    left-hand side,

    .. math::

        t_j = \sum_j o_{i, j} f_j

    it is now an tensor definition, which can be handled by the current class.

    A tensor definition is a subclass of tensor.  With explicit storage of a
    left-hand side, it can be convenient to be used for :py:meth:`Tensor.subst`
    or :py:meth:`Tensor.subst_all` method, or it can be directly indexed.  For
    example, with the above definition stored in ``t_def``, for tensor ``v``
    holding :math:`\sum_j u_{i, j} t_j`,

    ::

        v.subst(t_def.lhs, t_def.rhs)

    or ::

        v.subst_all([t_def])

    or, with the :py:meth:`act` method, ::

        t_def.act(v)

    we can get

    .. math::

        \sum_{j, k} u_{i, j} o_{j, k} f_k

    Tensor definition can also be directly subscripts, like ``t_def[1]``
    gives :math:`sum_i o_{1, i} f_i`.

    By being :py:class:`Tensor` subclass, all tensor manipulations are
    supported.  Just the result will not be an automatically a definition.

    """

    def __init__(self, base, exts, tensor: Tensor):
        """Initialize the tensor definition.

        In the same way as the initializer for the :py:class:`Tensor` class,
        this initializer is also unlikely to be used directly in user code.
        Drudge methods :py:meth:`Drudge.define` and
        :py:meth:`Drudge.define_einst`, and their wrapper
        :py:meth:`Drudge.def_` can be more convenient.  Inside drudge
        scripts, operators ``<<=`` or ``<=`` can be used, see:ref:`drs intro`
        and :py:meth:`Drudge.exec_drs`.

        Parameters
        ----------

        base
            The base for the definition.  It will be normalized to the correct
            type depending on the external indices and the right-hand side.

        exts
            The iterable for external indices.  They can be either symbol/range
            pairs for external indices with explicit range, or they can also be
            a plain symbol for generic definitions.

        tensor
            The RHS of the definition.

        """

        if isinstance(tensor, Tensor):
            super().__init__(tensor.drudge, tensor.terms)
        else:
            raise TypeError(
                'Invalid LHS for tensor definition', tensor,
                'expecting a tensor instance'
            )

        self._exts = []
        for i in exts:
            explicit_ext = (
                    isinstance(i, Sequence) and len(i) == 2 and
                    isinstance(i[0], Symbol) and isinstance(i[1], Range)
            )
            if explicit_ext:
                self._exts.append(tuple(i))
            elif isinstance(i, Expr):
                self._exts.append((
                    i, None
                ))
            else:
                raise TypeError(
                    'Invalid external index', i,
                    'expecting dummy/range pair or a dummy.'
                )
            continue

        # Normalize the base.
        base_name = str(base)
        is_scalar = self.is_scalar
        if base == 1:
            # Special case for identity definition.
            self._base = base
        elif is_scalar and len(self._exts) == 0:
            self._base = Symbol(base_name)
        elif is_scalar:
            self._base = IndexedBase(base_name)
        else:
            self._base = Vec(base_name)

    #
    # Basic properties.
    #

    @property
    def rhs(self):
        """Get the right-hand-side of the definition.

        The result is the definition itself.  Kept here for backward
        compatibility.
        """
        return self

    @property
    def rhs_terms(self):
        """Gather the terms on the right-hand-side of the definition."""
        return self.local_terms

    @property
    def lhs(self):
        """Get the standard left-hand-side of the definition."""
        if len(self._exts) == 0:
            return self._base
        else:
            return self._base[tuple(i[0] for i in self._exts)]

    @property
    def base(self):
        """The base of the tensor definition."""
        return self._base

    @property
    def exts(self):
        """The external indices."""
        return self._exts

    #
    # Simple operations.
    #

    def simplify(self):
        """Simplify the tensor in the definition.
        """

        reset = self.reset_dumms()
        return TensorDef(reset.base, reset.exts, Tensor.simplify(reset))

    def reset_dumms(self, excl=None):
        """Reset the dummies in the definition.

        The external indices will take higher precedence over the summed indices
        inside the right-hand side.
        """

        dumms = self.drudge.dumms

        exts, ext_substs, dummbegs = Term.reset_sums(
            self._exts, dumms.value, excl=excl
        )

        free_vars = self.free_vars
        if excl is None:
            excl = free_vars
        else:
            excl |= free_vars
        excl -= {i for i, _ in self._exts}

        tensor = Tensor(
            self.drudge,
            self.terms.map(lambda x: x.reset_dumms(
                dumms=dumms.value, excl=excl,
                dummbegs=dict(dummbegs), add_substs=ext_substs
            )[0])
        )

        return TensorDef(self._base, exts, tensor)

    def __eq__(self, other):
        """Compare two tensor definitions for equality.

        Note that similar to the equality comparison of tensors, here we only
        compare the syntactic equality rather than the mathematical equality.
        The left-hand side is put into consideration only for comparison with
        another tensor definition.
        """

        rhs_eq = super().__eq__(other)

        return rhs_eq and (
            self.lhs == other.lhs if isinstance(other, TensorDef) else True
        )

    #
    # Representations.
    #

    def __str__(self):
        """Form simple readable string for a definition.
        """

        return ' = '.join([str(self.lhs), super().__str__()])

    def latex(self, **kwargs):
        r"""Get the latex form for the tensor definition.

        The result will just be the form from :py:meth:`Tensor.latex` with the
        RHS prepended.

        Parameters
        ----------

        kwargs
            All keyword parameters are forwarded to the
            :py:meth:`Drudge.format_latex` method.

        """
        return self.drudge.format_latex(self, **kwargs)

    def display(self, if_return=True, **kwargs):
        """Display the tensor definition in interactive notebook sessions.

        The parameters here all have the same meaning as in
        :py:meth:`Tensor.display`.
        """

        form = Math(self.latex(**kwargs))
        if if_return:
            return form
        else:
            display(form)
            return

    #
    # Substitution.
    #

    def act(self, tensor, wilds=None, full_balance=False):
        """Act the definition on a tensor.

        This method is the active voice version of the :py:meth:`Tensor.subst`
        function.  All appearances of the defined object in the tensor will be
        substituted.

        """

        if not isinstance(tensor, Tensor):
            tensor = self.rhs.drudge.sum(tensor)

        return tensor.subst(
            self.lhs, self.rhs, wilds=wilds, full_balance=full_balance
        )

    def __getitem__(self, item):
        """Get the tensor when the definition is indexed.
        """

        if not isinstance(item, Sequence):
            item = (item,)

        n_exts = len(self._exts)
        if len(item) != n_exts:
            raise ValueError(
                'Invalid subscripts', item, 'expecting', n_exts
            )

        return self.act(self._base[item])

    #
    # Pickling support
    #

    def __getstate__(self):
        """Get the current state of the definition.
        """
        return self._base, self._exts, self.local_terms

    def __setstate__(self, state):
        """Set the state for the new definition.
        """
        super().__setstate__(state[2])
        self.__init__(state[0], state[1], self)
        return class Vec(ATerms, CantSympify):
    """Vectors.

    Vectors are the basic non-commutative quantities.  Its objects consist of an
    label for its base and some indices.  The label is allowed to be any
    hashable and ordered Python object, although small objects, like string, are
    advised.  The indices are always sympified into SymPy expressions.

    Its objects can be created directly by giving the label and indices, or
    existing vector objects can be subscribed to get new ones.  The semantics is
    similar to Haskell functions.

    Note that users cannot directly assign to the attributes of this class.

    This class can be used by itself, it can also be subclassed for special
    use cases.

    Despite very different internal data structure, the this class is attempted
    to emulate the behaviour of the SymPy ``IndexedBase`` class

    """

    __slots__ = ['_label', '_indices']

    def __init__(self, label, indices=()):
        """Initialize a vector.

        Atomic indices are added as the only index.  Iterable values will
        have all of its entries added.
        """
        self._label = label
        if not isinstance(indices, Iterable):
            indices = (indices,)
        self._indices = tuple(ensure_expr(i, 'vector index') for i in indices)

    @property
    def label(self):
        """The label for the base of the vector.
        """
        return self._label

    @property
    def base(self):
        """The base of the vector.

        This base can be subscribed to get other vectors.
        """
        return Vec(self._label, [])

    @property
    def indices(self):
        """The indices to the vector.
        """
        return self._indices

    def __getitem__(self, item):
        """Append the given indices to the vector.

        When multiple new indices are to be given, they have to be given as a
        tuple.
        """

        if not isinstance(item, tuple):
            item = (item,)

        # Pay attention to subclassing.
        return type(self)(self._label, itertools.chain(self._indices, item))

    def __repr__(self):
        """Form repr string form the vector."""
        return ''.join([
            type(self).__name__, '(', repr(self._label), ', (',
            ', '.join(repr(i) for i in self._indices),
            '))'
        ])

    def __str__(self):
        """Form a more readable string representation."""

        label = str(self._label)
        if len(self._indices) > 0:
            indices = ''.join([
                '[', ', '.join(str(i) for i in self._indices), ']'
            ])
        else:
            indices = ''

        return label + indices

    def __hash__(self):
        """Compute the hash value of a vector."""
        return hash((self._label, self._indices))

    def __eq__(self, other):
        """Compares the equality of two vectors."""
        return (
                (isinstance(self, type(other)) or
                 isinstance(other, type(self))) and
                self._label == other.label and self._indices == other.indices
        )

    @property
    def sort_key(self):
        """The sort key for the vector.

        This is a generic sort key for vectors.  Note that this is only useful
        for sorting the simplified terms and should not be used in the
        normal-ordering operations.
        """

        key = [self._label]
        key.extend(sympy_key(i) for i in self._indices)
        return key

    #
    # Misc facilities
    #

    def map(self, func):
        """Map the given function to indices."""
        return Vec(self._label, (func(i) for i in self._indices))

    @property
    def terms(self):
        """Get the terms from the vector.

        This is for the user input.
        """

        return [Term((), _UNITY, (self,))] class Range:
    """A symbolic range that can be summed over.

    This class is for symbolic ranges that is going to be summed over in
    tensors.  Each range should have a label, and optionally lower and upper
    bounds, which should be both given or absent.  The label can be any hashable
    and ordered Python type.  The bounds will not be directly used for symbolic
    computation, but rather designed for printers and conversion to SymPy
    summation.  Note that ranges are assumed to be atomic and disjoint.  Even in
    the presence of lower and upper bounds, unequal ranges are assumed to be
    disjoint.

    .. warning::

        Equality comparison and hashing of ranges are *completely* based on the
        label, without consideration of the bounds.  So ranges with the same
        label but different bounds will be considered equal.  This is to
        facilitate the usage of a same dummy set for ranges with slight
        variation in the bounds.

    .. warning::

         Unequal ranges are always assumed to be disjoint.

    """

    __slots__ = [
        '_label',
        '_lower',
        '_upper'
    ]

    def __init__(self, label, lower=None, upper=None):
        """Initialize the symbolic range.
        """
        self._label = label
        self._lower = (
            ensure_expr(lower, 'lower bound') if lower is not None else lower
        )

        if self._lower is None:
            if upper is not None:
                raise ValueError('lower range has not been given.')
            else:
                self._upper = None
        else:
            if upper is None:
                raise ValueError('upper range has not been given.')
            else:
                self._upper = ensure_expr(upper, 'upper bound')

    @property
    def label(self):
        """The label of the range."""
        return self._label

    @property
    def lower(self):
        """The lower bound of the range."""
        return self._lower

    @property
    def upper(self):
        """The upper bound of the range."""
        return self._upper

    @property
    def size(self):
        """The size of the range.

        This property given None for unbounded ranges.  For bounded ranges, it
        is the difference between the lower and upper bound.  Note that this
        contradicts the deeply entrenched mathematical convention of including
        other ends for a range.  But it does gives a lot of convenience and
        elegance.
        """

        return self._upper - self._lower if self.bounded else None

    @property
    def bounded(self):
        """If the range is explicitly bounded."""
        return self._lower is not None

    @property
    def args(self):
        """The arguments for range creation.

        When the bounds are present, we have a triple, or we have a singleton
        tuple of only the label.
        """

        if self.bounded:
            return self._label, self._lower, self._upper
        else:
            return self._label,

    def __hash__(self):
        """Hash the symbolic range.
        """
        return hash(self._label)

    def __eq__(self, other):
        """Compare equality of two ranges.
        """
        return isinstance(other, type(self)) and (
                self._label == other.label
        )

    def __repr__(self):
        """Form the representative string.
        """
        return ''.join([
            'Range(', ', '.join(repr(i) for i in self.args), ')'
        ])

    def __str__(self):
        """Form readable string representation.
        """
        return str(self._label)

    @property
    def sort_key(self):
        """The sort key for the range."""
        key = [self._label]
        if self.bounded:
            key.extend(sympy_key(i) for i in [self._lower, self._upper])
        return key

    def replace_label(self, new_label):
        """Replace the label of a given range.

        The bounds will be the same as the original range.
        """

        return Range(new_label, self._lower, self._upper)

    def __lt__(self, other):
        """Compare two ranges.

        This method is meant to skip explicit calling of the sort key when it is
        not convenient.
        """

        if not isinstance(other, Range):
            raise TypeError('Invalid range to compare', other)

        return self.sort_key < other.sort_key

    def __getitem__(self, item):
        """Get a new range with explicit bounds set.

        This can be used on both symbolic range to give them explicit bounds,
        and on bounded range to set new bounds for them.
        """
        if not isinstance(item, tuple) or len(item) != 2:
            raise TypeError(
                'Invalid bounds for range, expecting upper and lower range',
                item
            )

        lower, upper = [sympify(i) for i in item]
        return Range(self._label, lower=lower, upper=upper)

    def map(self, func):
        """Map the given function to the bounds of the range.

        For ranges without bounds, the same range is simply returned.
        """
        if self.bounded:
            return self[func(self.lower), func(self.upper)]
        else:
            return self def try_resolve_range(i, sums_dict, resolvers):
    """Attempt to resolve the range of an expression.

    None will be returned if it cannot be resolved.
    """

    for resolver in itertools.chain([sums_dict], resolvers):

        if isinstance(resolver, Mapping):
            if i in resolver:
                return resolver[i]
            else:
                continue
        elif isinstance(resolver, Callable):
            range_ = resolver(i)
            if range_ is None:
                continue
            else:
                if isinstance(range_, Range):
                    return range_
                else:
                    raise TypeError('Invalid range: ', range_,
                                    'from resolver', resolver,
                                    'expecting range or None')
        else:
            raise TypeError('Invalid resolver: ', resolver,
                            'expecting callable or mapping')

    # Never resolved nor error found.
    return None class Term(ATerms):
    """Terms in tensor expression.

    This is the core class for storing symbolic tensor expressions.  The
    actual symbolic tensor type is just a shallow wrapper over a list of
    terms.  It is basically comprised of three fields, a list of summations,
    a SymPy expression giving the amplitude, and a list of non-commutative
    vectors.
    """

    __slots__ = [
        '_sums',
        '_amp',
        '_vecs',
        '_free_vars',
        '_dumms'
    ]

    def __init__(
            self, sums: typing.Tuple[typing.Tuple[Symbol, Range], ...],
            amp: Expr, vecs: typing.Tuple[Vec, ...],
            free_vars: typing.FrozenSet[Symbol] = None,
            dumms: typing.Mapping[Symbol, Range] = None,
    ):
        """Initialize the tensor term.

        Users seldom have the need to create terms directly by this function.
        So this constructor is mostly a developer function, no sanity checking
        is performed on the input for performance.  Most importantly, this
        constructor does **not** copy either the summations or the vectors and
        directly expect them to be tuples (for hashability).  And the amplitude
        is **not** simpyfied.

        Also, it is important that the free variables and dummies dictionary be
        given only when they really satisfy what we got for them.

        """

        # For performance reason, no checking is done.
        #
        # Uncomment for debugging.
        # valid = (
        #     isinstance(sums, tuple) and isinstance(amp, Expr)
        #     and isinstance(vecs, tuple)
        # )
        # if not valid:
        #     raise TypeError('Invalid argument to term creation')

        self._sums = sums
        self._amp = amp
        self._vecs = vecs

        self._free_vars = free_vars
        self._dumms = dumms

    @property
    def sums(self):
        """The summations of the term."""
        return self._sums

    @property
    def amp(self) -> Expr:
        """The amplitude expression."""
        return self._amp

    @property
    def vecs(self):
        """The vectors in the term."""
        return self._vecs

    @property
    def is_scalar(self):
        """If the term is a scalar."""
        return len(self._vecs) == 0

    @property
    def args(self):
        """The triple of summations, amplitude, and vectors."""
        return self._sums, self._amp, self._vecs

    def __hash__(self):
        """Compute the hash of the term."""
        return hash(self.args)

    def __eq__(self, other):
        """Evaluate the equality with another term."""
        return isinstance(other, type(self)) and self.args == other.args

    def __repr__(self):
        """Form the representative string of a term."""
        return 'Term(sums=[{}], amp={}, vecs=[{}])'.format(
            ', '.join(repr(i) for i in self._sums),
            repr(self._amp),
            ', '.join(repr(i) for i in self._vecs)
        )

    def __str__(self):
        """Form the readable string representation of a term."""
        if len(self._sums) > 0:
            header = 'sum_{{{}}} '.format(
                ', '.join(str(i[0]) for i in self._sums))
        else:
            header = ''
        factors = [str(self._amp)]
        factors.extend(str(i) for i in self._vecs)
        return header + ' * '.join(factors)

    @property
    def sort_key(self):
        """The sort key for a term.

        This key attempts to sort the terms by complexity, with simpler terms
        coming earlier.  This capability of sorting the terms will make the
        equality comparison of multiple terms easier.

        This sort key also ensures that terms that can be merged are always put
        into adjacent positions.

        """

        vec_keys = [i.sort_key for i in self._vecs]
        sum_keys = [(i[1].sort_key, sympy_key(i[0])) for i in self._sums]

        return (
            len(vec_keys), vec_keys,
            len(sum_keys), sum_keys,
            sympy_key(self._amp)
        )

    @property
    def terms(self):
        """The singleton list of the current term.

        This property is for the rare cases where direct construction of tensor
        inputs from SymPy expressions and vectors are not sufficient.
        """
        return [self]

    def scale(self, factor):
        """Scale the term by a factor.
        """
        return Term(self._sums, self._amp * factor, self._vecs)

    def mul_term(self, other, dumms=None, excl=None):
        """Multiply with another tensor term.

        Note that by this function, the free symbols in the two operands are not
        automatically excluded.
        """
        lhs, rhs = self.reconcile_dumms(other, dumms, excl)
        return Term(
            lhs.sums + rhs.sums, lhs.amp * rhs.amp, lhs.vecs + rhs.vecs
        )

    def comm_term(self, other, dumms=None, excl=None):
        """Commute with another tensor term.

        In ths same way as the multiplication operation, here the free symbols
        in the operands are not automatically excluded.
        """
        lhs, rhs = self.reconcile_dumms(other, dumms, excl)
        sums = lhs.sums + rhs.sums
        amp0 = lhs.amp * rhs.amp
        return [
            Term(sums, amp0, lhs.vecs + rhs.vecs),
            Term(sums, -amp0, rhs.vecs + lhs.vecs)
        ]

    def reconcile_dumms(self, other, dumms, excl):
        """Reconcile the dummies in two terms."""
        lhs, dummbegs = self.reset_dumms(dumms, excl=excl)
        rhs, _ = other.reset_dumms(dumms, dummbegs=dummbegs, excl=excl)
        return lhs, rhs

    #
    # SymPy related
    #

    @property
    def exprs(self):
        """Loop over the sympy expression in the term.

        Note that the summation dummies are not looped over.
        """
        for _, range_ in self._sums:
            if range_.bounded:
                yield range_.lower
                yield range_.upper
            continue

        yield self._amp

        for vec in self._vecs:
            yield from vec.indices
            continue

    @property
    def free_vars(self):
        """The free symbols used in the term.
        """

        if self._free_vars is None:
            dumms = self.dumms
            self._free_vars = set(
                i for expr in self.exprs for i in expr.atoms(Symbol)
                if i not in dumms
            )
        return self._free_vars

    @property
    def dumms(self):
        """Get the mapping from dummies to their range.
        """

        if self._dumms is None:
            self._dumms = dict(self._sums)
        return self._dumms

    @property
    def amp_factors(self):
        """The factors in the amplitude expression.

        This is a convenience wrapper over :py:meth:`get_amp_factors` for the
        case of no special additional symbols.
        """

        return self.get_amp_factors()

    def get_amp_factors(self, *special_symbs, monom_only=True):
        """Get the factors in the amplitude and the coefficient.

        The indexed factors and factors involving dummies or the symbols in the
        given special symbols set will be returned as a list, with the rest
        returned as a single SymPy expression.

        When ``monom_only`` is set, Error will be raised if the amplitude is not
        a monomial.
        """

        amp = self._amp

        if monom_only and isinstance(amp, Add):
            raise ValueError('Invalid amplitude: ', amp, 'expecting monomial')

        if isinstance(amp, Mul):
            all_factors = amp.args
        else:
            all_factors = (amp,)

        dumms = self.dumms

        factors = []
        coeff = _UNITY
        for factor in all_factors:
            need_treatment = any(
                (i in dumms or any(i in j for j in special_symbs))
                for i in factor.atoms(Symbol)
            ) or factor.has(Indexed)
            if need_treatment:
                factors.append(factor)
            else:
                coeff *= factor
            continue

        return factors, coeff

    def map(
            self, func=lambda x: x, sums=None, amp=None, vecs=None,
            skip_vecs=False, skip_ranges=False
    ):
        """Map the given function to the SymPy expressions in the term.

        The given function will **not** be mapped to the dummies in the
        summations.  When operations on summations are needed, a **tuple**
        for the new summations can be given.

        By the default function of the identity function, this function can also
        be used to replace the summation list, the amplitude expression, or the
        vector part.
        """

        res_sums = self._sums if sums is None else sums
        if not skip_ranges:
            res_sums = tuple(
                (dumm, range_.map(func)) for dumm, range_ in res_sums
            )

        res_amp = func(self._amp if amp is None else amp)

        res_vecs = self._vecs if vecs is None else vecs
        if not skip_vecs:
            res_vecs = tuple(vec.map(func) for vec in res_vecs)

        return Term(res_sums, res_amp, res_vecs)

    def subst(self, substs, sums=None, amp=None, vecs=None, purge_sums=False):
        """Perform symbol substitution on the SymPy expressions.

        After the replacement of the fields given, the given substitutions are
        going to be performed using SymPy ``xreplace`` method simultaneously.

        If purge sums is set, the summations whose dummy is substituted is going
        to be removed.

        """

        if sums is None:
            sums = self._sums
        if purge_sums:
            sums = tuple(sum_ for sum_ in sums if sum_[0] not in substs)

        return self.map(
            lambda x: x.xreplace(substs), sums=sums, amp=amp, vecs=vecs,
            skip_ranges=False
        )

    def reset_dumms(self, dumms, dummbegs=None, excl=None, add_substs=None):
        """Reset the dummies in the term.

        The term with dummies reset will be returned alongside with the new
        dummy begins dictionary.  Note that the dummy begins dictionary will be
        mutated if one is given.

        ValueError will be raised when no more dummies are available.
        """

        new_sums, substs, dummbegs = self.reset_sums(
            self._sums, dumms, dummbegs, excl
        )
        if add_substs is not None:
            substs.update(add_substs)

        return self.subst(substs, new_sums), dummbegs

    @staticmethod
    def reset_sums(sums, dumms, dummbegs=None, excl=None):
        """Reset the given summations.

        The new summation list, substitution dictionary, and the new dummy begin
        dictionary will be returned.
        """

        if dummbegs is None:
            dummbegs = {}

        new_sums = []
        substs = {}
        for dumm_i, range_i in sums:

            # For linter.
            new_dumm = None
            new_beg = None

            beg = dummbegs[range_i] if range_i in dummbegs else 0
            for i in itertools.count(beg):

                try:
                    tentative = dumms[range_i][i]
                except KeyError:
                    raise ValueError('Dummies for range', range_i,
                                     'is not given')
                except IndexError:
                    raise ValueError('Dummies for range', range_i, 'is used up')

                if excl is None or tentative not in excl:
                    new_dumm = tentative
                    new_beg = i + 1
                    break
                else:
                    continue

            new_sums.append((new_dumm, range_i))
            substs[dumm_i] = new_dumm
            dummbegs[range_i] = new_beg

            continue

        return tuple(new_sums), substs, dummbegs

    #
    # Amplitude simplification
    #

    def simplify_deltas(self, resolvers):
        """Simplify deltas in the amplitude of the expression."""

        new_amp, substs = simplify_deltas_in_expr(
            self.dumms, self._amp, resolvers
        )

        # Note that here the substitutions needs to be performed in order.
        return self.subst(substs, purge_sums=True, amp=new_amp)

    def simplify_trivial_sums(self):
        """Simplify the trivial summations in the term.

        Trivial summations are the concrete summations that do not have any
        involvement at all.
        """

        involved = {
            i for expr in self.exprs for i in expr.atoms(Symbol)
        }

        new_sums = []
        factor = _UNITY
        dirty = False

        for symb, range_ in self._sums:
            if symb not in involved and range_.bounded:
                dirty = True
                factor *= range_.size
            else:
                new_sums.append((symb, range_))
            continue

        if dirty:
            return Term(tuple(new_sums), factor * self._amp, self._vecs)
        else:
            return self

    def expand(self):
        """Expand the term into many terms."""
        expanded_amp = self.amp.expand()

        if expanded_amp == 0:
            return []
        elif isinstance(expanded_amp, Add):
            amp_terms = expanded_amp.args
        else:
            amp_terms = (expanded_amp,)

        return [self.map(lambda x: x, amp=i) for i in amp_terms]

    #
    # Canonicalization.
    #

    def canon(self, symms=None, vec_colour=None):
        """Canonicalize the term.

        The given vector colour should be a callable accepting the index
        within vector list (under the keyword ``idx``) and the vector itself
        (under keyword ``vec``).  By default, vectors has colour the same as
        its index within the list of vectors.

        Note that whether or not colours for the vectors are given, the vectors
        are never permuted in the result.
        """

        if symms is None:
            symms = {}

        # Factors to canonicalize.
        factors = []

        # Additional information for factor reconstruction.
        #
        # It has integral placeholders for vectors and scalar factors without
        # any indexed quantity, the expression with (the only) indexed replaced
        # by the placeholder for factors with indexed.
        factors_info = []
        vec_factor = 1
        unindexed_factor = 2

        #
        # Get the factors in the amplitude.
        #

        # Cache globals for performance.
        wrapper_base = _WRAPPER_BASE
        treated_placeholder = _TREATED_PLACEHOLDER

        # Extractors for subexpressions requiring explicit treatment, defined
        # here to avoid repeated list and function creation for each factor.
        to_treats = []

        def replace_to_treats(expr: Expr):
            """Replace the quantities need treatment inside the factor.

            These quantities are explicitly indexed quantities and multi-valence
            functions either with explicit symmetry given or with multiple
            arguments involving dummies.
            """

            if_treat = False
            if expr.func == Indexed:
                if_treat = True
            elif len(expr.args) > 1:
                if expr.func in symms or (expr.func, len(expr.args)) in symms:
                    if_treat = True
                else:
                    if_treat = sum(1 for arg in expr.args if any(
                        arg.has(dumm) for dumm, _ in self.sums
                    )) > 1

            if if_treat:
                to_treats.append(expr)
                return treated_placeholder
            else:
                return expr

        amp_factors, coeff = self.amp_factors
        for i in amp_factors:

            amp_no_treated = i.replace(
                lambda _: True, NonsympifiableFunc(replace_to_treats)
            )

            n_to_treats = len(to_treats)
            if n_to_treats > 1:
                raise ValueError(
                    'Overly complex factor', i
                )
            elif n_to_treats == 1:
                to_treat = to_treats[0]
                to_treats.clear()  # Clean the container for the next factor.

                factors.append((
                    to_treat, (
                        _COMMUTATIVE,
                        sympy_key(to_treat.base.label)
                        if isinstance(to_treat, Indexed)
                        else sympy_key(to_treat.func),
                        sympy_key(amp_no_treated)
                    )
                ))
                factors_info.append(amp_no_treated)

            else:  # No part needing explicit treatment.

                # When the factor never has an indexed base, we treat it as
                # indexing a uni-valence internal indexed base.
                factors.append((
                    wrapper_base[i], (_COMMUTATIVE,)
                ))
                factors_info.append(unindexed_factor)

            continue

        #
        # Get the factors in the vectors.
        #

        for i, v in enumerate(self._vecs):
            colour = i if vec_colour is None else vec_colour(
                idx=i, vec=v, term=self
            )
            factors.append((
                v, (_NON_COMMUTATIVE, colour)
            ))
            factors_info.append(vec_factor)
            continue

        #
        # Invoke the core simplification.
        #

        res_sums, canoned_factors, canon_coeff = canon_factors(
            self._sums, factors, symms
        )

        #
        # Reconstruct the result
        #

        res_amp = coeff * canon_coeff
        res_vecs = []
        for i, j in zip(canoned_factors, factors_info):

            if j == vec_factor:
                # When we have a vector.
                res_vecs.append(i)
            elif j == unindexed_factor:
                res_amp *= i.indices[0]
            else:
                res_amp *= j.xreplace({treated_placeholder: i})
            continue

        return Term(tuple(res_sums), res_amp, tuple(res_vecs))

    def canon4normal(self, symms):
        """Canonicalize the term for normal-ordering.

        This is the preparation task for normal ordering.  The term will be
        canonicalized with all the vectors considered the same.  And the dummies
        will be reset internally according to the summation list.
        """

        # Make the internal dummies factory to canonicalize the vectors.
        dumms = collections.defaultdict(list)
        for i, v in self.sums:
            dumms[v].append(i)
        for i in dumms.values():
            # This is important for ordering vectors according to SymPy key in
            # normal ordering.
            i.sort(key=sympy_key)

        canon_term = (
            self.canon(symms=symms, vec_colour=lambda idx, vec, term: 0)
                .reset_dumms(dumms)[0]
        )

        return canon_term

    def has_base(self, base):
        """Test if the given base is present in the current term."""

        if isinstance(base, (IndexedBase, Symbol)):
            return self._amp.has(base)
        elif isinstance(base, Vec):
            label = base.label
            return any(
                i.label == label for i in self._vecs
            )
        else:
            raise TypeError('Invalid base to test presence', base) def sympy_key(expr):
    """Get the key for ordering SymPy expressions.

    This function assumes that the given expression is already sympified.
    """

    return count_ops(expr), default_sort_key(expr) def ensure_expr(obj, role=''):
    """Sympify the given object into an expression."""
    return ensure_sympify(obj, role, Expr) class EnumSymbs(AtomicExpr, metaclass=_EnumSymbsMeta):
    """Base class for enumeration symbols.

    Subclasses can set `_symbs_` inside the class body to be a sequence of
    string pairs.  Then attributes named after the first field of the pairs will
    be created, with the LaTeX form controlled by the second pair.

    The resulted values are valid SymPy expressions.  They are ordered according
    to their order in the given enumeration sequence.

    """

    _symbs_ = None

    _VAL_FIELD = '_val_index'
    __slots__ = [_VAL_FIELD]

    def __init__(self, val_index):
        """Initialize the concrete symbol object.
        """
        if self._symbs_ is None:
            raise ValueError('Base EnumSymbs class cannot be instantiated')
        setattr(self, self._VAL_FIELD, val_index)

    @property
    def args(self):
        """The argument for SymPy."""
        return Integer(getattr(self, self._VAL_FIELD)),

    def __str__(self):
        """Get the string representation of the symbol."""
        return self._symbs_[getattr(self, self._VAL_FIELD)][0]

    def __repr__(self):
        """Get the machine readable string representation."""
        return '.'.join([type(self).__name__, str(self)])

    _op_priority = 20.0

    def __eq__(self, other):
        """Test two values for equality."""
        return isinstance(other, type(self)) and self.args == other.args

    def __hash__(self):
        """Hash the concrete symbol object."""
        return hash(repr(self))

    def __lt__(self, other):
        """Test two values for less than order.

        The order will be based on the order given in the class.
        """
        return self.args < other.args

    def __gt__(self, other):
        """Test two values for greater than."""
        return self.args > other.args

    def __sub__(self, other: Expr):
        """Subtract the current value with another.

        This method is mainly to be able to work together with the Kronecker
        delta class from SymPy.  The difference is only guaranteed to have
        correct ``is_zero`` property.  The actual difference might not make
        mathematical sense.
        """

        if isinstance(other, type(self)):
            return self.args[0] - other.args[0]
        elif len(other.atoms(Symbol)) == 0:
            raise ValueError(
                'Invalid operation for ', (self, other),
                'concrete symbols can only be subtracted for the same type'
            )
        else:
            # We are having a symbolic value at the other expression.  We just
            # need to make sure that the result is fuzzy.
            assert other.is_zero is None
            return other

    def __rsub__(self, other):
        """Subtract the current value from the other expression.

        Only the ``is_zero`` property is guaranteed.
        """
        return self.__sub__(other)

    def sort_key(self, order=None):
        return (
            self.class_key(),
            (1, tuple(i.sort_key() for i in self.args)),
            S.One.sort_key(), S.One
        )

    def _latex(self, _):
        """Print itself as LaTeX code."""
        return self._symbs_[self.args[0]][1] class WickDrudge(Drudge, abc.ABC):
    """Drudge for Wick-style algebras.

    A Wick-style algebra is an algebraic system where the commutator between any
    generators of the algebra is a simple scalar value.  This drudge will
    attempt to put the vectors into normal order based on the given comparator
    and contractor by Wick theorem.

    Normally, subclasses need to override the properties :py:attr:`phase`,
    :py:attr:`contractor`, and :py:attr:`comparator` with domain-specific
    knowledge.
    """

    def __init__(self, *args, wick_parallel=0, **kwargs):
        """Initialize the Wick drudge.

        This level just have one option to handle, the parallelism option.
        """
        super().__init__(*args, **kwargs)
        self._wick_parallel = wick_parallel

    @property
    def wick_parallel(self):
        """Get the Wick parallelism level."""
        return self._wick_parallel

    @wick_parallel.setter
    def wick_parallel(self, level):
        """Set the Wick parallelism level.

        Valid values are ``0``, for normal problems, ``1``, for harder problems,
        and ``2``, for really hard expressions containing just a few terms.

        """

        if level not in {0, 1, 2}:
            raise ValueError(
                'Invalid parallel level for Wick expansion', level
            )
        self._wick_parallel = level

    @property
    @abc.abstractmethod
    def contractor(self) -> typing.Callable[[Vec, Vec, Term], Expr]:
        """Get the contractor for the algebraic system.

        The contractor is going to be called with two vectors to return the
        value of their contraction.

        """
        pass

    @property
    @abc.abstractmethod
    def phase(self):
        """Get the phase for the commutation rule.

        The phase should be a constant defining the phase of the commutation
        rule.
        """
        pass

    @property
    @abc.abstractmethod
    def comparator(self) -> typing.Callable[[Vec, Vec, Term], bool]:
        """Get the comparator for the canonicalized vectors.

        The normal ordering operation will be performed according to this
        comparator.  It will be called with two **canonicalized vectors** for a
        boolean value.  True should be returned if the first given vector is
        less than the second vector.  The two vectors will be attempted to be
        transposed when False is returned.

        """
        pass

    def normal_order(self, terms: RDD, **kwargs):
        """Normal order the terms according to generalized Wick theorem.

        The actual expansion is based on the information given in the subclasses
        by the abstract properties.

        """
        comparator = kwargs.pop('comparator', self.comparator)
        contractor = kwargs.pop('contractor', self.contractor)
        if len(kwargs) != 0:
            raise ValueError(
                'Invalid arguments to Wick normal order', kwargs
            )

        phase = self.phase
        symms = self.symms
        resolvers = self.resolvers

        terms.cache()
        terms_to_proc = terms.filter(lambda x: len(x.vecs) > 1)
        keep_top = 0 if comparator is None else 1
        terms_to_keep = terms.filter(lambda x: len(x.vecs) <= keep_top)
        terms_to_proc.cache()
        if terms_to_proc.count() == 0:
            return terms_to_keep

        # Triples: term, contractions, schemes.
        wick_terms = terms_to_proc.map(lambda x: _prepare_wick(
            x, comparator, contractor, symms.value, resolvers.value
        ))

        if self._wick_parallel == 0:

            normal_ordered = wick_terms.flatMap(lambda x: [
                _form_term_from_wick(x[0], x[1], phase, resolvers.value, i)
                for i in x[2]
            ])

        elif self._wick_parallel == 1:

            flattened = wick_terms.flatMap(
                lambda x: [(x[0], x[1], i) for i in x[2]]
            )
            if self._num_partitions is not None:
                flattened = flattened.repartition(self._num_partitions)

            normal_ordered = flattened.map(lambda x: _form_term_from_wick(
                x[0], x[1], phase, resolvers.value, x[2]
            ))

        elif self._wick_parallel == 2:

            # This level of parallelism is reserved for really hard problems.
            expanded = []
            for term, contrs, schemes in wick_terms.collect():
                # To work around a probable Spark bug.  Problem occurs when we
                # have closures inside a loop to be distributed out.
                form_term = functools.partial(
                    _form_term_from_wick_bcast, term, contrs, phase, resolvers
                )

                curr = self._ctx.parallelize(schemes).map(form_term)
                expanded.append(curr)
                continue

            normal_ordered = self._ctx.union(expanded)

        else:
            raise ValueError(
                'Invalid Wick expansion parallel level', self._wick_parallel
            )

        return terms_to_keep.union(normal_ordered)
import functools
import typing
import warnings
from pyspark import RDD
from sympy import (
    KroneckerDelta, IndexedBase, Expr, Symbol, Rational, symbols, conjugate,
    factorial
)
from ._tceparser import parse_tce_out
from .canon import NEG, IDENT
from .canonpy import Perm
from .drudge import Tensor, TensorDef
from .term import Vec, Range, try_resolve_range, Term
from .utils import sympy_key, ensure_expr, EnumSymbs
from .wick import WickDrudge
        end2 = 2 * n_body

        cycl = Perm(
            self._form_cycl(begin1, end1) + self._form_cycl(begin2, end2)
        )
        transp = Perm(
            self._form_transp(begin1, end1) + self._form_transp(begin2, end2)
        )

        self.set_symm(base, cycl, transp, valence=2 * n_body)

        return

    def set_dbbar_base(self, base: IndexedBase, n_body: int, n_body2=None):
        """Set an indexed base as a double-bar interaction.

        A double barred interaction has full permutation symmetry among its
        first half of slots and its second half individually.  For fermion
        field, the permutation is assumed to be anti-commutative.

        The size of the second half can be given by another optional argument,
        or it is assumed to have the same size as the first half.  It can also
        be zero, which gives one chunk of symmetric slots only.
        """

        n_body = int(n_body)

        n_body2 = n_body if n_body2 is None else int(n_body2)
        n_slots = n_body + n_body2

        transp_acc = NEG if self._exch == FERMI else IDENT
        cycl_accs = [
            NEG if self._exch == FERMI and i % 2 == 0 else IDENT
            for i in [n_body, n_body2]
        ]  # When either body is zero, this value is kinda wrong but not used.

        gens = []

        if n_body > 1:
            second_half = list(range(n_body, n_slots))
            gens.append(Perm(
                self._form_cycl(0, n_body) + second_half, cycl_accs[0]
            ))
            gens.append(Perm(
                self._form_transp(0, n_body) + second_half, transp_acc
            ))

        if n_body2 > 1:
            first_half = list(range(0, n_body))
            gens.append(Perm(
                first_half + self._form_cycl(n_body, n_slots), cycl_accs[1]
            ))
            gens.append(Perm(
                first_half + self._form_transp(n_body, n_slots), transp_acc
            ))

        self.set_symm(base, gens, valence=n_slots)

        return

    @staticmethod
    def _form_cycl(begin, end):
        """Form the pre-image for a cyclic permutation over the given range."""
        before_end = end - 1
        res = [before_end]
        res.extend(range(begin, before_end))
        return res

    @staticmethod
    def _form_transp(begin, end):
        """Form a pre-image array with the first two points transposed."""
        res = list(range(begin, end))
        res[0], res[1] = res[1], res[0]
        return res

    def _latex_vec(self, vec):
        """Get the LaTeX form of field operators.
        """

        head = r'{}^{{{}}}'.format(vec.label, self._latex_sympy(vec.indices[0]))
        indices = ', '.join(self._latex_sympy(i) for i in vec.indices[1:])
        return r'{}_{{{}}}'.format(head, indices)

    _latex_vec_mul = ' '


def parse_field_op(op: Vec, _: Term):
    """Get the operator label, character and actual indices.

    ValueError will be raised if the given operator does not satisfy the format
    for field operators.
    """

    indices = op.indices
    if len(indices) < 1 or (indices[0] != CR and indices[0] != AN):
        raise ValueError('Invalid field operator', op,
                         'expecting operator character')

    return op.label, indices[0], indices[1:]


def _compare_field_ops(
        op1: Vec, op2: Vec, term: Term,
        op_parser: FockDrudge.OP_PARSER
):
    """Compare the given field operators.

    Here we try to emulate physicists' convention as much as possible.  The
    annihilation operators are ordered in reversed direction.
    """

    label1, char1, indices1 = op_parser(op1, term)
    label2, char2, indices2 = op_parser(op2, term)

    if char1 == CR and char2 == AN:
        return True
    elif char1 == AN and char2 == CR:
        return False

    key1 = (label1, [sympy_key(i) for i in indices1])

    key2 = (label2, [sympy_key(i) for i in indices2])