05-模型安全与治理

"""
SHAP (SHapley Additive exPlanations) 解释器
基于博弈论的特征归因方法
"""

import shap
import numpy as np
import pandas as pd
from typing import Dict, List, Any, Optional, Union
from dataclasses import dataclass
import matplotlib.pyplot as plt
from abc import ABC, abstractmethod
import json


@dataclass
class ExplanationResult:
    """解释结果"""
    feature_names: List[str]
    shap_values: np.ndarray
    base_value: float
    expected_value: float
    feature_importance: Dict[str, float]
    prediction: Any


class SHAPExplainer:
    """SHAP 解释器"""

    def __init__(self, model, model_type: str = "tree"):
        """
        初始化 SHAP 解释器

        Args:
            model: 训练好的模型
            model_type: 模型类型 (tree, kernel, deep, linear)
        """
        self.model = model
        self.model_type = model_type
        self.explainer = None
        self.background_data = None

    def fit(self, background_data: Union[np.ndarray, pd.DataFrame],
            sample_size: int = 100):
        """
        拟合解释器

        Args:
            background_data: 背景数据用于计算期望值
            sample_size: 采样数量
        """
        # 采样背景数据
        if len(background_data) > sample_size:
            indices = np.random.choice(
                len(background_data),
                sample_size,
                replace=False
            )
            if isinstance(background_data, pd.DataFrame):
                self.background_data = background_data.iloc[indices]
            else:
                self.background_data = background_data[indices]
        else:
            self.background_data = background_data

        # 创建解释器
        if self.model_type == "tree":
            self.explainer = shap.TreeExplainer(self.model)
        elif self.model_type == "kernel":
            self.explainer = shap.KernelExplainer(
                self.model.predict_proba if hasattr(self.model, 'predict_proba')
                else self.model.predict,
                self.background_data
            )
        elif self.model_type == "deep":
            self.explainer = shap.DeepExplainer(
                self.model,
                self.background_data
            )
        elif self.model_type == "linear":
            self.explainer = shap.LinearExplainer(
                self.model,
                self.background_data
            )
        else:
            # 默认使用 Kernel 解释器
            self.explainer = shap.KernelExplainer(
                self.model.predict,
                self.background_data
            )

    def explain(self, X: Union[np.ndarray, pd.DataFrame],
                feature_names: Optional[List[str]] = None) -> ExplanationResult:
        """
        解释预测结果

        Args:
            X: 待解释的样本
            feature_names: 特征名称

        Returns:
            ExplanationResult: 解释结果
        """
        if self.explainer is None:
            raise ValueError("请先调用 fit() 方法")

        # 计算 SHAP 值
        shap_values = self.explainer.shap_values(X)

        # 获取期望值
        if hasattr(self.explainer, 'expected_value'):
            expected_value = self.explainer.expected_value
            if isinstance(expected_value, np.ndarray):
                expected_value = expected_value[0]
        else:
            expected_value = 0.0

        # 获取特征名称
        if feature_names is None:
            if isinstance(X, pd.DataFrame):
                feature_names = X.columns.tolist()
            else:
                feature_names = [f"feature_{i}" for i in range(X.shape[1])]

        # 处理多分类情况
        if isinstance(shap_values, list):
            shap_values = shap_values[1]  # 取正类

        # 计算特征重要性
        if len(shap_values.shape) == 1:
            importance = dict(zip(feature_names, np.abs(shap_values)))
        else:
            importance = dict(zip(
                feature_names,
                np.mean(np.abs(shap_values), axis=0)
            ))

        # 获取预测结果
        if hasattr(self.model, 'predict_proba'):
            prediction = self.model.predict_proba(X)
        else:
            prediction = self.model.predict(X)

        return ExplanationResult(
            feature_names=feature_names,
            shap_values=shap_values,
            base_value=expected_value,
            expected_value=expected_value,
            feature_importance=importance,
            prediction=prediction
        )

    def explain_instance(self, x: np.ndarray,
                         feature_names: Optional[List[str]] = None,
                         top_k: int = 10) -> Dict[str, Any]:
        """
        解释单个实例

        Args:
            x: 单个样本
            feature_names: 特征名称
            top_k: 返回前 k 个重要特征

        Returns:
            Dict: 解释详情
        """
        if x.ndim == 1:
            x = x.reshape(1, -1)

        result = self.explain(x, feature_names)

        # 排序特征重要性
        sorted_importance = sorted(
            result.feature_importance.items(),
            key=lambda x: abs(x[1]),
            reverse=True
        )[:top_k]

        # 构建解释
        explanation = {
            "prediction": float(result.prediction[0]) if isinstance(result.prediction[0], (np.floating, float)) else result.prediction[0].tolist(),
            "base_value": float(result.base_value),
            "top_features": [
                {
                    "name": name,
                    "shap_value": float(value),
                    "contribution": "positive" if value > 0 else "negative"
                }
                for name, value in sorted_importance
            ],
            "feature_impacts": {}
        }

        # 计算每个特征的影响
        for i, name in enumerate(result.feature_names):
            shap_val = result.shap_values[0, i] if result.shap_values.ndim > 1 else result.shap_values[i]
            explanation["feature_impacts"][name] = {
                "shap_value": float(shap_val),
                "feature_value": float(x[0, i])
            }

        return explanation

    def plot_summary(self, X: Union[np.ndarray, pd.DataFrame],
                     feature_names: Optional[List[str]] = None,
                     max_display: int = 20,
                     save_path: Optional[str] = None):
        """
        绘制 SHAP 摘要图

        Args:
            X: 数据集
            feature_names: 特征名称
            max_display: 最大显示特征数
            save_path: 保存路径
        """
        result = self.explain(X, feature_names)

        plt.figure(figsize=(10, 8))
        shap.summary_plot(
            result.shap_values,
            X,
            feature_names=result.feature_names,
            max_display=max_display,
            show=False
        )

        if save_path:
            plt.savefig(save_path, bbox_inches='tight', dpi=150)
        plt.close()

    def plot_force(self, x: np.ndarray,
                   feature_names: Optional[List[str]] = None,
                   save_path: Optional[str] = None):
        """
        绘制力图

        Args:
            x: 单个样本
            feature_names: 特征名称
            save_path: 保存路径
        """
        if x.ndim == 1:
            x = x.reshape(1, -1)

        result = self.explain(x, feature_names)

        shap.force_plot(
            result.base_value,
            result.shap_values[0] if result.shap_values.ndim > 1 else result.shap_values,
            x[0],
            feature_names=result.feature_names,
            matplotlib=True,
            show=False
        )

        if save_path:
            plt.savefig(save_path, bbox_inches='tight', dpi=150)
        plt.close()


