def compile_drs(src, filename):
    """Compile the drudge script."""

    root = ast.parse(src, filename=filename)

    for i in _FIXERS:
        root = i.visit(root)
        continue

    return compile(root, filename, mode='exec') class DrsEnv(dict):
    """The global scope for drudge script execution.
    """

    def __init__(self, dr, specials=None):
        """Initialize the scope."""

        super().__init__()

        self._drudge = dr
        path = [(dr.names, frozenset())]
        self._path = path

        if specials is not None:
            path.append((specials, frozenset()))

        path.append((dr, frozenset()))
        import drudge
        path.append((drudge, frozenset()))

        try:
            import gristmill
        except ModuleNotFoundError:
            pass
        else:
            path.append((gristmill, frozenset()))

        import sympy
        path.append((sympy, frozenset([
            'N'
        ])))

        import builtins
        path.append((builtins, frozenset()))

    def __missing__(self, key: str):
        """Get the missing name.

        The missing name will be returned.  The result is not cached to avoid
        surprise.
        """

        if key.startswith('__') and key.endswith('__'):
            raise KeyError(key)

        for entry, excl in self._path:
            if hasattr(entry, key) and key not in excl:
                resolv = getattr(entry, key)
                break
            else:
                continue
        else:
            resolv = DrsSymbol(self._drudge, key)
        return resolv class DrsSymbol(_Definable, Symbol):
    """Symbols used in drudge scripts.

    The drudge symbol needs to behave as similar to the actual symbol as
    possible, because it is possible that they are used for keys in
    dictionaries.
    """

    __slots__ = [
        '_drudge',
        '_orig'
    ]

    def __new__(cls, drudge, name):
        """Create a symbol object."""
        symb = super().__new__(cls, name)
        return symb

    def __init__(self, drudge, name):
        """Initialize the symbol object."""
        self._drudge = drudge
        self._orig = Symbol(name)

    def __eq__(self, other):
        """Make equality comparison."""
        return self._orig == other

    def __hash__(self):
        """Compute the hash."""
        return hash(self._orig)

    def _hashable_content(self):
        """Hashable content for SymPy usages."""
        return self._orig._hashable_content()

    @classmethod
    def class_key(cls):
        return Symbol.class_key()

    def __getitem__(self, indices):
        """Index the given symbol.

        In drudge scripts, all symbols are by itself indexed bases.
        """
        base = IndexedBase(self._orig)
        if isinstance(indices, collections.Sequence):
            return DrsIndexed(self._drudge, base, *indices)
        else:
            return DrsIndexed(self._drudge, base, indices)

    def __iter__(self):
        """Disable iterability of the symbol.

        Or a default implementation from ``__getitem__`` will be used,
        which makes the symbols unable to be used as subscripts for indexed
        objects.
        """
        raise TypeError('Drudge script symbol cannot be iterated over.')

    def __call__(self, *args, **kwargs):
        """Make a call to a tensor method."""
        name = self.name

        if len(args) == 0:
            raise NameError('Undefined function', name)
        else:
            target = args[0]
            rest = args[1:]

            err = NameError('Invalid method', name, 'for', str(type(target)))
            if not hasattr(target, name):
                # It is possible that a tensor method is tried to be called on
                # an input.
                if isinstance(target, ATerms):
                    tensor = self._drudge.sum(target)
                    if hasattr(tensor, name):
                        target = tensor
                    else:
                        raise err
                else:
                    raise err

            meth = getattr(target, name)
            if inspect.ismethod(meth):
                # It has a caveat that static methods might not be able to be
                # called.
                return meth(*rest, **kwargs)
            elif len(rest) == 0 and len(kwargs) == 0:
                return meth
            else:
                raise err

    # Pickle support.

    def __getnewargs__(self):
        """Get the arguments for __new__."""
        return None, self.name

    def __getstate__(self):
        """Get the state for pickling."""
        return None

    def __setstate__(self, state):
        """Set the state according to pickled content."""
        from .drudge import current_drudge
        if current_drudge is None:
            raise ValueError(_PICKLE_ENV_ERR)
        self.__init__(current_drudge, self.name)

    # Better error reporting.
    def __getattr__(self, item):
        """Report undefined attributes."""
        raise AttributeError(
            'Invalid operation `{}` on Drudge symbol `{}`'.format(
                item, self.name
            )
        ) class Report:
    """Simple report for output drudge results.

    This class helps to write symbolic results to files for batch processing
    jobs.  It is not recommended to be used directly.  Users should use the
    method provided by drudge class instead in ``with`` statements.
    """

    def __init__(self, filename: str, title: str, drudge: str):
        """Initialize the report object."""

        self._filename = filename
        self._drudge = drudge

        filename_parts = filename.split('.')
        if len(filename_parts) < 2:
            raise ValueError(
                'Invalid filename, unable to determine type', filename
            )
        ext = filename_parts[-1].lower()
        if ext not in {'html', 'tex', 'pdf'}:
            raise ValueError('Invalid extension', ext, 'in', filename)
        self._ext = ext
        self._basename = '.'.join(filename_parts[:-1])

        self._sects = []
        self._ctx = {
            'title': title,
            'sects': self._sects
        }

    def add(
            self, title=None, content=None, description=None,
            env='[', **kwargs
    ):
        r"""Add a section to the result.

        Parameters
        ----------

        title

            The title of the equation.  It will be used as a section header.
            When it is given as a None, the section header will not be added.

        content

            The actual tensor or tensor definition to be printed.  It can be
            given as a None to skip any equation rendering.

        description

            A verbal description of the content.  It will be typeset before the
            actual equation as normal text.  A None value will cause it to be
            suppressed.

        env

            The environment to put the equation in.  A value of ``'['`` will use
            ``\[`` and ``\]`` as the deliminator of the math environment.  Other
            values will be put inside the common ``\begin{}`` and ``\end{}``
            tags of LaTeX.

        kwargs

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

        Note
        ----

        **For long equations in LaTeX environments,** normally ``env='align'``
        and ``sep_lines=True`` can be set to allow each term to occupy separate
        lines, automatic page break will be inserted, or ``env='dmath'`` and
        ``sep_lines=False`` can be used to use ``breqn`` package to
        automatically flow the terms.

        """

        if env == '[':
            opening, closing = r'\[', r'\]'
        else:
            opening, closing = [
                r'\{}{{{}}}'.format(i, env)
                for i in ['begin', 'end']
            ]

        if content is None:
            expr = None
        elif hasattr(content, 'latex'):
            expr = content.latex(**kwargs)
        elif isinstance(content, ATerms):
            expr = self._drudge.sum(content).latex(**kwargs)
        else:
            # Try to use raw SymPy LaTeX printing.
            expr = ScalarLatexPrinter().doprint(content)

        self._sects.append({
            'title': title,
            'description': description,
            'expr': expr,
            'opening': opening,
            'closing': closing
        })

    def write(self):
        """Write the report.

        Note that this method also closes the output file.
        """

        env = Environment(
            loader=PackageLoader('drudge'),
            lstrip_blocks=True, trim_blocks=True
        )

        if self._ext == 'html':
            templ_name = 'report.html'
            filename = self._filename
        elif self._ext == 'tex' or self._ext == 'pdf':
            templ_name = 'report.tex'
            filename = self._basename + '.tex'
        else:
            assert False

        templ = env.get_template(templ_name)

        with open(filename, 'w') as fp:
            templ.stream(self._ctx).dump(fp)

        if self._ext == 'pdf':
            if shutil.which(_PDFLATEX) is not None:
                stat = subprocess.run(
                    [_PDFLATEX, filename],
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                # Do not crash program only because LaTeX does not compile.
                if stat.returncode != 0:
                    err_msg = '{} failed for {}.  Error: \n{}\n{}'.format(
                        _PDFLATEX, filename, stat.stdout, stat.stderr
                    )
                    warnings.warn(err_msg)
            else:
                warnings.warn('{} cannot be found.'.format(_PDFLATEX)) class ScalarLatexPrinter(LatexPrinter):
    """Specialized LaTeX printers for usage in drudge.

    Basically this class tries to fix some problems with using the original
    LaTeX printer from SymPy in common drudge tasks.

    Specifically, for indexed objects, if the base already contains a subscript,
    it will be raised into a superscript wrapped inside a pair of parenthesis.
    """

    def _print_Indexed(self, expr):
        base = self._print(expr.base)
        match = re.match(r'(.*)_\{(.*)\}', base)
        if match and len(match.group(2)) > 0:
            base = ''.join([
                match.group(1), '^{(', match.group(2), ')}'
            ])
        if base.startswith('_'):
            base = base[1:]

        indices = ','.join(self._print(i) for i in expr.indices)
        return ''.join([
            base, '_{', indices, '}'
        ]) _UNITY = Integer(1)
_NEG_UNITY = Integer(-1)
_NAUGHT = Integer(0)
_WRAPPER_BASE = IndexedBase(
    'internalWrapper', shape=('internalShape',)
)
_TREATED_PLACEHOLDER = Symbol('internalTreatedPlaceholder')
_COMMUTATIVE = 1
_NON_COMMUTATIVE = 0
_GRAD_SYMB_FMT = 'internalGradient{tag}Placeholder'
_GRAD_SYMB = Symbol(_GRAD_SYMB_FMT.format(tag=''))
_GRAD_CONJ_SYMB = Symbol(_GRAD_SYMB_FMT.format(tag='Conj'))
_GRAD_REAL_SYMB = Symbol(
    _GRAD_SYMB_FMT.format(tag='Real'), real=True
)
_GRAD_WILD_FMT = 'InternalWildSymbol{}'
class Range:
class ATerms(abc.ABC):
class Terms(ATerms):
class Vec(ATerms, CantSympify):
class Term(ATerms):
    def __init__(self, label, lower=None, upper=None):
    def label(self):
    def lower(self):
    def upper(self):
    def size(self):
    def bounded(self):
    def args(self):
    def __hash__(self):
    def __eq__(self, other):
    def __repr__(self):
    def __str__(self):
    def sort_key(self):
    def replace_label(self, new_label):
    def __lt__(self, other):
    def __getitem__(self, item):
    def map(self, func):
    def terms(self) -> typing.List['Term']:
    def __mul__(self, other):
    def __rmul__(self, other):
    def _mul(left_terms, right_terms):
    def __add__(self, other):
    def __radd__(self, other):
    def __sub__(self, other):
    def __rsub__(self, other):
    def __neg__(self):
    def _add(left_terms, right_terms):
    def _neg_terms(terms: typing.Iterable['Term']):
    def __truediv__(self, other):
    def __rtruediv__(self, other):
    def __or__(self, other):
    def __ror__(self, other):
    def __init__(self, terms: typing.Iterable['Term']):
    def terms(self):
def parse_terms(obj) -> typing.List['Term']:
    def __init__(self, label, indices=()):
    def label(self):
    def base(self):
    def indices(self):
    def __getitem__(self, item):
    def __repr__(self):
    def __str__(self):
    def __hash__(self):
    def __eq__(self, other):
    def sort_key(self):
    def map(self, func):
    def terms(self):
    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,
    ):
    def sums(self):
    def amp(self) -> Expr:
    def vecs(self):
    def is_scalar(self):
    def args(self):
    def __hash__(self):
    def __eq__(self, other):
    def __repr__(self):
    def __str__(self):
    def sort_key(self):
    def terms(self):
    def scale(self, factor):
    def mul_term(self, other, dumms=None, excl=None):
    def comm_term(self, other, dumms=None, excl=None):
    def reconcile_dumms(self, other, dumms, excl):
    def exprs(self):
    def free_vars(self):
    def dumms(self):
    def amp_factors(self):
    def get_amp_factors(self, *special_symbs, monom_only=True):
    def map(
            self, func=lambda x: x, sums=None, amp=None, vecs=None,
            skip_vecs=False, skip_ranges=False
    ):
    def subst(self, substs, sums=None, amp=None, vecs=None, purge_sums=False):
    def reset_dumms(self, dumms, dummbegs=None, excl=None, add_substs=None):
    def reset_sums(sums, dumms, dummbegs=None, excl=None):
    def simplify_deltas(self, resolvers):
    def simplify_trivial_sums(self):
    def expand(self):
    def canon(self, symms=None, vec_colour=None):
        def replace_to_treats(expr: Expr):
    def canon4normal(self, symms):
    def has_base(self, base):
