import numpy as np
from typing import Dict, Any, Optional, List, Union
from ..utility_evaluator import UtilityEvaluator


class QueryError(UtilityEvaluator):
    """
    查询误差效用评估器。
    计算原始数据和扰动数据在指定聚合查询上的相对误差。
    查询类型支持 'mean', 'sum', 'count', 'variance' 等。

    参数:
        columns: List[int] 或 int，指定要查询的列索引（默认 None，表示所有列）
        query_type: str，查询类型，可选 'mean', 'sum', 'count', 'var'（默认 'mean'）
        epsilon: float，防止除零的小常数（默认 1e-10）
    """

    def __init__(self,
                 columns: Optional[Union[int, List[int]]] = None,
                 query_type: str = 'mean',
                 epsilon: float = 1e-10):
        self.columns = columns
        self.query_type = query_type
        self.epsilon = epsilon

    def evaluate(self,
                 original: np.ndarray,
                 perturbed: np.ndarray,
                 task: Optional[Dict[str, Any]] = None) -> float:
        """
        计算查询的相对误差。
        返回 1 - 相对误差（归一化到 [0,1]，越大越好）。

        Parameters
        ----------
        original : np.ndarray
            原始数据。
        perturbed : np.ndarray
            扰动后的数据。
        task : dict, optional
            可覆盖默认的 columns 和 query_type。

        Returns
        -------
        float
            效用值，范围 [0,1]，值越大表示查询结果越接近。
        """
        # 确定查询列
        cols = self.columns
        if task and 'columns' in task:
            cols = task['columns']

        # 确定查询类型
        qtype = self.query_type
        if task and 'query_type' in task:
            qtype = task['query_type']

        # 选择数据列
        if cols is None:
            orig_sel = original
            pert_sel = perturbed
        else:
            if isinstance(cols, int):
                cols = [cols]
            orig_sel = original[:, cols]
            pert_sel = perturbed[:, cols]

        # 计算原始数据的查询值
        if qtype == 'mean':
            orig_val = np.mean(orig_sel, axis=0)
            pert_val = np.mean(pert_sel, axis=0)
        elif qtype == 'sum':
            orig_val = np.sum(orig_sel, axis=0)
            pert_val = np.sum(pert_sel, axis=0)
        elif qtype == 'count':
            orig_val = orig_sel.shape[0]
            pert_val = pert_sel.shape[0]
        elif qtype == 'var':
            orig_val = np.var(orig_sel, axis=0)
            pert_val = np.var(pert_sel, axis=0)
        else:
            raise ValueError(f"不支持的查询类型: {qtype}")

        # 计算相对误差，然后转换为效用（1 - 归一化误差）
        # 对于多列，取平均
        abs_diff = np.abs(orig_val - pert_val)
        denom = np.abs(orig_val) + self.epsilon
        rel_error = abs_diff / denom
        # 限制最大误差为1，避免负值
        rel_error = np.clip(rel_error, 0, 1)
        utility = 1 - np.mean(rel_error)

        return float(utility)