import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from typing import Optional, Dict, Any
from ..risk_meter import RiskMeter


class EmpiricalRiskMeter(RiskMeter):
    """
    经验风险度量器。
    通过训练一个KNN分类器来区分原始数据和扰动数据，
    以分类准确率作为隐私风险值（越高表示风险越大）。

    参数:
        test_size: float, 用于评估的测试集比例（默认0.2）
        random_state: int, 随机种子（默认42）
        n_neighbors: int, KNN的邻居数（默认5）
    """

    def __init__(self, test_size: float = 0.2, random_state: int = 42, n_neighbors: int = 5):
        self.test_size = test_size
        self.random_state = random_state
        self.n_neighbors = n_neighbors

    def measure(self,
                original: np.ndarray,
                perturbed: np.ndarray,
                params: Optional[Dict[str, Any]] = None) -> float:
        """
        测量隐私风险。

        Parameters
        ----------
        original : np.ndarray
            原始数据。
        perturbed : np.ndarray
            扰动后的数据。

        Returns
        -------
        float
            风险值，范围 [0, 1]。值越大，表示原始数据和扰动数据越可区分，隐私风险越高。
            如果返回值为0.5，表示随机猜测水平，风险较低。
        """
        # 构造标签：原始数据为1，扰动数据为0
        n_orig = original.shape[0]
        n_pert = perturbed.shape[0]

        # 合并数据
        X = np.vstack([original, perturbed])
        y = np.hstack([np.ones(n_orig), np.zeros(n_pert)])

        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=self.test_size, random_state=self.random_state, stratify=y
        )

        # 训练KNN分类器
        knn = KNeighborsClassifier(n_neighbors=self.n_neighbors)
        knn.fit(X_train, y_train)

        # 预测并计算准确率
        y_pred = knn.predict(X_test)
        acc = accuracy_score(y_test, y_pred)

        return acc