import os
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM
from modelscope import snapshot_download
import torch.nn.functional as F
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score
# 全局配置
CACHE_DIR = "D:\\modelscope\\hub"
os.environ["MODELSCOPE_CACHE"] = CACHE_DIR
# 设置设备（优先GPU）
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"使用设备：{DEVICE}")
    
def build_annotated_test_set():
    """
    构建标注好的意图识别测试集（真实场景样本）
    :return: DataFrame，包含text、true_intent列
    """
    # 标注好的测试样本（覆盖不同表达形式）
    test_data = [
        # 问候意图
        {"text": "你好呀，今天过得怎么样", "true_intent": "问候"},
        {"text": "嗨，好久不见", "true_intent": "问候"},
        {"text": "早上好，麻烦问个问题", "true_intent": "问候"},
        {"text": "大家好，我是新来的", "true_intent": "问候"},
        {"text": "哈喽，在吗", "true_intent": "问候"},
        
        # 天气查询意图
        {"text": "今天外面会不会下雨", "true_intent": "天气查询"},
        {"text": "明天的气温大概是多少度", "true_intent": "天气查询"},
        {"text": "出门需要带伞吗", "true_intent": "天气查询"},
        {"text": "未来三天的天气预报怎么样", "true_intent": "天气查询"},
        {"text": "现在外面冷不冷", "true_intent": "天气查询"},
        
        # 计算意图
        {"text": "帮我算一下5乘以8等于多少", "true_intent": "计算"},
        {"text": "100减去25是多少", "true_intent": "计算"},
        {"text": "9除以3的结果是啥", "true_intent": "计算"},
        {"text": "12加18等于几", "true_intent": "计算"},
        {"text": "算一下平方：5的平方是多少", "true_intent": "计算"},
        
        # 翻译意图
        {"text": "把“我爱中国”翻译成英文", "true_intent": "翻译"},
        {"text": "“Hello World”的中文意思是什么", "true_intent": "翻译"},
        {"text": "帮我翻译一句日语：こんにちは", "true_intent": "翻译"},
        {"text": "“谢谢”的英文怎么说", "true_intent": "翻译"},
        {"text": "翻译“明天见”成法语", "true_intent": "翻译"},
        
        # 告别意图
        {"text": "再见，下次再聊", "true_intent": "告别"},
        {"text": "我先撤了，拜拜", "true_intent": "告别"},
        {"text": "晚安，早点休息", "true_intent": "告别"},
        {"text": "就到这里吧，再见", "true_intent": "告别"},
        {"text": "回见，祝你愉快", "true_intent": "告别"},
        
        # 混合意图（测试边界）
        {"text": "你好，帮我算一下10加20", "true_intent": "计算"},
        {"text": "再见，顺便问下明天天气", "true_intent": "天气查询"},
        {"text": "帮我翻译“谢谢”，早上好", "true_intent": "翻译"},
    ]
    
    df = pd.DataFrame(test_data)
    # 重置索引，方便后续处理
    df = df.reset_index(drop=True)
    return df
# 构建测试集
test_df = build_annotated_test_set()
print("标注测试集构建完成，共{}条样本".format(len(test_df)))
print("测试集前5条：")
print(test_df.head())

# 大模型的意图列表
INTENT_LIST = ["问候", "天气查询", "计算", "翻译", "告别"]
def plot_evaluation_results(semantic_metrics, qwen_metrics, eval_df):
    """
    可视化评估结果（真实模型版）
    """
    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 1. 混淆矩阵对比图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    
    # 语义模型混淆矩阵
    sns.heatmap(
        semantic_metrics["confusion_matrix"],
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=INTENT_LIST,
        yticklabels=INTENT_LIST,
        ax=ax1
    )
    ax1.set_title("Qwen1.5-1.8B 语义向量模型混淆矩阵")
    ax1.set_xlabel("预测意图")
    ax1.set_ylabel("真实意图")
    
    # Qwen模型混淆矩阵
    sns.heatmap(
        qwen_metrics["confusion_matrix"],
        annot=True,
        fmt="d",
        cmap="Greens",
        xticklabels=INTENT_LIST,
        yticklabels=INTENT_LIST,
        ax=ax2
    )
    ax2.set_title("Qwen1.5-1.8B 混淆矩阵")
    ax2.set_xlabel("预测意图")
    ax2.set_ylabel("真实意图")
    
    plt.tight_layout()
    plt.savefig("Qwen1.5-1.8B 混淆矩阵 confusion_matrix_comparison.png", dpi=300, bbox_inches='tight')
    plt.show()

