transformer.py

import re

import sympy
from sympy.external import import_module
from sympy.parsing.latex.errors import LaTeXParsingError

lark = import_module("lark")

if lark:
    from lark import Transformer, Token  # type: ignore
else:
    class Transformer:  # type: ignore
        def transform(self, *args):
            pass


    class Token:  # type: ignore
        pass


# noinspection PyPep8Naming,PyMethodMayBeStatic
class TransformToSymPyExpr(Transformer):
    """Returns a SymPy expression that is generated by traversing the ``lark.Tree``
    passed to the ``.transform()`` function.

    Notes
    =====

    **This class is never supposed to be used directly.**

    In order to tweak the behavior of this class, it has to be subclassed and then after
    the required modifications are made, the name of the new class should be passed to
    the :py:class:`LarkLaTeXParser` class by using the ``transformer`` argument in the
    constructor.

    Parameters
    ==========

    visit_tokens : bool, optional
        For information about what this option does, see `here
        <https://lark-parser.readthedocs.io/en/latest/visitors.html#lark.visitors.Transformer>`_.

        Note that the option must be set to ``True`` for the default parser to work.
    """

    SYMBOL = sympy.Symbol
    DIGIT = sympy.core.numbers.Integer

    def CMD_INFTY(self, tokens):
        return sympy.oo

    def GREEK_SYMBOL(self, tokens):
        # we omit the first character because it is a backslash. Also, if the variable name has "var" in it,
        # like "varphi" or "varepsilon", we remove that too
        variable_name = re.sub("var", "", tokens[1:])

        return sympy.Symbol(variable_name)

    def BASIC_SUBSCRIPTED_SYMBOL(self, tokens):
        symbol, sub = tokens.value.split("_")
        if sub.startswith("{"):
            return sympy.Symbol("%s_{%s}" % (symbol, sub[1:-1]))
        else:
            return sympy.Symbol("%s_{%s}" % (symbol, sub))

    def GREEK_SUBSCRIPTED_SYMBOL(self, tokens):
        greek_letter, sub = tokens.value.split("_")
        greek_letter = re.sub("var", "", greek_letter[1:])

        if sub.startswith("{"):
            return sympy.Symbol("%s_{%s}" % (greek_letter, sub[1:-1]))
        else:
            return sympy.Symbol("%s_{%s}" % (greek_letter, sub))

    def SYMBOL_WITH_GREEK_SUBSCRIPT(self, tokens):
        symbol, sub = tokens.value.split("_")
        if sub.startswith("{"):
            greek_letter = sub[2:-1]
            greek_letter = re.sub("var", "", greek_letter)

            return sympy.Symbol("%s_{%s}" % (symbol, greek_letter))
        else:
            greek_letter = sub[1:]
            greek_letter = re.sub("var", "", greek_letter)

            return sympy.Symbol("%s_{%s}" % (symbol, greek_letter))

    def multi_letter_symbol(self, tokens):
        return sympy.Symbol(tokens[2])

    def number(self, tokens):
        if "." in tokens[0]:
            return sympy.core.numbers.Float(tokens[0])
        else:
            return sympy.core.numbers.Integer(tokens[0])

    def latex_string(self, tokens):
        return tokens[0]

    def group_round_parentheses(self, tokens):
        return tokens[1]

    def group_square_brackets(self, tokens):
        return tokens[1]

    def group_curly_parentheses(self, tokens):
        return tokens[1]

    def eq(self, tokens):
        return sympy.Eq(tokens[0], tokens[2])

    def ne(self, tokens):
        return sympy.Ne(tokens[0], tokens[2])

    def lt(self, tokens):
        return sympy.Lt(tokens[0], tokens[2])

    def lte(self, tokens):
        return sympy.Le(tokens[0], tokens[2])

    def gt(self, tokens):
        return sympy.Gt(tokens[0], tokens[2])

    def gte(self, tokens):
        return sympy.Ge(tokens[0], tokens[2])

    def add(self, tokens):
        return sympy.Add(tokens[0], tokens[2])

    def sub(self, tokens):
        if len(tokens) == 2:
            return -tokens[1]
        elif len(tokens) == 3:
            return sympy.Add(tokens[0], -tokens[2])

    def mul(self, tokens):
        return sympy.Mul(tokens[0], tokens[2])

    def div(self, tokens):
        return sympy.Mul(tokens[0], sympy.Pow(tokens[2], -1))

    def adjacent_expressions(self, tokens):
        # Most of the time, if two expressions are next to each other, it means implicit multiplication,
        # but not always
        from sympy.physics.quantum import Bra, Ket
        if isinstance(tokens[0], Ket) and isinstance(tokens[1], Bra):
            from sympy.physics.quantum import OuterProduct
            return OuterProduct(tokens[0], tokens[1])
        elif tokens[0] == sympy.Symbol("d"):
            # If the leftmost token is a "d", then it is highly likely that this is a differential
            return tokens[0], tokens[1]
        elif isinstance(tokens[0], tuple):
            # then we have a derivative
            return sympy.Derivative(tokens[1], tokens[0][1])
        else:
            return sympy.Mul(tokens[0], tokens[1])

    def superscript(self, tokens):
        return sympy.Pow(tokens[0], tokens[2])

    def fraction(self, tokens):
        numerator = tokens[1]
        if isinstance(tokens[2], tuple):
            # we only need the variable w.r.t. which we are differentiating
            _, variable = tokens[2]

            # we will pass this information upwards
            return "derivative", variable
        else:
            denominator = tokens[2]
            return sympy.Mul(numerator, sympy.Pow(denominator, -1))

    def binomial(self, tokens):
        return sympy.binomial(tokens[1], tokens[2])

    def normal_integral(self, tokens):
        underscore_index = None
        caret_index = None

        if "_" in tokens:
            # we need to know the index because the next item in the list is the
            # arguments for the lower bound of the integral
            underscore_index = tokens.index("_")

        if "^" in tokens:
            # we need to know the index because the next item in the list is the
            # arguments for the upper bound of the integral
            caret_index = tokens.index("^")

        lower_bound = tokens[underscore_index + 1] if underscore_index else None
        upper_bound = tokens[caret_index + 1] if caret_index else None

        differential_symbol = self._extract_differential_symbol(tokens)

        if differential_symbol is None:
            raise LaTeXParsingError("Differential symbol was not found in the expression."
                                    "Valid differential symbols are \"d\", \"\\text{d}, and \"\\mathrm{d}\".")

        # else we can assume that a differential symbol was found
        differential_variable_index = tokens.index(differential_symbol) + 1
        differential_variable = tokens[differential_variable_index]

        # we can't simply do something like `if (lower_bound and not upper_bound) ...` because this would
        # evaluate to `True` if the `lower_bound` is 0 and upper bound is non-zero
        if lower_bound is not None and upper_bound is None:
            # then one was given and the other wasn't
            raise LaTeXParsingError("Lower bound for the integral was found, but upper bound was not found.")

        if upper_bound is not None and lower_bound is None:
            # then one was given and the other wasn't
            raise LaTeXParsingError("Upper bound for the integral was found, but lower bound was not found.")

        # check if any expression was given or not. If it wasn't, then set the integrand to 1.
        if underscore_index is not None and underscore_index == differential_variable_index - 3:
            # The Token at differential_variable_index - 2 should be the integrand. However, if going one more step
            # backwards after that gives us the underscore, then that means that there _was_ no integrand.
            # Example: \int^7_0 dx
            integrand = 1
        elif caret_index is not None and caret_index == differential_variable_index - 3:
            # The Token at differential_variable_index - 2 should be the integrand. However, if going one more step
            # backwards after that gives us the caret, then that means that there _was_ no integrand.
            # Example: \int_0^7 dx
            integrand = 1
        elif differential_variable_index == 2:
            # this means we have something like "\int dx", because the "\int" symbol will always be
            # at index 0 in `tokens`
            integrand = 1
        else:
            # The Token at differential_variable_index - 1 is the differential symbol itself, so we need to go one
            # more step before that.
            integrand = tokens[differential_variable_index - 2]

        if lower_bound is not None:
            # then we have a definite integral

            # we can assume that either both the lower and upper bounds are given, or
            # neither of them are
            return sympy.Integral(integrand, (differential_variable, lower_bound, upper_bound))
        else:
            # we have an indefinite integral
            return sympy.Integral(integrand, differential_variable)

    def group_curly_parentheses_int(self, tokens):
        # return signature is a tuple consisting of the expression in the numerator, along with the variable of
        # integration
        if len(tokens) == 3:
            return 1, tokens[1]
        elif len(tokens) == 4:
            return tokens[1], tokens[2]
        # there are no other possibilities

    def special_fraction(self, tokens):
        numerator, variable = tokens[1]
        denominator = tokens[2]

        # We pass the integrand, along with information about the variable of integration, upw
        return sympy.Mul(numerator, sympy.Pow(denominator, -1)), variable

    def integral_with_special_fraction(self, tokens):
        underscore_index = None
        caret_index = None

        if "_" in tokens:
            # we need to know the index because the next item in the list is the
            # arguments for the lower bound of the integral
            underscore_index = tokens.index("_")

        if "^" in tokens:
            # we need to know the index because the next item in the list is the
            # arguments for the upper bound of the integral
            caret_index = tokens.index("^")

        lower_bound = tokens[underscore_index + 1] if underscore_index else None
        upper_bound = tokens[caret_index + 1] if caret_index else None

        # we can't simply do something like `if (lower_bound and not upper_bound) ...` because this would
        # evaluate to `True` if the `lower_bound` is 0 and upper bound is non-zero
        if lower_bound is not None and upper_bound is None:
            # then one was given and the other wasn't
            raise LaTeXParsingError("Lower bound for the integral was found, but upper bound was not found.")

        if upper_bound is not None and lower_bound is None:
            # then one was given and the other wasn't
            raise LaTeXParsingError("Upper bound for the integral was found, but lower bound was not found.")

        integrand, differential_variable = tokens[-1]

        if lower_bound is not None:
            # then we have a definite integral

            # we can assume that either both the lower and upper bounds are given, or
            # neither of them are
            return sympy.Integral(integrand, (differential_variable, lower_bound, upper_bound))
        else:
            # we have an indefinite integral
            return sympy.Integral(integrand, differential_variable)

    def group_curly_parentheses_special(self, tokens):
        underscore_index = tokens.index("_")
        caret_index = tokens.index("^")

        # given the type of expressions we are parsing, we can assume that the lower limit
        # will always use braces around its arguments. This is because we don't support
        # converting unconstrained sums into SymPy expressions.

        # first we isolate the bottom limit
        left_brace_index = tokens.index("{", underscore_index)
        right_brace_index = tokens.index("}", underscore_index)

        bottom_limit = tokens[left_brace_index + 1: right_brace_index]

        # next, we isolate the upper limit
        top_limit = tokens[caret_index + 1:]

        # the code below will be useful for supporting things like `\sum_{n = 0}^{n = 5} n^2`
        # if "{" in top_limit:
        #     left_brace_index = tokens.index("{", caret_index)
        #     if left_brace_index != -1:
        #         # then there's a left brace in the string, and we need to find the closing right brace
        #         right_brace_index = tokens.index("}", caret_index)
        #         top_limit = tokens[left_brace_index + 1: right_brace_index]

        # print(f"top  limit = {top_limit}")

        index_variable = bottom_limit[0]
        lower_limit = bottom_limit[-1]
        upper_limit = top_limit[0]  # for now, the index will always be 0

        # print(f"return value = ({index_variable}, {lower_limit}, {upper_limit})")

        return index_variable, lower_limit, upper_limit

    def summation(self, tokens):
        return sympy.Sum(tokens[2], tokens[1])

    def product(self, tokens):
        return sympy.Product(tokens[2], tokens[1])

    def limit_dir_expr(self, tokens):
        caret_index = tokens.index("^")

        if "{" in tokens:
            left_curly_brace_index = tokens.index("{", caret_index)
            direction = tokens[left_curly_brace_index + 1]
        else:
            direction = tokens[caret_index + 1]

        if direction == "+":
            return tokens[0], "+"
        elif direction == "-":
            return tokens[0], "-"
        else:
            return tokens[0], "+-"

    def group_curly_parentheses_lim(self, tokens):
        limit_variable = tokens[1]
        if isinstance(tokens[3], tuple):
            destination, direction = tokens[3]
        else:
            destination = tokens[3]
            direction = "+-"

        return limit_variable, destination, direction

    def limit(self, tokens):
        limit_variable, destination, direction = tokens[2]

        return sympy.Limit(tokens[-1], limit_variable, destination, direction)

    def differential(self, tokens):
        return tokens[1]

    def derivative(self, tokens):
        return sympy.Derivative(tokens[-1], tokens[5])

    def list_of_expressions(self, tokens):
        if len(tokens) == 1:
            # we return it verbatim because the function_applied node expects
            # a list
            return tokens
        else:
            def remove_tokens(args):
                if isinstance(args, Token):
                    if args.type != "COMMA":
                        # An unexpected token was encountered
                        raise LaTeXParsingError("A comma token was expected, but some other token was encountered.")
                    return False
                return True

            return filter(remove_tokens, tokens)

    def function_applied(self, tokens):
        return sympy.Function(tokens[0])(*tokens[2])

    def min(self, tokens):
        return sympy.Min(*tokens[2])

    def max(self, tokens):
        return sympy.Max(*tokens[2])

    def bra(self, tokens):
        from sympy.physics.quantum import Bra
        return Bra(tokens[1])

    def ket(self, tokens):
        from sympy.physics.quantum import Ket
        return Ket(tokens[1])

    def inner_product(self, tokens):
        from sympy.physics.quantum import Bra, Ket, InnerProduct
        return InnerProduct(Bra(tokens[1]), Ket(tokens[3]))

    def sin(self, tokens):
        return sympy.sin(tokens[1])

    def cos(self, tokens):
        return sympy.cos(tokens[1])

    def tan(self, tokens):
        return sympy.tan(tokens[1])

    def csc(self, tokens):
        return sympy.csc(tokens[1])

    def sec(self, tokens):
        return sympy.sec(tokens[1])

    def cot(self, tokens):
        return sympy.cot(tokens[1])

    def sin_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.asin(tokens[-1])
        else:
            return sympy.Pow(sympy.sin(tokens[-1]), exponent)

    def cos_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.acos(tokens[-1])
        else:
            return sympy.Pow(sympy.cos(tokens[-1]), exponent)

    def tan_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.atan(tokens[-1])
        else:
            return sympy.Pow(sympy.tan(tokens[-1]), exponent)

    def csc_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.acsc(tokens[-1])
        else:
            return sympy.Pow(sympy.csc(tokens[-1]), exponent)

    def sec_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.asec(tokens[-1])
        else:
            return sympy.Pow(sympy.sec(tokens[-1]), exponent)

    def cot_power(self, tokens):
        exponent = tokens[2]
        if exponent == -1:
            return sympy.acot(tokens[-1])
        else:
            return sympy.Pow(sympy.cot(tokens[-1]), exponent)

    def arcsin(self, tokens):
        return sympy.asin(tokens[1])

    def arccos(self, tokens):
        return sympy.acos(tokens[1])

    def arctan(self, tokens):
        return sympy.atan(tokens[1])

    def arccsc(self, tokens):
        return sympy.acsc(tokens[1])

    def arcsec(self, tokens):
        return sympy.asec(tokens[1])

    def arccot(self, tokens):
        return sympy.acot(tokens[1])

    def sinh(self, tokens):
        return sympy.sinh(tokens[1])

    def cosh(self, tokens):
        return sympy.cosh(tokens[1])

    def tanh(self, tokens):
        return sympy.tanh(tokens[1])

    def asinh(self, tokens):
        return sympy.asinh(tokens[1])

    def acosh(self, tokens):
        return sympy.acosh(tokens[1])

    def atanh(self, tokens):
        return sympy.atanh(tokens[1])

    def abs(self, tokens):
        return sympy.Abs(tokens[1])

    def floor(self, tokens):
        return sympy.floor(tokens[1])

    def ceil(self, tokens):
        return sympy.ceiling(tokens[1])

    def factorial(self, tokens):
        return sympy.factorial(tokens[0])

    def conjugate(self, tokens):
        return sympy.conjugate(tokens[1])

    def square_root(self, tokens):
        if len(tokens) == 2:
            # then there was no square bracket argument
            return sympy.sqrt(tokens[1])
        elif len(tokens) == 3:
            # then there _was_ a square bracket argument
            return sympy.root(tokens[2], tokens[1])

    def exponential(self, tokens):
        return sympy.exp(tokens[1])

    def log(self, tokens):
        if tokens[0].type == "FUNC_LG":
            # we don't need to check if there's an underscore or not because having one
            # in this case would be meaningless
            # TODO: ANTLR refers to ISO 80000-2:2019. should we keep base 10 or base 2?
            return sympy.log(tokens[1], 10)
        elif tokens[0].type == "FUNC_LN":
            return sympy.log(tokens[1])
        elif tokens[0].type == "FUNC_LOG":
            # we check if a base was specified or not
            if "_" in tokens:
                # then a base was specified
                return sympy.log(tokens[3], tokens[2])
            else:
                # a base was not specified
                return sympy.log(tokens[1])

    def _extract_differential_symbol(self, s: str):
        differential_symbols = {"d", r"\text{d}", r"\mathrm{d}"}

        differential_symbol = next((symbol for symbol in differential_symbols if symbol in s), None)

        return differential_symbol