from math import *
import ast
import sys
import matplotlib.pyplot as plt
from io import StringIO

def execute_code(code_str, **args):
    code_str = str(code_str)
    "executes any sort of python code"
    stdout = sys.stdout
    sys.stdout = capture_output = StringIO()
    
    try:
        code_ast = ast.parse(code_str, mode='exec')
        exec(compile(code_ast, filename="<ast>", mode="exec"), globals(), locals())

        for obj in list(locals().values()):
            if isinstance(obj, plt.Figure):
                obj.show()
    except Exception as e:
        print("An error occurred while executing the code:")
        print(e)
    sys.stdout = stdout
    output = capture_output.getvalue()
    capture_output.close()

    return {"llm_output": str(output), "display": None, "type": "console"}
    
def calculator(equation, **args):
    equation = str(equation)
    "This function simply solves any equation using math library"
    equation = equation.replace("x", "*")
    equation = equation.replace("^", "**")
    try:
        result = eval(equation)
    except Exception as e:
        result = "Please format the equation correctly. It is not formatted correctly."
    return {"llm_output": result, "display": None, "type": "text"}

def none(some_value, **args):
    "This function simply returns True"
    return {"llm_output": "True", "display": None, "type": "text"}
from math import *

def solve_expression(expression):
    lines = expression.split('\n')
    variables = {}
    
    for line in lines:
        parts = line.split('=')
        if len(parts) != 2:
            print(f"Ignoring invalid expression: '{line.strip()}'")
            continue

        var_name = parts[0].strip()
        expr = parts[1].strip()
        
        # Replace variable names with their values
        for var, val in variables.items():
            expr = expr.replace(var, str(val))

        # Evaluate the expression using symbols from math module
        variables[var_name] = eval(expr)

    return variables

# Example usage:
expression = """
x = 99
y = 99 - x
invalid_expression_without_equal_sign
"""