def calculate_intent_metrics(eval_df, true_col, pred_col, intent_list):
    """
    计算意图识别的核心指标（多分类场景）
    :param eval_df: 评估结果DataFrame
    :param true_col: 真实意图列名
    :param pred_col: 预测意图列名
    :param intent_list: 意图列表
    :return: 详细指标字典
    """
    # 1. 构建混淆矩阵
    cm = confusion_matrix(
        eval_df[true_col], 
        eval_df[pred_col], 
        labels=intent_list
    )
    
    # 2. 计算每个意图的精准率、召回率、F1值
    per_intent_metrics = {}
    for i, intent in enumerate(intent_list):
        # 提取该意图的TP、FP、FN
        TP = cm[i, i]
        FP = cm[:, i].sum() - TP  # 所有行的第i列之和 - TP
        FN = cm[i, :].sum() - TP  # 第i行的所有列之和 - TP
        TN = cm.sum() - TP - FP - FN
        
        # 计算指标（避免除以0）
        precision = TP / (TP + FP) if (TP + FP) > 0 else 0.0
        recall = TP / (TP + FN) if (TP + FN) > 0 else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
        
        per_intent_metrics[intent] = {
            "TP": TP,
            "FP": FP,
            "FN": FN,
            "TN": TN,
            "Precision": round(precision, 4),
            "Recall": round(recall, 4),
            "F1": round(f1, 4)
        }
    
    # 3. 计算宏平均和微平均
    # 宏平均：先算每个意图的指标，再平均
    macro_precision = np.mean([m["Precision"] for m in per_intent_metrics.values()])
    macro_recall = np.mean([m["Recall"] for m in per_intent_metrics.values()])
    macro_f1 = np.mean([m["F1"] for m in per_intent_metrics.values()])
    
    # 微平均：先汇总所有TP/FP/FN，再计算
    total_TP = sum([m["TP"] for m in per_intent_metrics.values()])
    total_FP = sum([m["FP"] for m in per_intent_metrics.values()])
    total_FN = sum([m["FN"] for m in per_intent_metrics.values()])
    
    micro_precision = total_TP / (total_TP + total_FP) if (total_TP + total_FP) > 0 else 0.0
    micro_recall = total_TP / (total_TP + total_FN) if (total_TP + total_FN) > 0 else 0.0
    micro_f1 = 2 * micro_precision * micro_recall / (micro_precision + micro_recall) if (micro_precision + micro_recall) > 0 else 0.0
    
    # 4. 计算整体准确率
    accuracy = (eval_df[true_col] == eval_df[pred_col]).sum() / len(eval_df)
    
    # 5. 计算响应时间指标（如果包含）
    rt_col = pred_col.replace("_intent", "_response_time_ms")
    if rt_col in eval_df.columns:
        rt_metrics = {
            "avg_response_time_ms": round(eval_df[rt_col].mean(), 2),
            "p95_response_time_ms": round(np.percentile(eval_df[rt_col], 95), 2),
            "max_response_time_ms": round(eval_df[rt_col].max(), 2),
            "min_response_time_ms": round(eval_df[rt_col].min(), 2)
        }
    else:
        rt_metrics = {}
    
    return {
        "confusion_matrix": cm,
        "per_intent": per_intent_metrics,
        "macro": {
            "Precision": round(macro_precision, 4),
            "Recall": round(macro_recall, 4),
            "F1": round(macro_f1, 4)
        },
        "micro": {
            "Precision": round(micro_precision, 4),
            "Recall": round(micro_recall, 4),
            "F1": round(micro_f1, 4)
        },
        "accuracy": round(accuracy, 4),
        "response_time": rt_metrics
    }
