import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, roc_auc_score
from typing import Tuple, List


class MembershipInferenceAttack:
    """
    成员推断攻击验证器。
    使用影子模型训练攻击模型，评估目标模型的隐私泄露风险。
    """

    def __init__(self, random_state: int = 42):
        self.random_state = random_state
        self.attack_model = None
        self.target_model = None

    def _train_shadow_models(self, X, y, n_shadow=5, test_size=0.3):
        """训练影子模型并生成攻击训练数据"""
        shadow_models = []
        attack_X = []
        attack_y = []

        for i in range(n_shadow):
            # 划分影子模型的训练集和测试集（非成员）
            X_shadow, X_nonmember, y_shadow, y_nonmember = train_test_split(
                X, y, test_size=test_size, random_state=self.random_state + i
            )

            # 训练影子模型（与目标模型同类型）
            model = RandomForestClassifier(n_estimators=10, random_state=self.random_state)
            model.fit(X_shadow, y_shadow)
            shadow_models.append(model)

            # 生成成员数据（训练集）的预测结果
            pred_member = model.predict_proba(X_shadow)
            # 生成非成员数据（测试集）的预测结果
            pred_nonmember = model.predict_proba(X_nonmember)

            # 将预测结果和真实标签拼接作为攻击模型的特征
            # 特征：[真实标签, 预测概率0, 预测概率1, 正确预测的置信度]
            for j in range(len(X_shadow)):
                true_label = y_shadow[j]
                probs = pred_member[j]
                confidence = probs[true_label]  # 正确类的概率
                attack_X.append([true_label, probs[0], probs[1], confidence])
                attack_y.append(1)  # 1表示成员

            for j in range(len(X_nonmember)):
                true_label = y_nonmember[j]
                probs = pred_nonmember[j]
                confidence = probs[true_label]
                attack_X.append([true_label, probs[0], probs[1], confidence])
                attack_y.append(0)  # 0表示非成员

        return np.array(attack_X), np.array(attack_y)

    def fit_attack_model(self, X, y):
        """训练攻击模型"""
        attack_X, attack_y = self._train_shadow_models(X, y)
        self.attack_model = RandomForestClassifier(n_estimators=20, random_state=self.random_state)
        self.attack_model.fit(attack_X, attack_y)
        return self.attack_model

    def evaluate_target_model(self, target_model, X_member, X_nonmember):
        """
        评估目标模型的成员推断风险。
        返回攻击准确率和AUC。
        """
        if self.attack_model is None:
            raise ValueError("必须先训练攻击模型")

        # 获取目标模型对成员和非成员的预测
        pred_member = target_model.predict_proba(X_member)
        pred_nonmember = target_model.predict_proba(X_nonmember)

        # 构造攻击测试数据
        attack_test_X = []
        attack_test_y = []

        # 真实标签需要从目标模型获取（这里假设我们已知标签，但在实际攻击中可能未知，简化处理）
        # 为了演示，我们使用随机生成的标签（真实场景中攻击者可能不知道真实标签，只能利用预测结果）
        # 这里我们简化：直接使用真实标签（需要传入 y_member 和 y_nonmember）
        # 为了不增加参数，我们假设标签是预测的最高概率类（攻击者视角）

        for i in range(len(X_member)):
            probs = pred_member[i]
            true_label = np.argmax(probs)  # 攻击者假设的最高概率类
            confidence = probs[true_label]
            attack_test_X.append([true_label, probs[0], probs[1], confidence])
            attack_test_y.append(1)

        for i in range(len(X_nonmember)):
            probs = pred_nonmember[i]
            true_label = np.argmax(probs)
            confidence = probs[true_label]
            attack_test_X.append([true_label, probs[0], probs[1], confidence])
            attack_test_y.append(0)

        attack_test_X = np.array(attack_test_X)
        attack_test_y = np.array(attack_test_y)

        # 攻击模型预测
        pred = self.attack_model.predict(attack_test_X)
        acc = accuracy_score(attack_test_y, pred)
        auc = roc_auc_score(attack_test_y, self.attack_model.predict_proba(attack_test_X)[:, 1])

        return acc, auc

    def validate_risk_consistency(self, target_model, X_member, X_nonmember, risk_meter_value):
        """
        验证攻击准确率与风险度量值的一致性。
        返回攻击准确率，并打印对比。
        """
        acc, auc = self.evaluate_target_model(target_model, X_member, X_nonmember)
        print(f"成员推断攻击准确率: {acc:.4f}, AUC: {auc:.4f}")
        print(f"风险度量器输出: {risk_meter_value:.4f}")
        if acc > risk_meter_value + 0.1:
            print("警告：攻击准确率显著高于风险度量值，可能低估了风险")
        elif acc < risk_meter_value - 0.1:
            print("注意：攻击准确率低于风险度量值，风险度量可能偏保守")
        else:
            print("风险度量与攻击结果基本一致")
        return acc, auc