import pandas as pd
import numpy as np
from typing import Dict, List, Tuple, Optional, Callable, Union


class FeatureAnalyzer:
    """分析数据集特征，为隐私机制提供先验信息"""

    def __init__(self,
                 data: pd.DataFrame,
                 sensitivity_type: Union[str, Callable] = 'range'):
        """
        参数:
            data: 输入数据
            sensitivity_type: 敏感度计算方式。
                - 'range': max - min
                - 'std': 标准差
                - callable: 自定义函数，接受列数据返回敏感度
        """
        self.data = data
        self.column_types = self._detect_column_types()
        self.stats = self._compute_statistics()
        self.sensitivity = self._compute_sensitivity(sensitivity_type)

    def _detect_column_types(self) -> Dict[str, str]:
        """检测每列的类型：'numeric' 或 'categorical'"""
        types = {}
        for col in self.data.columns:
            if pd.api.types.is_numeric_dtype(self.data[col]):
                types[col] = 'numeric'
            else:
                types[col] = 'categorical'
        return types

    def _compute_statistics(self) -> Dict[str, Dict]:
        """计算每列的基本统计信息"""
        stats = {}
        for col, col_type in self.column_types.items():
            if col_type == 'numeric':
                stats[col] = {
                    'mean': self.data[col].mean(),
                    'std': self.data[col].std(),
                    'min': self.data[col].min(),
                    'max': self.data[col].max(),
                    'q1': self.data[col].quantile(0.25),
                    'median': self.data[col].median(),
                    'q3': self.data[col].quantile(0.75)
                }
            else:
                # 对于类别列，记录唯一值和频率
                value_counts = self.data[col].value_counts()
                stats[col] = {
                    'unique_values': value_counts.index.tolist(),
                    'frequencies': value_counts.values.tolist()
                }
        return stats

    def _compute_sensitivity(self, method: Union[str, Callable]) -> Dict[str, float]:
        """计算每列的全局敏感度"""
        sensitivity = {}
        for col, col_type in self.column_types.items():
            if col_type == 'numeric':
                if callable(method):
                    # 自定义函数
                    sensitivity[col] = method(self.data[col])
                elif method == 'range':
                    sensitivity[col] = self.stats[col]['max'] - self.stats[col]['min']
                elif method == 'std':
                    sensitivity[col] = self.stats[col]['std']
                else:
                    raise ValueError(f"Unknown sensitivity method: {method}")
            else:
                # 对于类别列，敏感度通常为1（改变一个记录最多改变一个类别）
                sensitivity[col] = 1.0
        return sensitivity

    def get_numeric_columns(self) -> List[str]:
        return [col for col, typ in self.column_types.items() if typ == 'numeric']

    def get_categorical_columns(self) -> List[str]:
        return [col for col, typ in self.column_types.items() if typ == 'categorical']