from sympy import *
from matplotlib.pyplot import *
from convert_formula import *

def basic_caculate(formula):
    return eval(convert_formula(formula).replace('*j', 'j'))

def solve_formula(formula, independent_variable=1):
    """
    求解方程
    :param formula: 方程列表
    :param independent_variable: 变量个数
    :return: 方程解
    """
    llist = []
    x, y, z = symbols('x y z')
    for i in formula:
        l = i.split('=')
        left = l[0] + '-(' + l[1] + ')'
        llist.append(sympify(convert_formula(left)))
    if independent_variable == 1:
        return solve(llist[0], x)
    elif independent_variable == 2:
        return solve((llist[0], llist[1]), (x, y))
    elif independent_variable == 3:
        return solve((llist[0], llist[1], llist[2]), (x, y, z))

def plot_function(formula, length):
    """
    绘制函数图像
    :param formula: 函数表达式
    :param length: 自变量取值范围
    :return: 图像
    """
    x = np.linspace(length[0], length[1], 1000)
    independent, formula = split_function(convert_formula(formula))
    y = eval('lambda %s: %s' % (independent, formula))  # 使用eval直接生成lambda函数
    plot(x, y(x))
    show()

def plot_implicit_function(formula, x_range=(-5, 5), y_range=(-5, 5), density=200):
    '''
    绘制隐函数图像（如 x²+y²=1）
    :param formula: 隐函数表达式（字符串，如 "x**2 + y**2 - 1"）
    :param x_range: x轴取值范围（元组，如 (-5,5)）
    :param y_range: y轴取值范围（元组，如 (-5,5)）
    :param density: 网格点密度（数值越大越平滑）
    '''
    # 生成x和y的网格点
    x = np.linspace(x_range[0], x_range[1], density)
    y = np.linspace(y_range[0], y_range[1], density)
    X, Y = np.meshgrid(x, y)
    # 转换表达式（复用现有convert_formula处理符号）
    formula = convert_formula(formula)
    l = formula.split('=')
    formula = l[0]+'-('+l[1]+')'
    # 计算网格点的函数值F(x,y)
    try:
        F = eval(formula.replace('x', 'X').replace('y', 'Y'))  # 将表达式中的x/y替换为二维网格X/Y
    except Exception as e:
        raise ValueError(f"表达式错误：{e}")
    # 绘制F(x,y)=0的等高线（即隐函数图像）
    contour(X, Y, F, levels=[0], colors='blue')
    grid(True)
    show()

def UI():
    while True:
        command = input('请选择功能：1.基本计算 2.解方程 3.绘制函数图像 4.绘制平面解析几何图像 5.退出\n')
        if command == '1':
            formula = input('请输入算式：')
            print(basic_caculate(formula))
        elif command == '2':
            independent_variable = input('请输入方程未知量个数：')
            formula = input('请输入方程组（方程之间用英文逗号分隔）：')
            print(solve_formula(formula.split(','), int(independent_variable)))
        elif command == '3':
            formula = input('请输入函数表达式（如 f(x)=x^2）：')
            length = input('请输入自变量取值范围（区间形式，不支持半开半闭区间）：')
            plot_function(formula, eval(length))
        elif command == '4':
            formula = input('请输入隐函数表达式（如 x^2+y^2=1）：')
            plot_implicit_function(formula)
        elif command == '5':
            print('谢谢使用！')
            break
        else:
            print('输入错误，请重新输入！')

if __name__ == '__main__':
    UI()