import math

from lark import Transformer, v_args

@v_args(inline=True)
class ExpressionTransformer(Transformer):


    def add(self, left, right) -> str:
        result = f"{left} + {right}"
        return result

    def sub(self, left, right) -> str:
        result = f"{left} - {right}"
        return result

    def mul(self, left, right) -> str:
        result = f"{left} * {right}"
        return result

    def div(self, left, right) -> str:
        result = f"{left} / {right}"
        return result

    def exponent(self, left, right) -> str:
        result = f"{left} ** {right}"
        return result

    def number(self, n) -> str:
        return f"{n}"

    def neg(self, expr) -> str:
        return f"-{expr}"

    def var(self, id:str) -> str:
        if id == 'x':
            return id
        elif hasattr(math, id):
            return f"math.{id}"
        else:
            raise ValueError(f"Variable must be 'x' one of constants in math module, got {id}")

    def function(self, name, argv) -> str:
        if hasattr(math, name):
            return f"math.{name}({argv})"
        else:
            raise ValueError(f"'{name}' is not a valid function name in math module")

    def argv(self, *arg) -> str:
        return f"{', '.join(arg)}"