# 计算语义向量模型的指标
semantic_metrics = calculate_intent_metrics(
    eval_df, 
    true_col="true_intent", 
    pred_col="semantic_intent", 
    intent_list=INTENT_LIST
)
# 计算Qwen大模型的指标
qwen_metrics = calculate_intent_metrics(
    eval_df, 
    true_col="true_intent", 
    pred_col="qwen_intent", 
    intent_list=INTENT_LIST
)
# 打印指标结果
print("\n===== 语义向量模型评估指标 =====")
print(f"整体准确率：{semantic_metrics['accuracy']:.2%}")
print(f"宏平均F1值：{semantic_metrics['macro']['F1']:.2%}")
print(f"微平均F1值：{semantic_metrics['micro']['F1']:.2%}")
print(f"平均响应时间：{semantic_metrics['response_time']['avg_response_time_ms']} ms")
print(f"95分位响应时间：{semantic_metrics['response_time']['p95_response_time_ms']} ms")
print("\n===== Qwen大模型评估指标 =====")
print(f"整体准确率：{qwen_metrics['accuracy']:.2%}")
print(f"宏平均F1值：{qwen_metrics['macro']['F1']:.2%}")
print(f"微平均F1值：{qwen_metrics['micro']['F1']:.2%}")
print(f"平均响应时间：{qwen_metrics['response_time']['avg_response_time_ms']} ms")
print(f"95分位响应时间：{qwen_metrics['response_time']['p95_response_time_ms']} ms")
# 使用之前计算的语义和Qwen指标（semantic_metrics和qwen_metrics已在前面计算）
print("\n" + "="*50)
print("                    对比结果摘要")
print("="*50)
print(f"{'模型':<15} | {'准确率':<8} | {'Macro F1':<8} | {'平均耗时(ms)':<12}")
print("-" * 55)
print(f"{'语义向量':<15} | {semantic_metrics['accuracy']:.2%}     | {semantic_metrics['macro']['F1']:.2f}     | {semantic_metrics['response_time']['avg_response_time_ms']:.2f}")
print(f"{'Qwen1.5-1.8B':<15} | {qwen_metrics['accuracy']:.2%}     | {qwen_metrics['macro']['F1']:.2f}     | {qwen_metrics['response_time']['avg_response_time_ms']:.2f}")
print("="*50)
# 打印各意图详细指标
print("\n各意图详细指标（语义向量模型）:")
print(f"{'意图':<12} | {'精准率':<10} | {'召回率':<10} | {'F1值':<10} | {'TP':<8} | {'FP':<8}| {'FN':<8}")
print("-" * 50)
for intent, metrics in semantic_metrics["per_intent"].items():
    print(f"{intent:<12} | {metrics['Precision']:.2%}      | {metrics['Recall']:.2%}      | {metrics['F1']:.2%}      | {metrics['TP']:<8}      | {metrics['FP']:<8}      | {metrics['FN']:<8}")
print("\n各意图详细指标（Qwen1.5-1.8B）:")
print(f"{'意图':<12} | {'精准率':<10} | {'召回率':<10} | {'F1值':<10} | {'TP':<8} | {'FP':<8}| {'FN':<8}")
print("-" * 50)
for intent, metrics in qwen_metrics["per_intent"].items():
    print(f"{intent:<12} | {metrics['Precision']:.2%}      | {metrics['Recall']:.2%}      | {metrics['F1']:.2%}      | {metrics['TP']:<8}      | {metrics['FP']:<8}      | {metrics['FN']:<8}")

