'''
Write a function to convert python code to scala= code using AST trees. It should support math, variables and methods. Show a test case
'''

import ast

# Define a dictionary to map Python operators to Scala operators
OPERATOR_MAP = {
    '+': '+',
    '-': '-',
    '*': '*',
    '/': '/',
    '%': '%',
    '**': '^',
    '//': '/'
}

# Define a dictionary to map Python built-in functions to Scala functions
FUNCTION_MAP = {
    'print': 'println'
}

# Define a dictionary to store variable assignments
VARIABLES = {}

def convert_python_to_scala(python_code):
    # Parse the Python code into an AST tree
    tree = ast.parse(python_code)
    
    # Traverse the AST tree to generate the equivalent Scala code
    result = []
    for node in tree.body:
        if isinstance(node, ast.Expr):
            result.append(convert_expr(node.value))
        elif isinstance(node, ast.Assign):
            variable_name = node.targets[0].id
            variable_value = convert_expr(node.value)
            VARIABLES[variable_name] = variable_value
        elif isinstance(node, ast.FunctionDef):
            function_name = node.name
            function_args = ','.join([arg.arg for arg in node.args.args])
            function_body = '\n'.join([convert_expr(expr) for expr in node.body])
            result.append(f'def {function_name}({function_args}): \n{function_body}')
    
    # Add variable assignments to the result
    if VARIABLES:
        variables_code = '\n'.join([f'val {name} = {value}' for name, value in VARIABLES.items()])
        result.insert(0, variables_code)
    
    # Return the generated Scala code
    return '\n'.join(result)
    

def convert_expr(expr):
    if isinstance(expr, ast.Num):
        return str(expr.n)
    elif isinstance(expr, ast.Name):
        return VARIABLES.get(expr.id, expr.id)
    elif isinstance(expr, ast.BinOp):
        left = convert_expr(expr.left)
        operator = OPERATOR_MAP[expr.op.__class__.__name__]
        right = convert_expr(expr.right)
        return f'({left} {operator} {right})'
    elif isinstance(expr, ast.Call):
        function_name = FUNCTION_MAP.get(expr.func.id, expr.func.id)
        args = ','.join([convert_expr(arg) for arg in expr.args])
        return f'{function_name}({args})'
    else:
        raise NotImplementedError(f'Expression type {type(expr).__name__} not supported')

# Test case
python_code = '''
a = 5
b = 7
c = a + b
print(c)
def add_numbers(x, y):
    return x + y
'''
scala_code = convert_python_to_scala(python_code)
print(scala_code)