class LIMEExplainer:
    """LIME (Local Interpretable Model-agnostic Explanations) 解释器"""

    def __init__(self, model, mode: str = "classification"):
        """
        初始化 LIME 解释器

        Args:
            model: 训练好的模型
            mode: 模式 (classification, regression)
        """
        from lime.lime_tabular import LimeTabularExplainer

        self.model = model
        self.mode = mode
        self.explainer = None

    def fit(self, training_data: np.ndarray,
            feature_names: Optional[List[str]] = None,
            categorical_features: Optional[List[int]] = None,
            class_names: Optional[List[str]] = None):
        """
        拟合 LIME 解释器

        Args:
            training_data: 训练数据
            feature_names: 特征名称
            categorical_features: 分类特征索引
            class_names: 类别名称
        """
        from lime.lime_tabular import LimeTabularExplainer

        self.explainer = LimeTabularExplainer(
            training_data,
            feature_names=feature_names,
            categorical_features=categorical_features,
            class_names=class_names,
            mode=self.mode
        )

    def explain_instance(self, x: np.ndarray,
                         num_features: int = 10,
                         num_samples: int = 5000) -> Dict[str, Any]:
        """
        解释单个实例

        Args:
            x: 单个样本
            num_features: 解释特征数
            num_samples: 采样数量

        Returns:
            Dict: 解释结果
        """
        if self.explainer is None:
            raise ValueError("请先调用 fit() 方法")

        # 获取预测函数
        if self.mode == "classification":
            predict_fn = self.model.predict_proba
        else:
            predict_fn = self.model.predict

        # 生成解释
        exp = self.explainer.explain_instance(
            x,
            predict_fn,
            num_features=num_features,
            num_samples=num_samples
        )

        # 提取结果
        explanation = {
            "prediction": predict_fn(x.reshape(1, -1))[0].tolist() if self.mode == "classification" else float(predict_fn(x.reshape(1, -1))[0]),
            "intercept": float(exp.intercept[1]) if self.mode == "classification" else float(exp.intercept),
            "features": [
                {
                    "description": feature,
                    "weight": float(weight)
                }
                for feature, weight in exp.as_list()
            ],
            "local_prediction": float(exp.local_pred[0]) if hasattr(exp, 'local_pred') else None
        }

        return explanation


class IntegratedGradientsExplainer:
    """积分梯度解释器 - 用于深度学习模型"""

    def __init__(self, model, baseline: Optional[np.ndarray] = None):
        """
        初始化积分梯度解释器

        Args:
            model: PyTorch/TensorFlow 模型
            baseline: 基线输入
        """
        self.model = model
        self.baseline = baseline

    def compute_gradients(self, inputs, target_class: Optional[int] = None):
        """计算梯度"""
        import torch

        inputs = torch.tensor(inputs, requires_grad=True, dtype=torch.float32)
        outputs = self.model(inputs)

        if target_class is not None:
            outputs = outputs[:, target_class]
        else:
            outputs = outputs.sum()

        outputs.backward()
        return inputs.grad.numpy()

    def explain(self, x: np.ndarray,
                steps: int = 50,
                target_class: Optional[int] = None) -> Dict[str, Any]:
        """
        计算积分梯度

        Args:
            x: 输入样本
            steps: 积分步数
            target_class: 目标类别

        Returns:
            Dict: 归因结果
        """
        import torch

        # 设置基线
        if self.baseline is None:
            baseline = np.zeros_like(x)
        else:
            baseline = self.baseline

        # 生成插值路径
        alphas = np.linspace(0, 1, steps)
        interpolated_inputs = np.array([
            baseline + alpha * (x - baseline)
            for alpha in alphas
        ])

        # 计算路径上的梯度
        gradients = []
        for inp in interpolated_inputs:
            grad = self.compute_gradients(
                inp.reshape(1, -1),
                target_class
            )
            gradients.append(grad)

        gradients = np.array(gradients)

        # 计算积分梯度
        avg_gradients = np.mean(gradients, axis=0)
        integrated_gradients = (x - baseline) * avg_gradients

        return {
            "attributions": integrated_gradients.flatten().tolist(),
            "baseline": baseline.flatten().tolist(),
            "delta": (x - baseline).flatten().tolist()
        }

"""
大语言模型可解释性工具
"""

import torch
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from transformers import AutoModelForCausalLM, AutoTokenizer
import matplotlib.pyplot as plt