# 5. 计算响应时间指标（如果包含）
    rt_col = pred_col.replace("_intent", "_response_time_ms")
    if rt_col in eval_df.columns:
        rt_metrics = {
            "avg_response_time_ms": round(eval_df[rt_col].mean(), 2),
            "p95_response_time_ms": round(np.percentile(eval_df[rt_col], 95), 2),
            "max_response_time_ms": round(eval_df[rt_col].max(), 2),
            "min_response_time_ms": round(eval_df[rt_col].min(), 2)
        }
    else:
        rt_metrics = {}
# 2. 响应时间对比图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    
    # 响应时间分布直方图
    ax1.hist(eval_df["semantic_response_time_ms"], bins=10, alpha=0.7, label="语义向量模型", color="blue")
    ax1.hist(eval_df["qwen_response_time_ms"], bins=10, alpha=0.7, label="Qwen大模型", color="green")
    ax1.axvline(eval_df["semantic_response_time_ms"].mean(), color="blue", linestyle="--", label=f"语义模型均值：{eval_df['semantic_response_time_ms'].mean():.2f}ms")
    ax1.axvline(eval_df["qwen_response_time_ms"].mean(), color="green", linestyle="--", label=f"Qwen模型均值：{eval_df['qwen_response_time_ms'].mean():.2f}ms")
    ax1.set_xlabel("响应时间（ms）")
    ax1.set_ylabel("样本数")
    ax1.set_title("Qwen1.5-1.8B 响应时间分布对比")
    ax1.legend()
    ax1.grid(alpha=0.3)
    
    # 各意图平均响应时间对比
    intent_rt_semantic = eval_df.groupby("true_intent")["semantic_response_time_ms"].mean()
    intent_rt_qwen = eval_df.groupby("true_intent")["qwen_response_time_ms"].mean()
    
    x = np.arange(len(INTENT_LIST))
    width = 0.35
    
    ax2.bar(x - width/2, intent_rt_semantic[INTENT_LIST], width, label="语义向量模型", color="blue")
    ax2.bar(x + width/2, intent_rt_qwen[INTENT_LIST], width, label="Qwen大模型", color="green")
    
    ax2.set_xlabel("意图类型")
    ax2.set_ylabel("平均响应时间（ms）")
    ax2.set_title("Qwen1.5-1.8B 各意图平均响应时间对比")
    ax2.set_xticks(x)
    ax2.set_xticklabels(INTENT_LIST)
    ax2.legend()
    ax2.grid(alpha=0.3)
    
    plt.tight_layout()
    plt.savefig("Qwen1.5-1.8B 各意图平均响应时间对比 response_time_comparison.png", dpi=300, bbox_inches='tight')
    plt.show()

# 3. F1值对比图
    fig, ax = plt.subplots(figsize=(12, 6))
    
    # 提取各意图的F1值
    semantic_f1 = [semantic_metrics["per_intent"][intent]["F1"] for intent in INTENT_LIST]
    qwen_f1 = [qwen_metrics["per_intent"][intent]["F1"] for intent in INTENT_LIST]
    
    x = np.arange(len(INTENT_LIST))
    width = 0.35
    
    ax.bar(x - width/2, semantic_f1, width, label="语义向量模型", color="blue")
    ax.bar(x + width/2, qwen_f1, width, label="Qwen大模型", color="green")
    
    # 添加数值标签
    for i, v in enumerate(semantic_f1):
        ax.text(i - width/2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9)
    for i, v in enumerate(qwen_f1):
        ax.text(i + width/2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9)
    
    ax.set_xlabel("意图类型")
    ax.set_ylabel("F1值")
    ax.set_title("Qwen1.5-1.8B 各意图F1值对比")
    ax.set_xticks(x)
    ax.set_xticklabels(INTENT_LIST)
    ax.set_ylim(0, 1.1)
    ax.legend()
    ax.grid(alpha=0.3)
    
    plt.tight_layout()
    plt.savefig("Qwen1.5-1.8B 各意图F1值对比 f1_score_comparison.png", dpi=300, bbox_inches='tight')
    plt.show()