from sympy import sympify, lambdify, Symbol
from sympy.logic.boolalg import Boolean
from pandas import DataFrame
import pandas as pd
from typing import Union
import numpy as np
from .object import BaseException

class MathsException(BaseException):
    pass

class ExpressionCache:
    """用于缓存表达式和对应矢量化函数的类"""
    _cache = {}

    @classmethod
    def get_calculator(cls, expression, safe_keys = []):
        """获取或创建表达式的矢量化计算器"""
        if expression not in cls._cache:
            cls._cache[expression] = VectorizedExpressionCalculator(expression, safe_keys)
        return cls._cache[expression]


class VectorizedExpressionCalculator:
    """矢量化表达式计算类"""
    def __init__(self, expression, safe_keys = []):
        # 将表达式字符串转换为 SymPy 表达式
        self.expression = sympify(expression, locals={k: Symbol(k) for k in safe_keys})
        # 将 SymPy 表达式转换为矢量化函数（支持 NumPy/Pandas 运算）
        self.variables = sorted(self.expression.free_symbols, key=lambda x: str(x)) # 这里要用sorted很关键，因为set是没有顺序的，会导致值计算错误
        self.vectorized_func = lambdify(
            self.variables,
            self.expression,
            modules=[{
                'max': np.maximum,
                'min': np.minimum,
                '&': np.logical_and,
                '|': np.logical_or,
                '==': np.equal,
                '>': np.greater,
                '<': np.less,
                '>=': np.greater_equal,
                '<=': np.less_equal
            }, 'numpy']
        )

    def calculate_dataframe(self, df):
        # 确保 DataFrame 中包含表达式需要的所有变量
        required_columns = [str(x) for x in self.variables]
        missing_columns = set(required_columns) - set(df.columns)
        if missing_columns:
            raise ValueError(f"Missing required columns: {missing_columns}")
        
        # 使用矢量化函数计算
        return self.vectorized_func(*(df[col] for col in required_columns))


def calculate_expression(expression, df):
    """接口函数，直接传入表达式和 DataFrame 进行计算"""
    calculator = ExpressionCache.get_calculator(expression)
    return calculator.calculate_dataframe(df)

def cal_exp(exp: str, values: Union[dict, DataFrame], safe_keys: list[str]=[]):
    try:
        single_run = isinstance(values, dict)
        if single_run:
            values = pd.DataFrame([values])
        calculator = ExpressionCache.get_calculator(exp, safe_keys)
        series = pd.Series(calculator.calculate_dataframe(values))
        return series if not single_run else (series.iloc[0] if len(series) else None)
    except Exception as e:
        raise MathsException(f"表达式{exp}计算错误：{str(e)}", data={'values': values})

# t = cal_exp('a+b+c', {'a': 1, 'b': 2, 'c': 3})
# print(t)