evaluator.py

"""
Evaluator for AST nodes generated by Parser.
Evaluates math expressions with variable and function support.
"""

import math
from typing import Dict, Any, Union
from expression_parser.parser import ASTNode, NumberNode, VariableNode, UnaryOpNode, BinaryOpNode, FunctionCallNode


class EvaluationError(Exception):
    pass


class Evaluator:
    BUILTIN_FUNCTIONS = {
        "sin": math.sin,
        "cos": math.cos,
        "tan": math.tan,
        "asin": math.asin,
        "acos": math.acos,
        "atan": math.atan,
        "sqrt": math.sqrt,
        "abs": abs,
        "log": math.log,
        "exp": math.exp,
        "floor": math.floor,
        "ceil": math.ceil,
        "min": min,
        "max": max,
    }

    BUILTIN_CONSTANTS = {
        "pi": math.pi,
        "e": math.e,
        "tau": math.tau,
    }

    def __init__(self, variables: Dict[str, Union[int, float]] = None):
        self.variables = self.BUILTIN_CONSTANTS.copy()
        if variables:
            self.variables.update(variables)

    def evaluate(self, node: ASTNode) -> Union[int, float]:
        if isinstance(node, NumberNode):
            return node.value

        if isinstance(node, VariableNode):
            if node.name in self.variables:
                return self.variables[node.name]
            raise EvaluationError(f"Undefined variable or constant: '{node.name}'")

        if isinstance(node, UnaryOpNode):
            val = self.evaluate(node.operand)
            if node.operator == '+':
                return +val
            elif node.operator == '-':
                return -val
            raise EvaluationError(f"Unsupported unary operator: '{node.operator}'")

        if isinstance(node, BinaryOpNode):
            left_val = self.evaluate(node.left)
            right_val = self.evaluate(node.right)

            if node.operator == '+':
                return left_val + right_val
            elif node.operator == '-':
                return left_val - right_val
            elif node.operator == '*':
                return left_val * right_val
            elif node.operator == '/':
                if right_val == 0:
                    raise EvaluationError("Division by zero")
                return left_val / right_val
            elif node.operator == '%':
                if right_val == 0:
                    raise EvaluationError("Modulo by zero")
                return left_val % right_val
            elif node.operator in ('^', '**'):
                try:
                    res = left_val ** right_val
                    if isinstance(res, complex):
                        raise EvaluationError("Complex numbers result not supported in real domain")
                    return res
                except OverflowError:
                    raise EvaluationError("Numerical overflow in exponentiation")

            raise EvaluationError(f"Unsupported binary operator: '{node.operator}'")

        if isinstance(node, FunctionCallNode):
            fn_name = node.name.lower()
            if fn_name not in self.BUILTIN_FUNCTIONS:
                raise EvaluationError(f"Unknown or unsupported function: '{node.name}'")

            args = [self.evaluate(arg) for arg in node.args]
            func = self.BUILTIN_FUNCTIONS[fn_name]

            try:
                result = func(*args)
                return result
            except TypeError as e:
                raise EvaluationError(f"Invalid arguments for function '{node.name}': {e}")
            except ValueError as e:
                raise EvaluationError(f"Mathematical domain error in function '{node.name}': {e}")

        raise EvaluationError(f"Unknown AST node type: {type(node).__name__}")


def evaluate_expression(expression: str, variables: Dict[str, Union[int, float]] = None) -> Union[int, float]:
    """Helper function to tokenize, parse, and evaluate a math expression string."""
    from expression_parser.tokenizer import Tokenizer
    from expression_parser.parser import Parser

    tokenizer = Tokenizer(expression)
    tokens = tokenizer.tokenize()
    parser = Parser(tokens)
    ast = parser.parse()
    evaluator = Evaluator(variables)
    return evaluator.evaluate(ast)