class LLMExplainer:
    """大语言模型解释器"""

    def __init__(self, model_name: str):
        """
        初始化 LLM 解释器

        Args:
            model_name: HuggingFace 模型名称
        """
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            output_attentions=True,
            output_hidden_states=True
        )
        self.model.eval()

    def get_attention_weights(self, text: str) -> Dict[str, Any]:
        """
        获取注意力权重

        Args:
            text: 输入文本

        Returns:
            Dict: 注意力权重信息
        """
        # 分词
        inputs = self.tokenizer(text, return_tensors="pt")
        tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

        # 前向传播
        with torch.no_grad():
            outputs = self.model(**inputs)

        # 提取注意力权重
        attentions = outputs.attentions  # (num_layers, batch, num_heads, seq_len, seq_len)

        # 平均所有层和头的注意力
        avg_attention = torch.stack(attentions).mean(dim=[0, 1, 2])

        return {
            "tokens": tokens,
            "attention_matrix": avg_attention.numpy().tolist(),
            "layer_attentions": [
                att[0].mean(dim=0).numpy().tolist()
                for att in attentions
            ]
        }

    def token_importance(self, text: str,
                         target_position: int = -1) -> List[Dict[str, float]]:
        """
        计算每个 token 对目标位置的重要性

        Args:
            text: 输入文本
            target_position: 目标 token 位置

        Returns:
            List: token 重要性列表
        """
        attention_info = self.get_attention_weights(text)
        tokens = attention_info["tokens"]
        attention_matrix = np.array(attention_info["attention_matrix"])

        # 获取目标位置对所有输入的注意力
        importance = attention_matrix[target_position]

        return [
            {"token": token, "importance": float(imp)}
            for token, imp in zip(tokens, importance)
        ]

    def analyze_generation(self, prompt: str,
                           max_new_tokens: int = 50) -> Dict[str, Any]:
        """
        分析生成过程

        Args:
            prompt: 输入提示
            max_new_tokens: 最大生成 token 数

        Returns:
            Dict: 生成分析结果
        """
        inputs = self.tokenizer(prompt, return_tensors="pt")

        # 逐 token 生成并记录
        generated_tokens = []
        token_probs = []
        attention_history = []

        input_ids = inputs["input_ids"]

        for _ in range(max_new_tokens):
            with torch.no_grad():
                outputs = self.model(input_ids, output_attentions=True)

            # 获取下一个 token 的概率分布
            logits = outputs.logits[0, -1]
            probs = torch.softmax(logits, dim=-1)

            # 采样下一个 token
            next_token_id = torch.argmax(probs).unsqueeze(0).unsqueeze(0)
            next_token = self.tokenizer.decode(next_token_id[0])

            # 记录
            generated_tokens.append(next_token)
            token_probs.append({
                "token": next_token,
                "probability": float(probs[next_token_id[0, 0]]),
                "top_5": [
                    {
                        "token": self.tokenizer.decode([idx]),
                        "probability": float(probs[idx])
                    }
                    for idx in torch.topk(probs, 5).indices
                ]
            })

            # 更新输入
            input_ids = torch.cat([input_ids, next_token_id], dim=1)

            # 检查是否结束
            if next_token_id[0, 0] == self.tokenizer.eos_token_id:
                break

        return {
            "prompt": prompt,
            "generated_text": "".join(generated_tokens),
            "token_analysis": token_probs
        }

    def counterfactual_analysis(self, text: str,
                                 target_word: str,
                                 replacements: List[str]) -> Dict[str, Any]:
        """
        反事实分析 - 替换关键词观察输出变化

        Args:
            text: 原始文本
            target_word: 目标词
            replacements: 替换词列表

        Returns:
            Dict: 分析结果
        """
        results = {"original": {}, "counterfactuals": []}

        # 原始输出
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=50,
                output_scores=True,
                return_dict_in_generate=True
            )
        results["original"] = {
            "text": text,
            "output": self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
        }

        # 反事实分析
        for replacement in replacements:
            modified_text = text.replace(target_word, replacement)
            inputs = self.tokenizer(modified_text, return_tensors="pt")

            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=50
                )

            results["counterfactuals"].append({
                "replacement": replacement,
                "text": modified_text,
                "output": self.tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
            })

        return results


class AttentionVisualizer:
    """注意力可视化工具"""

    @staticmethod
    def plot_attention_heatmap(tokens: List[str],
                               attention_matrix: np.ndarray,
                               save_path: Optional[str] = None):
        """
        绘制注意力热力图

        Args:
            tokens: token 列表
            attention_matrix: 注意力矩阵
            save_path: 保存路径
        """
        fig, ax = plt.subplots(figsize=(12, 10))

        im = ax.imshow(attention_matrix, cmap='Blues')

        ax.set_xticks(range(len(tokens)))
        ax.set_yticks(range(len(tokens)))
        ax.set_xticklabels(tokens, rotation=45, ha='right')
        ax.set_yticklabels(tokens)

        plt.colorbar(im)
        plt.title('Attention Weights')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()

    @staticmethod
    def plot_layer_attention(layer_attentions: List[np.ndarray],
                             tokens: List[str],
                             save_path: Optional[str] = None):
        """
        绘制各层注意力分布

        Args:
            layer_attentions: 各层注意力
            tokens: token 列表
            save_path: 保存路径
        """
        num_layers = len(layer_attentions)
        fig, axes = plt.subplots(
            (num_layers + 3) // 4, 4,
            figsize=(16, 4 * ((num_layers + 3) // 4))
        )
        axes = axes.flatten()

        for i, attention in enumerate(layer_attentions):
            im = axes[i].imshow(np.array(attention), cmap='Blues')
            axes[i].set_title(f'Layer {i + 1}')
            axes[i].set_xticks([])
            axes[i].set_yticks([])

        # 隐藏多余的子图
        for i in range(num_layers, len(axes)):
            axes[i].axis('off')

        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()

"""
对抗攻击实现
包括 FGSM, PGD, C&W 等经典攻击方法
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional, Tuple, Callable
from abc import ABC, abstractmethod


class AdversarialAttack(ABC):
    """对抗攻击基类"""

    def __init__(self, model: nn.Module, device: str = "cuda"):
        self.model = model
        self.device = device
        self.model.to(device)
        self.model.eval()

    @abstractmethod
    def generate(self, x: torch.Tensor, y: torch.Tensor,
                 **kwargs) -> torch.Tensor:
        """生成对抗样本"""
        pass

    def _get_loss(self, x: torch.Tensor, y: torch.Tensor,
                  targeted: bool = False) -> torch.Tensor:
        """计算损失"""
        outputs = self.model(x)
        loss = F.cross_entropy(outputs, y)
        return -loss if targeted else loss


class FGSM(AdversarialAttack):
    """Fast Gradient Sign Method"""

    def generate(self, x: torch.Tensor, y: torch.Tensor,
                 epsilon: float = 0.03,
                 targeted: bool = False) -> torch.Tensor:
        """
        生成 FGSM 对抗样本

        Args:
            x: 原始输入
            y: 标签
            epsilon: 扰动大小
            targeted: 是否为目标攻击

        Returns:
            对抗样本
        """
        x = x.clone().detach().to(self.device)
        y = y.clone().detach().to(self.device)
        x.requires_grad = True

        # 计算损失
        loss = self._get_loss(x, y, targeted)

        # 计算梯度
        loss.backward()

        # 生成对抗样本
        x_adv = x + epsilon * x.grad.sign()
        x_adv = torch.clamp(x_adv, 0, 1)

        return x_adv.detach()


class PGD(AdversarialAttack):
    """Projected Gradient Descent"""

    def generate(self, x: torch.Tensor, y: torch.Tensor,
                 epsilon: float = 0.03,
                 alpha: float = 0.007,
                 num_steps: int = 40,
                 random_start: bool = True,
                 targeted: bool = False) -> torch.Tensor:
        """
        生成 PGD 对抗样本

        Args:
            x: 原始输入
            y: 标签
            epsilon: 最大扰动
            alpha: 步长
            num_steps: 迭代次数
            random_start: 是否随机初始化
            targeted: 是否为目标攻击

        Returns:
            对抗样本
        """
        x = x.clone().detach().to(self.device)
        y = y.clone().detach().to(self.device)

        # 随机初始化
        if random_start:
            x_adv = x + torch.empty_like(x).uniform_(-epsilon, epsilon)
            x_adv = torch.clamp(x_adv, 0, 1)
        else:
            x_adv = x.clone()

        # 迭代优化
        for _ in range(num_steps):
            x_adv.requires_grad = True

            loss = self._get_loss(x_adv, y, targeted)
            loss.backward()

            with torch.no_grad():
                # 梯度更新
                x_adv = x_adv + alpha * x_adv.grad.sign()

                # 投影到 epsilon 球内
                delta = torch.clamp(x_adv - x, -epsilon, epsilon)
                x_adv = torch.clamp(x + delta, 0, 1)

        return x_adv.detach()


class CarliniWagner(AdversarialAttack):
    """Carlini & Wagner L2 Attack"""

    def generate(self, x: torch.Tensor, y: torch.Tensor,
                 targeted: bool = False,
                 c: float = 1.0,
                 kappa: float = 0,
                 max_iterations: int = 1000,
                 learning_rate: float = 0.01,
                 binary_search_steps: int = 9) -> torch.Tensor:
        """
        生成 C&W 对抗样本

        Args:
            x: 原始输入
            y: 标签 (targeted 时为目标标签)
            targeted: 是否为目标攻击
            c: 置信度参数
            kappa: 攻击置信度
            max_iterations: 最大迭代次数
            learning_rate: 学习率
            binary_search_steps: 二分搜索步数

        Returns:
            对抗样本
        """
        x = x.clone().detach().to(self.device)
        y = y.clone().detach().to(self.device)

        # 变换到 tanh 空间
        x_tanh = self._to_tanh_space(x)

        # 初始化
        w = x_tanh.clone().requires_grad_(True)
        optimizer = torch.optim.Adam([w], lr=learning_rate)

        best_adv = x.clone()
        best_l2 = float('inf')

        # 二分搜索 c
        c_low = 0.0
        c_high = c

        for _ in range(binary_search_steps):
            c_mid = (c_low + c_high) / 2

            for _ in range(max_iterations):
                optimizer.zero_grad()

                # 从 tanh 空间转换回来
                x_adv = self._from_tanh_space(w)

                # 计算 L2 距离
                l2_dist = torch.sum((x_adv - x) ** 2)

                # 计算 f 函数
                outputs = self.model(x_adv)
                one_hot = F.one_hot(y, outputs.shape[-1]).float()

                real = (one_hot * outputs).sum(dim=1)
                other = ((1 - one_hot) * outputs - one_hot * 10000).max(dim=1)[0]

                if targeted:
                    f = torch.clamp(other - real + kappa, min=0)
                else:
                    f = torch.clamp(real - other + kappa, min=0)

                # 总损失
                loss = l2_dist + c_mid * f.sum()
                loss.backward()
                optimizer.step()

                # 更新最佳结果
                with torch.no_grad():
                    x_adv = self._from_tanh_space(w)
                    pred = self.model(x_adv).argmax(dim=1)

                    if targeted:
                        success = (pred == y).all()
                    else:
                        success = (pred != y).all()

                    if success:
                        current_l2 = torch.sum((x_adv - x) ** 2).item()
                        if current_l2 < best_l2:
                            best_l2 = current_l2
                            best_adv = x_adv.clone()

            # 调整 c
            if best_l2 < float('inf'):
                c_high = c_mid
            else:
                c_low = c_mid

        return best_adv

    def _to_tanh_space(self, x: torch.Tensor) -> torch.Tensor:
        """转换到 tanh 空间"""
        return torch.atanh(2 * x - 1 + 1e-7)

    def _from_tanh_space(self, w: torch.Tensor) -> torch.Tensor:
        """从 tanh 空间转换回来"""
        return (torch.tanh(w) + 1) / 2


class TextAdversarialAttack:
    """文本对抗攻击"""

    def __init__(self, model, tokenizer,
                 similarity_threshold: float = 0.8):
        """
        初始化文本对抗攻击

        Args:
            model: 文本分类模型
            tokenizer: 分词器
            similarity_threshold: 语义相似度阈值
        """
        self.model = model
        self.tokenizer = tokenizer
        self.similarity_threshold = similarity_threshold

        # 加载同义词库
        self.synonyms = self._load_synonyms()

    def _load_synonyms(self) -> dict:
        """加载同义词库"""
        # 实际应用中可使用 WordNet 或自定义同义词库
        return {}

    def get_synonyms(self, word: str) -> list:
        """获取同义词"""
        return self.synonyms.get(word.lower(), [])

    def textfooler_attack(self, text: str,
                          label: int,
                          max_perturbations: int = 10) -> dict:
        """
        TextFooler 攻击

        Args:
            text: 原始文本
            label: 原始标签
            max_perturbations: 最大扰动次数

        Returns:
            攻击结果
        """
        words = text.split()
        word_importance = self._compute_word_importance(words, label)

        # 按重要性排序
        sorted_indices = sorted(
            range(len(words)),
            key=lambda i: word_importance[i],
            reverse=True
        )

        perturbed_words = words.copy()
        num_perturbations = 0

        for idx in sorted_indices:
            if num_perturbations >= max_perturbations:
                break

            word = words[idx]
            synonyms = self.get_synonyms(word)

            for synonym in synonyms:
                # 尝试替换
                perturbed_words[idx] = synonym
                perturbed_text = " ".join(perturbed_words)

                # 检查是否攻击成功
                pred = self._predict(perturbed_text)
                if pred != label:
                    return {
                        "success": True,
                        "original_text": text,
                        "perturbed_text": perturbed_text,
                        "perturbations": num_perturbations + 1,
                        "original_pred": label,
                        "perturbed_pred": pred
                    }

            # 恢复原词
            perturbed_words[idx] = word
            num_perturbations += 1

        return {
            "success": False,
            "original_text": text,
            "perturbed_text": " ".join(perturbed_words),
            "perturbations": num_perturbations
        }

    def _compute_word_importance(self, words: list, label: int) -> list:
        """计算词重要性"""
        importance = []
        original_prob = self._get_probability(words, label)

        for i in range(len(words)):
            # 删除词后的概率
            words_without = words[:i] + words[i+1:]
            prob_without = self._get_probability(words_without, label)
            importance.append(original_prob - prob_without)

        return importance

    def _get_probability(self, words: list, label: int) -> float:
        """获取预测概率"""
        text = " ".join(words)
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = F.softmax(outputs.logits, dim=-1)
        return probs[0, label].item()

    def _predict(self, text: str) -> int:
        """预测"""
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.logits.argmax(dim=-1).item()

"""
对抗防御策略实现
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Optional, Tuple, List
from dataclasses import dataclass


class AdversarialTraining:
    """对抗训练"""

    def __init__(self, model: nn.Module,
                 attack_method: str = "pgd",
                 epsilon: float = 0.03,
                 alpha: float = 0.007,
                 num_steps: int = 7):
        """
        初始化对抗训练

        Args:
            model: 模型
            attack_method: 攻击方法
            epsilon: 扰动大小
            alpha: 步长
            num_steps: PGD 步数
        """
        self.model = model
        self.attack_method = attack_method
        self.epsilon = epsilon
        self.alpha = alpha
        self.num_steps = num_steps

    def generate_adversarial(self, x: torch.Tensor,
                             y: torch.Tensor) -> torch.Tensor:
        """生成对抗样本用于训练"""
        self.model.eval()

        if self.attack_method == "fgsm":
            x_adv = self._fgsm_attack(x, y)
        elif self.attack_method == "pgd":
            x_adv = self._pgd_attack(x, y)
        else:
            x_adv = x

        self.model.train()
        return x_adv

    def _fgsm_attack(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """FGSM 攻击"""
        x = x.clone().requires_grad_(True)
        outputs = self.model(x)
        loss = F.cross_entropy(outputs, y)
        loss.backward()

        x_adv = x + self.epsilon * x.grad.sign()
        x_adv = torch.clamp(x_adv, 0, 1)
        return x_adv.detach()

    def _pgd_attack(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """PGD 攻击"""
        x_adv = x.clone().detach()
        x_adv = x_adv + torch.empty_like(x_adv).uniform_(-self.epsilon, self.epsilon)
        x_adv = torch.clamp(x_adv, 0, 1)

        for _ in range(self.num_steps):
            x_adv.requires_grad_(True)
            outputs = self.model(x_adv)
            loss = F.cross_entropy(outputs, y)
            loss.backward()

            with torch.no_grad():
                x_adv = x_adv + self.alpha * x_adv.grad.sign()
                delta = torch.clamp(x_adv - x, -self.epsilon, self.epsilon)
                x_adv = torch.clamp(x + delta, 0, 1)

        return x_adv.detach()

    def train_step(self, x: torch.Tensor, y: torch.Tensor,
                   optimizer: torch.optim.Optimizer,
                   mix_ratio: float = 0.5) -> Tuple[float, float]:
        """
        对抗训练步骤

        Args:
            x: 输入
            y: 标签
            optimizer: 优化器
            mix_ratio: 对抗样本比例

        Returns:
            (clean_loss, adv_loss)
        """
        self.model.train()
        optimizer.zero_grad()

        # 干净样本损失
        outputs_clean = self.model(x)
        loss_clean = F.cross_entropy(outputs_clean, y)

        # 对抗样本损失
        x_adv = self.generate_adversarial(x, y)
        outputs_adv = self.model(x_adv)
        loss_adv = F.cross_entropy(outputs_adv, y)

        # 混合损失
        total_loss = (1 - mix_ratio) * loss_clean + mix_ratio * loss_adv
        total_loss.backward()
        optimizer.step()

        return loss_clean.item(), loss_adv.item()


class InputPreprocessing:
    """输入预处理防御"""

    @staticmethod
    def jpeg_compression(x: torch.Tensor, quality: int = 75) -> torch.Tensor:
        """
        JPEG 压缩防御

        Args:
            x: 输入图像
            quality: 压缩质量

        Returns:
            压缩后的图像
        """
        from PIL import Image
        import io

        # 转换为 PIL 图像
        x_np = (x.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
        img = Image.fromarray(x_np)

        # JPEG 压缩
        buffer = io.BytesIO()
        img.save(buffer, format='JPEG', quality=quality)
        buffer.seek(0)
        img_compressed = Image.open(buffer)

        # 转回张量
        x_compressed = torch.from_numpy(
            np.array(img_compressed).astype(np.float32) / 255
        ).permute(2, 0, 1).unsqueeze(0)

        return x_compressed

    @staticmethod
    def bit_depth_reduction(x: torch.Tensor, bits: int = 4) -> torch.Tensor:
        """
        位深度降低

        Args:
            x: 输入
            bits: 目标位深度

        Returns:
            处理后的输入
        """
        max_val = 2 ** bits - 1
        x_reduced = torch.round(x * max_val) / max_val
        return x_reduced

    @staticmethod
    def spatial_smoothing(x: torch.Tensor,
                          kernel_size: int = 3) -> torch.Tensor:
        """
        空间平滑

        Args:
            x: 输入
            kernel_size: 核大小

        Returns:
            平滑后的输入
        """
        padding = kernel_size // 2
        kernel = torch.ones(1, 1, kernel_size, kernel_size) / (kernel_size ** 2)

        # 对每个通道应用平滑
        channels = []
        for c in range(x.shape[1]):
            channel = x[:, c:c+1, :, :]
            smoothed = F.conv2d(channel, kernel.to(x.device), padding=padding)
            channels.append(smoothed)

        return torch.cat(channels, dim=1)


class DefenseDistillation:
    """防御蒸馏"""

    def __init__(self, teacher_model: nn.Module,
                 student_model: nn.Module,
                 temperature: float = 20.0):
        """
        初始化防御蒸馏

        Args:
            teacher_model: 教师模型
            student_model: 学生模型
            temperature: 温度参数
        """
        self.teacher = teacher_model
        self.student = student_model
        self.temperature = temperature

    def distillation_loss(self, x: torch.Tensor,
                          y: torch.Tensor,
                          alpha: float = 0.5) -> torch.Tensor:
        """
        计算蒸馏损失

        Args:
            x: 输入
            y: 硬标签
            alpha: 软标签权重

        Returns:
            蒸馏损失
        """
        # 教师软标签
        with torch.no_grad():
            teacher_logits = self.teacher(x)
            soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)

        # 学生输出
        student_logits = self.student(x)

        # 软损失
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.temperature, dim=1),
            soft_targets,
            reduction='batchmean'
        ) * (self.temperature ** 2)

        # 硬损失
        hard_loss = F.cross_entropy(student_logits, y)

        return alpha * soft_loss + (1 - alpha) * hard_loss

    def train_step(self, x: torch.Tensor, y: torch.Tensor,
                   optimizer: torch.optim.Optimizer) -> float:
        """蒸馏训练步骤"""
        self.student.train()
        optimizer.zero_grad()

        loss = self.distillation_loss(x, y)
        loss.backward()
        optimizer.step()

        return loss.item()


class AdversarialDetector:
    """对抗样本检测器"""

    def __init__(self, model: nn.Module,
                 threshold: float = 0.5):
        """
        初始化检测器

        Args:
            model: 原始分类模型
            threshold: 检测阈值
        """
        self.model = model
        self.threshold = threshold
        self.statistics = {}

    def fit(self, clean_data: torch.Tensor):
        """
        在干净数据上计算统计量

        Args:
            clean_data: 干净样本
        """
        self.model.eval()

        with torch.no_grad():
            outputs = self.model(clean_data)
            probs = F.softmax(outputs, dim=1)

        # 计算置信度统计
        confidences = probs.max(dim=1)[0]
        self.statistics['mean_confidence'] = confidences.mean().item()
        self.statistics['std_confidence'] = confidences.std().item()

        # 计算熵统计
        entropies = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
        self.statistics['mean_entropy'] = entropies.mean().item()
        self.statistics['std_entropy'] = entropies.std().item()

    def detect(self, x: torch.Tensor) -> Tuple[bool, float]:
        """
        检测是否为对抗样本

        Args:
            x: 输入样本

        Returns:
            (是否为对抗样本, 异常分数)
        """
        self.model.eval()

        with torch.no_grad():
            outputs = self.model(x)
            probs = F.softmax(outputs, dim=1)

        # 置信度检测
        confidence = probs.max(dim=1)[0].item()
        conf_zscore = abs(confidence - self.statistics['mean_confidence']) / self.statistics['std_confidence']

        # 熵检测
        entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1).item()
        ent_zscore = abs(entropy - self.statistics['mean_entropy']) / self.statistics['std_entropy']

        # 综合分数
        anomaly_score = (conf_zscore + ent_zscore) / 2
        is_adversarial = anomaly_score > self.threshold

        return is_adversarial, anomaly_score

"""
偏见检测与公平性评估
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from sklearn.metrics import confusion_matrix
import scipy.stats as stats


@dataclass
class FairnessMetrics:
    """公平性指标结果"""
    statistical_parity_difference: float
    equal_opportunity_difference: float
    average_odds_difference: float
    predictive_parity_difference: float
    disparate_impact: float
    group_metrics: Dict[str, Dict[str, float]]


class BiasDetector:
    """偏见检测器"""

    def __init__(self, protected_attribute: str):
        """
        初始化偏见检测器

        Args:
            protected_attribute: 受保护属性名称
        """
        self.protected_attribute = protected_attribute

    def compute_fairness_metrics(self,
                                  y_true: np.ndarray,
                                  y_pred: np.ndarray,
                                  protected: np.ndarray) -> FairnessMetrics:
        """
        计算公平性指标

        Args:
            y_true: 真实标签
            y_pred: 预测标签
            protected: 受保护属性

        Returns:
            FairnessMetrics: 公平性指标
        """
        # 获取不同组
        groups = np.unique(protected)
        if len(groups) != 2:
            raise ValueError("目前只支持二元受保护属性")

        group_0_mask = protected == groups[0]
        group_1_mask = protected == groups[1]

        # 各组的预测率
        pr_0 = y_pred[group_0_mask].mean()
        pr_1 = y_pred[group_1_mask].mean()

        # 统计均等差异
        spd = pr_1 - pr_0

        # 各组的真正率 (TPR)
        tpr_0 = self._compute_tpr(y_true[group_0_mask], y_pred[group_0_mask])
        tpr_1 = self._compute_tpr(y_true[group_1_mask], y_pred[group_1_mask])

        # 机会均等差异
        eod = tpr_1 - tpr_0

        # 各组的假正率 (FPR)
        fpr_0 = self._compute_fpr(y_true[group_0_mask], y_pred[group_0_mask])
        fpr_1 = self._compute_fpr(y_true[group_1_mask], y_pred[group_1_mask])

        # 平均赔率差异
        aod = ((tpr_1 - tpr_0) + (fpr_1 - fpr_0)) / 2

        # 预测均等差异
        ppv_0 = self._compute_ppv(y_true[group_0_mask], y_pred[group_0_mask])
        ppv_1 = self._compute_ppv(y_true[group_1_mask], y_pred[group_1_mask])
        ppd = ppv_1 - ppv_0

        # 差异影响比
        di = pr_0 / pr_1 if pr_1 > 0 else float('inf')

        return FairnessMetrics(
            statistical_parity_difference=spd,
            equal_opportunity_difference=eod,
            average_odds_difference=aod,
            predictive_parity_difference=ppd,
            disparate_impact=di,
            group_metrics={
                str(groups[0]): {
                    'positive_rate': pr_0,
                    'tpr': tpr_0,
                    'fpr': fpr_0,
                    'ppv': ppv_0
                },
                str(groups[1]): {
                    'positive_rate': pr_1,
                    'tpr': tpr_1,
                    'fpr': fpr_1,
                    'ppv': ppv_1
                }
            }
        )

    def _compute_tpr(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        """计算真正率"""
        tp = np.sum((y_true == 1) & (y_pred == 1))
        fn = np.sum((y_true == 1) & (y_pred == 0))
        return tp / (tp + fn) if (tp + fn) > 0 else 0.0

    def _compute_fpr(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        """计算假正率"""
        fp = np.sum((y_true == 0) & (y_pred == 1))
        tn = np.sum((y_true == 0) & (y_pred == 0))
        return fp / (fp + tn) if (fp + tn) > 0 else 0.0

    def _compute_ppv(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        """计算阳性预测值"""
        tp = np.sum((y_true == 1) & (y_pred == 1))
        fp = np.sum((y_true == 0) & (y_pred == 1))
        return tp / (tp + fp) if (tp + fp) > 0 else 0.0

    def detect_bias(self,
                    y_true: np.ndarray,
                    y_pred: np.ndarray,
                    protected: np.ndarray,
                    threshold: float = 0.1) -> Dict[str, Any]:
        """
        检测偏见

        Args:
            y_true: 真实标签
            y_pred: 预测标签
            protected: 受保护属性
            threshold: 偏见阈值

        Returns:
            Dict: 偏见检测结果
        """
        metrics = self.compute_fairness_metrics(y_true, y_pred, protected)

        # 判断是否存在偏见
        biases = {
            'statistical_parity': abs(metrics.statistical_parity_difference) > threshold,
            'equal_opportunity': abs(metrics.equal_opportunity_difference) > threshold,
            'average_odds': abs(metrics.average_odds_difference) > threshold,
            'predictive_parity': abs(metrics.predictive_parity_difference) > threshold,
            'disparate_impact': metrics.disparate_impact < 0.8 or metrics.disparate_impact > 1.25
        }

        return {
            'has_bias': any(biases.values()),
            'bias_types': [k for k, v in biases.items() if v],
            'metrics': {
                'statistical_parity_difference': metrics.statistical_parity_difference,
                'equal_opportunity_difference': metrics.equal_opportunity_difference,
                'average_odds_difference': metrics.average_odds_difference,
                'predictive_parity_difference': metrics.predictive_parity_difference,
                'disparate_impact': metrics.disparate_impact
            },
            'group_metrics': metrics.group_metrics,
            'recommendation': self._generate_recommendation(biases, metrics)
        }

    def _generate_recommendation(self, biases: Dict[str, bool],
                                  metrics: FairnessMetrics) -> List[str]:
        """生成改进建议"""
        recommendations = []

        if biases['statistical_parity']:
            recommendations.append(
                f"模型在不同群体间的预测率差异为 {metrics.statistical_parity_difference:.3f}，"
                "考虑使用重采样或重加权方法"
            )

        if biases['equal_opportunity']:
            recommendations.append(
                f"真正率在群体间存在差异 ({metrics.equal_opportunity_difference:.3f})，"
                "考虑调整分类阈值或使用后处理方法"
            )

        if biases['disparate_impact']:
            recommendations.append(
                f"差异影响比为 {metrics.disparate_impact:.3f}，"
                "不满足 4/5 规则，建议审查特征选择"
            )

        return recommendations


class FairnessMitigation:
    """公平性缓解方法"""

    @staticmethod
    def reweighting(X: np.ndarray, y: np.ndarray,
                    protected: np.ndarray) -> np.ndarray:
        """
        重加权方法

        Args:
            X: 特征
            y: 标签
            protected: 受保护属性

        Returns:
            样本权重
        """
        n = len(y)
        weights = np.ones(n)

        groups = np.unique(protected)
        labels = np.unique(y)

        for g in groups:
            for l in labels:
                mask = (protected == g) & (y == l)
                n_gl = mask.sum()
                n_g = (protected == g).sum()
                n_l = (y == l).sum()

                expected = n_g * n_l / n
                weight = expected / n_gl if n_gl > 0 else 1.0
                weights[mask] = weight

        return weights

    @staticmethod
    def threshold_optimization(y_true: np.ndarray,
                               y_prob: np.ndarray,
                               protected: np.ndarray,
                               metric: str = 'equal_opportunity') -> Dict[int, float]:
        """
        阈值优化

        Args:
            y_true: 真实标签
            y_prob: 预测概率
            protected: 受保护属性
            metric: 目标指标

        Returns:
            各组的最优阈值
        """
        groups = np.unique(protected)
        thresholds = {}

        # 对每组寻找最优阈值
        for g in groups:
            mask = protected == g
            best_threshold = 0.5
            best_score = float('inf')

            for t in np.arange(0.1, 0.9, 0.01):
                y_pred = (y_prob[mask] >= t).astype(int)

                if metric == 'equal_opportunity':
                    # 最大化 TPR
                    tpr = np.sum((y_true[mask] == 1) & (y_pred == 1)) / max(np.sum(y_true[mask] == 1), 1)
                    score = -tpr  # 负数因为要最小化
                elif metric == 'statistical_parity':
                    # 目标预测率
                    target_rate = y_true.mean()
                    actual_rate = y_pred.mean()
                    score = abs(actual_rate - target_rate)

                if score < best_score:
                    best_score = score
                    best_threshold = t

            thresholds[g] = best_threshold

        return thresholds

    @staticmethod
    def adversarial_debiasing(X: np.ndarray, y: np.ndarray,
                               protected: np.ndarray,
                               adversary_weight: float = 1.0):
        """
        对抗去偏 - 训练时使用

        返回可用于 PyTorch 训练的对抗网络架构
        """
        import torch
        import torch.nn as nn

        class AdversarialDebiasingModel(nn.Module):
            def __init__(self, input_dim: int, hidden_dim: int = 128):
                super().__init__()

                # 主预测器
                self.predictor = nn.Sequential(
                    nn.Linear(input_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Linear(hidden_dim, hidden_dim),
                    nn.ReLU(),
                    nn.Linear(hidden_dim, 1),
                    nn.Sigmoid()
                )

                # 对抗器 (预测受保护属性)
                self.adversary = nn.Sequential(
                    nn.Linear(hidden_dim, hidden_dim // 2),
                    nn.ReLU(),
                    nn.Linear(hidden_dim // 2, 1),
                    nn.Sigmoid()
                )

                self.hidden_dim = hidden_dim

            def forward(self, x):
                # 中间表示
                hidden = x
                for layer in list(self.predictor.children())[:-2]:
                    hidden = layer(hidden)

                # 主任务预测
                prediction = self.predictor(x)

                # 对抗预测
                adversary_pred = self.adversary(hidden)

                return prediction, adversary_pred

        return AdversarialDebiasingModel(X.shape[1])