def subst_vec_term(
        term: Term, lhs: typing.Tuple[Vec], rhs_terms: typing.List[Term],
        dumms, dummbegs, excl
):
def subst_factor_term(
        term: Term, lhs, rhs_terms: typing.List[Term], dumms, dummbegs, excl,
        full_simplify=True
):
        def query_func(expr):
        def replace_func(_):
        def query_func(expr):
        def replace_func(expr):
    def decouple_pow(base, e):
def _match_indices(target, expr):
def _update_substs(substs, new):
def _prepare_subst_states(rhs_terms, substs, dumms, dummbegs, excl):
def rewrite_term(
        term: Term, vecs: typing.Sequence[Vec], new_amp: Expr
) -> typing.Tuple[typing.Optional[Term], Term]:
def expand_sums_term(
        term: Term, range_: Range, expander: Sum_expander,
        exts=None, conv_accs=None
):
    def check_expr(expr: Expr):
def sum_term(sum_args, summand, predicate=None) -> typing.List[Term]:
def _parse_sums(args):
def _cat_sums(sums1, sums2):
def einst_term(term: Term, resolvers) -> typing.Tuple[
    def proc_indexed(*args):
def parse_term(term):
def simplify_deltas_in_expr(sums_dict, amp, resolvers):
def compose_simplified_delta(amp, new_substs, substs, sums_dict, resolvers):
def proc_delta(arg1, arg2, sums_dict, resolvers):
def _proc_delta_in_amp(sums_dict, resolvers, substs, *args):
def _try_simpl_unresolved_deltas(amp: Expr):
def _parse_factor_symb(expr: Expr):
def simplify_amp_sums_term(term: Term, simplifiers, excl_bases, resolvers):
def diff_term(term: Term, variable, real, wirtinger_conj):
def try_resolve_range(i, sums_dict, resolvers): def ensure_symb(obj, role=''):
    """Sympify the given object into a symbol."""
    return ensure_sympify(obj, role, Symbol) class BCastVar:
    """Automatically broadcast variables.

    This class is a shallow encapsulation of a variable and its broadcast
    into the spark context.  The variable can be redistributed automatically
    after any change.

    """

    __slots__ = [
        '_ctx',
        '_var',
        '_bcast'
    ]

    def __init__(self, ctx: SparkContext, var):
        """Initialize the broadcast variable."""
        self._ctx = ctx
        self._var = var
        self._bcast = None

    @property
    def var(self):
        """Get the variable to mutate."""
        self._bcast = None
        return self._var

    @property
    def ro(self):
        """Get the variable, read-only.

        Note that this function only prevents the redistribution of the
        variable.  It cannot force the variable not be mutated.
        """
        return self._var

    @property
    def bcast(self):
        """Get the broadcast variable."""
        if self._bcast is None:
            self._bcast = self._ctx.broadcast(self._var)
        return self._bcast def nest_bind(rdd: RDD, func, full_balance=True):
    """Nest the flat map of the given function.

    When an entry no longer need processing, None can be returned by the call
    back function.

    """

    if full_balance:
        return _nest_bind_full_balance(rdd, func)
    else:
        return _nest_bind_no_balance(rdd, func) 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) class SymbResolver:
    """Resolver based on symbols.

    It can be given an iterable of range/symbols pairs telling that the symbols
    are associated with the key ranges.  In strict mode, only the given symbols
    can be resolved to be in the given range.  In non-strict mode, all
    expressions having one of the symbols will be resolved to be in the range of
    the symbol.

    Behaviour is undefined if we have non-disjoint symbol sets for different
    ranges or when we have expression containing symbols for multiple known
    ranges.

    """

    __slots__ = [
        '_known',
        '_strict'
    ]

    def __init__(self, range_symbs, strict):
        """Initialize the resolver."""

        known = {}
        self._known = known
        for range_, dumms in range_symbs:
            for i in dumms:
                known[i] = range_
                continue
            continue

        self._strict = strict

    def __call__(self, expr: Expr):
        """Try to resolve an expression."""

        known = self._known
        if self._strict:
            if expr in known:
                return known[expr]
        else:
            for i in expr.atoms(Symbol):
                if i in known:
                    return known[i]
                continue

        return None
import collections
import contextlib
import functools
import inspect
import operator
import pickle
import sys
import types
import typing
import warnings
from collections.abc import Iterable, Sequence
from IPython.display import Math, display
from pyspark import RDD, SparkContext
from sympy import (
    IndexedBase, Symbol, Indexed, Wild, symbols, sympify, Expr, Add, Matrix, Mul
)
from sympy.concrete.summations import eval_sum_symbolic
from .canonpy import Perm, Group
from .drs import compile_drs, DrsEnv, DrsSymbol
from .report import Report, ScalarLatexPrinter
from .term import (
    Range, sum_term, Term, Vec, subst_factor_term, subst_vec_term, parse_terms,
    einst_term, diff_term, try_resolve_range, rewrite_term, Sum_expander,
    expand_sums_term, ATerms, simplify_amp_sums_term
)
from .utils import (
    ensure_symb, BCastVar, nest_bind, sympy_key, SymbResolver
)

        .. 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: