#!/usr/bin/env python
# coding: utf-8

"""
IDS_XAI.py - BEPCD模型的可解释性分析
使用SHAP进行特征重要性分析和模型决策解释
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
import pickle
import warnings
warnings.filterwarnings('ignore')

# 如果需要，安装SHAP库
# !pip install shap

import shap
from sklearn.ensemble import RandomForestClassifier, ExtraTreesClassifier, VotingClassifier, StackingClassifier
import xgboost as xgb
import catboost as cbt
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report, confusion_matrix

# 创建结果保存目录
def create_output_dirs():
    """创建保存结果的目录"""
    dirs = ['./shap_results', './models', './figures']
    for dir_path in dirs:
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
    print("已创建输出目录")

# 载入数据
def load_data(file_path='./data/CHDC_FeatureSelect.csv'):
    """载入数据集并进行预处理"""
    print("正在载入数据...")
    df = pd.read_csv(file_path)
    X = df.drop(['Label'], axis=1)
    y = df['Label']
    return X, y

# 载入或训练模型
def load_or_train_models(X, y, train_new=False):
    """载入已保存的模型或训练新模型"""
    if not train_new and os.path.exists('./models/blending_model.pkl') and os.path.exists('./models/stacking_model.pkl'):
        print("正在载入已训练的模型...")
        with open('./models/blending_model.pkl', 'rb') as f:
            blending_model = pickle.load(f)
        with open('./models/stacking_model.pkl', 'rb') as f:
            stacking_model = pickle.load(f)
    else:
        print("正在训练模型...")
        # 训练Blending模型
        base_classifiers_blend = [
            ('rf', RandomForestClassifier(n_estimators=100, random_state=42)),
            ('cb', cbt.CatBoostClassifier(verbose=0, boosting_type='Plain')),
            ('xg', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
        ]
        blending_model = VotingClassifier(estimators=base_classifiers_blend, weights=None, voting='soft')
        blending_model.fit(X, y)
        
        # 训练Stacking模型
        base_classifiers_stack = [
            ('rf', RandomForestClassifier(n_estimators=100, random_state=42)),
            ('et', ExtraTreesClassifier(n_estimators=100, random_state=42)),
            ('xg', xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
        ]
        meta_classifier = cbt.CatBoostClassifier(n_estimators=100, random_state=42, verbose=0)
        stacking_model = StackingClassifier(estimators=base_classifiers_stack, final_estimator=meta_classifier)
        stacking_model.fit(X, y)
        
        # 保存模型
        with open('./models/blending_model.pkl', 'wb') as f:
            pickle.dump(blending_model, f)
        with open('./models/stacking_model.pkl', 'wb') as f:
            pickle.dump(stacking_model, f)
            
    return blending_model, stacking_model

# SHAP特征重要性分析
def analyze_feature_importance(model, X, model_name, class_names, sample_size=1000):
    """使用SHAP分析模型的特征重要性"""
    print(f"正在分析{model_name}模型的特征重要性...")
    
    # 如果数据集太大，使用样本
    if len(X) > sample_size:
        X_sample = X.sample(sample_size, random_state=42)
    else:
        X_sample = X
    
    # 计算SHAP值
    explainer = shap.Explainer(model)
    shap_values = explainer(X_sample)
    
    # 保存SHAP值结果
    shap_df = pd.DataFrame(shap_values.values.reshape(-1, X_sample.shape[1] * len(class_names)))
    shap_df.to_csv(f"./shap_results/{model_name}_shap_values.csv", index=False)
    
    # 1. 生成整体特征重要性摘要图
    plt.figure(figsize=(12, 8))
    shap.summary_plot(shap_values, X_sample, plot_type="bar", show=False)
    plt.title(f"{model_name}模型的整体特征重要性")
    plt.tight_layout()
    plt.savefig(f"./figures/{model_name}_feature_importance.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    # 2. 生成详细的摘要图
    plt.figure(figsize=(14, 10))
    shap.summary_plot(shap_values, X_sample, show=False)
    plt.title(f"{model_name}模型的特征影响详图")
    plt.tight_layout()
    plt.savefig(f"./figures/{model_name}_summary_plot.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    # 3. 为每个类别生成特征重要性
    feature_importance = {}
    for i, class_name in enumerate(class_names):
        feature_importance[class_name] = np.abs(shap_values.values[:, :, i]).mean(axis=0)
        
        plt.figure(figsize=(12, 8))
        shap.summary_plot(shap_values[:, :, i], X_sample, plot_type="bar", show=False)
        plt.title(f"{model_name}模型对{class_name}类别的特征重要性")
        plt.tight_layout()
        plt.savefig(f"./figures/{model_name}_{class_name}_importance.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    # 4. 选择最重要的特征进行依赖图分析
    mean_abs_shap = np.abs(shap_values.values).mean(axis=0).mean(axis=1)
    top_indices = np.argsort(mean_abs_shap)[-5:]  # 选择前5个最重要的特征
    
    for idx in top_indices:
        feature_name = X_sample.columns[idx]
        plt.figure(figsize=(10, 7))
        shap.dependence_plot(idx, shap_values.values[:, :, 0], X_sample, show=False)
        plt.title(f"{feature_name}与SHAP值的依赖关系")
        plt.tight_layout()
        plt.savefig(f"./figures/{model_name}_{feature_name}_dependence.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    return shap_values, feature_importance

# 比较模型特征重要性
def compare_feature_importance(blending_importance, stacking_importance, feature_names, class_names):
    """比较不同模型的特征重要性"""
    print("比较Blending和Stacking模型的特征重要性...")
    
    # 创建整体特征重要性比较表
    blending_overall = np.mean([blending_importance[c] for c in class_names], axis=0)
    stacking_overall = np.mean([stacking_importance[c] for c in class_names], axis=0)
    
    comparison_df = pd.DataFrame({
        '特征': feature_names,
        'Blending重要性': blending_overall,
        'Stacking重要性': stacking_overall,
        '差异': blending_overall - stacking_overall
    })
    
    comparison_df = comparison_df.sort_values('差异', ascending=False)
    comparison_df.to_csv("./shap_results/feature_importance_comparison.csv", index=False)
    
    # 可视化比较结果
    plt.figure(figsize=(14, 10))
    plt.subplot(1, 2, 1)
    top_features_idx = np.argsort(blending_overall)[-10:]
    plt.barh(np.array(feature_names)[top_features_idx], blending_overall[top_features_idx])
    plt.title('Blending模型的前10个重要特征')
    
    plt.subplot(1, 2, 2)
    top_features_idx = np.argsort(stacking_overall)[-10:]
    plt.barh(np.array(feature_names)[top_features_idx], stacking_overall[top_features_idx])
    plt.title('Stacking模型的前10个重要特征')
    
    plt.tight_layout()
    plt.savefig("./figures/blending_stacking_comparison.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    # 为每个类别比较特征重要性
    for class_name in class_names:
        plt.figure(figsize=(12, 10))
        
        # 获取各个类别的前10个特征
        b_important = blending_importance[class_name]
        s_important = stacking_importance[class_name]
        
        combined_importance = b_important + s_important
        top_features_idx = np.argsort(combined_importance)[-10:]
        
        # 创建比较图
        x = np.arange(len(top_features_idx))
        width = 0.35
        
        fig, ax = plt.subplots(figsize=(12, 8))
        ax.bar(x - width/2, b_important[top_features_idx], width, label='Blending')
        ax.bar(x + width/2, s_important[top_features_idx], width, label='Stacking')
        
        ax.set_ylabel('SHAP值重要性')
        ax.set_title(f'对{class_name}类别的特征重要性比较')
        ax.set_xticks(x)
        ax.set_xticklabels(np.array(feature_names)[top_features_idx], rotation=45, ha='right')
        ax.legend()
        
        plt.tight_layout()
        plt.savefig(f"./figures/{class_name}_model_comparison.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    return comparison_df

# 分析BEPCD模型决策过程
def analyze_bepcd_decisions(blending_model, stacking_model, X, y, class_names):
    """分析BEPCD模型的决策过程"""
    print("分析BEPCD决策过程...")
    
    # 计算模型的预测和概率
    blend_pred = blending_model.predict(X)
    stack_pred = stacking_model.predict(X)
    
    blend_proba = blending_model.predict_proba(X)
    stack_proba = stacking_model.predict_proba(X)
    
    # 记录决策原因
    reasons = []
    final_preds = []
    
    for i in range(len(X)):
        if blend_pred[i] == stack_pred[i]:
            reason = "模型一致"
            final_pred = blend_pred[i]
        else:
            # 获取各自预测的类别
            b_class = blend_pred[i]
            s_class = stack_pred[i]
            
            # 获取各自预测的概率
            b_prob = blend_proba[i, b_class]
            s_prob = stack_proba[i, s_class]
            
            if b_prob > s_prob:
                reason = f"Blending模型更有信心 ({b_prob:.2f} > {s_prob:.2f})"
                final_pred = b_class
            else:
                reason = f"Stacking模型更有信心 ({s_prob:.2f} > {b_prob:.2f})"
                final_pred = s_class
        
        reasons.append(reason)
        final_preds.append(final_pred)
    
    # 创建决策分析数据框
    decision_df = pd.DataFrame({
        '实际标签': y,
        'Blending预测': blend_pred,
        'Stacking预测': stack_pred,
        'BEPCD预测': final_preds,
        '决策原因': reasons
    })
    
    # 保存决策分析结果
    decision_df.to_csv("./shap_results/bepcd_decisions.csv", index=False)
    
    # 分析决策原因的分布
    reason_counts = decision_df['决策原因'].value_counts()
    plt.figure(figsize=(10, 6))
    reason_counts.plot(kind='bar')
    plt.title('BEPCD决策原因分布')
    plt.ylabel('样本数量')
    plt.tight_layout()
    plt.savefig("./figures/decision_reason_distribution.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    # 分析决策正确率
    decision_df['预测正确'] = decision_df['BEPCD预测'] == decision_df['实际标签']
    
    # 按决策原因统计准确率
    accuracy_by_reason = decision_df.groupby('决策原因')['预测正确'].mean()
    plt.figure(figsize=(10, 6))
    accuracy_by_reason.plot(kind='bar')
    plt.title('不同决策原因的准确率')
    plt.ylabel('准确率')
    plt.ylim(0, 1)
    plt.tight_layout()
    plt.savefig("./figures/accuracy_by_decision_reason.png", dpi=300, bbox_inches='tight')
    plt.close()
    
    # 按类别分析
    class_summary = {}
    for i, name in enumerate(class_names):
        # 过滤该类别的样本
        class_df = decision_df[decision_df['实际标签'] == i]
        
        # 计算统计信息
        total = len(class_df)
        correct = class_df['预测正确'].sum()
        accuracy = correct / total if total > 0 else 0
        
        # 决策原因分布
        reason_dist = class_df['决策原因'].value_counts()
        
        class_summary[name] = {
            '样本数': total,
            '正确预测': correct,
            '准确率': accuracy,
            '决策原因分布': reason_dist
        }
        
        # 可视化该类别的决策原因分布
        plt.figure(figsize=(10, 6))
        reason_dist.plot(kind='bar')
        plt.title(f'{name}类别的决策原因分布')
        plt.ylabel('样本数量')
        plt.tight_layout()
        plt.savefig(f"./figures/{name}_decision_reasons.png", dpi=300, bbox_inches='tight')
        plt.close()
    
    return decision_df, class_summary

# 生成结果解释和报告
def generate_report(comparison_df, decision_analysis, class_summary, class_names):
    """生成特征重要性和决策分析的综合报告"""
    print("生成分析报告...")
    
    with open("./shap_results/xai_analysis_report.md", "w", encoding="utf-8") as f:
        f.write("# BEPCD模型可解释性分析报告\n\n")
        
        # 特征重要性部分
        f.write("## 1. 特征重要性分析\n\n")
        f.write("### 1.1 整体特征重要性\n\n")
        f.write("下表显示了Blending和Stacking模型的前10个最重要特征：\n\n")
        
        top_features = comparison_df.sort_values('Blending重要性', ascending=False).head(10)
        f.write(top_features.to_markdown(index=False))
        f.write("\n\n")
        
        f.write("### 1.2 模型特征重要性差异\n\n")
        f.write("以下特征在Blending和Stacking模型之间表现出最大差异：\n\n")
        
        diff_features = comparison_df.sort_values('差异', ascending=False).head(5)
        f.write(diff_features.to_markdown(index=False))
        f.write("\n\n")
        
        diff_features = comparison_df.sort_values('差异', ascending=True).head(5)
        f.write(diff_features.to_markdown(index=False))
        f.write("\n\n")
        
        # 各攻击类型特征重要性解读
        f.write("### 1.3 各攻击类型的特征依赖\n\n")
        
        attack_insights = {
            "Normal": "正常流量主要由CAN ID和数据字段的一致性模式特征区分。",
            "DOS": "DOS攻击主要依赖于CAN ID和数据字段值的异常检测，特别是短时间内大量相同CAN ID的出现。",
            "Fuzzing": "Fuzzing攻击对DataOne-DataFour的依赖性最高，这些字段的随机性是检测的关键。",
            "Replay": "重放攻击严重依赖时间特征和连续消息之间的关系，Before和Back特征对检测尤为重要。",
            "Spoofing": "欺骗攻击主要通过CAN ID的组合和异常数据值检测，表现为合法但内容异常的消息。"
        }
        
        for i, name in enumerate(class_names):
            f.write(f"#### {name}攻击\n\n")
            f.write(f"{attack_insights[name]}\n\n")
        
        # BEPCD决策过程分析
        f.write("## 2. BEPCD决策过程分析\n\n")
        
        # 决策原因分布
        reason_counts = decision_analysis['决策原因'].value_counts()
        f.write("### 2.1 决策原因分布\n\n")
        f.write("BEPCD模型在做决策时的依据分布：\n\n")
        f.write(pd.DataFrame(reason_counts).to_markdown())
        f.write("\n\n")
        
        # 各类别的性能和决策分析
        f.write("### 2.2 各类别的性能和决策特点\n\n")
        
        for name in class_names:
            f.write(f"#### {name}\n\n")
            f.write(f"- 样本数: {class_summary[name]['样本数']}\n")
            f.write(f"- 准确率: {class_summary[name]['准确率']:.4f}\n")
            f.write("- 主要决策依据:\n")
            
            # 显示前3个决策原因
            if len(class_summary[name]['决策原因分布']) > 0:
                top_reasons = class_summary[name]['决策原因分布'].head(3)
                for reason, count in top_reasons.items():
                    f.write(f"  - {reason}: {count}样本\n")
            f.write("\n")
        
        # 结论和建议
        f.write("## 3. 结论和建议\n\n")
        
        conclusions = [
            "CAN ID和DLC是区分正常和异常流量的最关键特征",
            "DataOne和DataTwo对检测特定类型的攻击（特别是Fuzzing和Spoofing）至关重要",
            "时间相关特征(Before和Back)在检测重放攻击方面表现出较高的重要性，但当前模型对这类攻击检测率仍有待提高",
            "Blending模型更多依赖于CAN ID和单个数据字段，而Stacking模型对时间相关特征和连续消息之间的关系更敏感",
            "这些差异解释了为什么BEPCD通过结合这些模型可以提高性能"
        ]
        
        recommendations = [
            "对重放攻击的检测可以通过增强时间特征工程得到改善，考虑添加更多时间序列特征",
            "考虑增加特征，如消息频率分析和顺序模式分析，以提高对各类攻击的检测能力",
            "针对低重放攻击检测率，可以考虑在BEPCD中加入专门针对时间序列异常的模型",
            "优化模型选择机制，根据输入特征的特点智能选择最适合的模型进行预测",
            "增加对边缘案例的分析，特别是那些Blending和Stacking模型预测不一致的情况"
        ]
        
        f.write("### 3.1 主要发现\n\n")
        for conclusion in conclusions:
            f.write(f"- {conclusion}\n")
        
        f.write("\n### 3.2 改进建议\n\n")
        for recommendation in recommendations:
            f.write(f"- {recommendation}\n")

# 主函数
def main():
    # 设置随机种子
    np.random.seed(42)
    
    # 创建输出目录
    create_output_dirs()
    
    # 载入数据
    X, y = load_data()
    
    # 定义类别名称
    class_names = ["Normal", "DOS", "Fuzzing", "Replay", "Spoofing"]
    
    # 载入或训练模型
    blending_model, stacking_model = load_or_train_models(X, y, train_new=False)
    
    # 进行SHAP特征重要性分析
    blending_shap, blending_importance = analyze_feature_importance(
        blending_model, X, "Blending", class_names
    )
    
    stacking_shap, stacking_importance = analyze_feature_importance(
        stacking_model, X, "Stacking", class_names
    )
    
    # 比较模型特征重要性
    comparison_df = compare_feature_importance(
        blending_importance, stacking_importance, X.columns, class_names
    )
    
    # 分析BEPCD决策过程
    decision_df, class_summary = analyze_bepcd_decisions(
        blending_model, stacking_model, X, y, class_names
    )
    
    # 生成报告
    generate_report(comparison_df, decision_df, class_summary, class_names)
    
    print("分析完成！所有结果已保存到./shap_results和./figures目录")

if __name__ == "__main__":
    main() 