import torch
import numpy as np
from scipy.spatial.distance import cosine
from modelscope.hub.snapshot_download import snapshot_download
from transformers import AutoTokenizer, AutoModel
import time
from typing import List, Dict, Tuple
# ===================== 全局配置 =====================
CACHE_DIR = "D:\\modelscope\\hub"
# 推荐使用多语言模型以支持跨境客服，此处演示使用轻量级英文模型
MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
TORCH_DTYPE = torch.float16
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ===================== 核心工具函数 =====================
def load_model_and_tokenizer(model_id: str, cache_dir: str) -> Tuple[AutoTokenizer, AutoModel]:
    """加载模型和分词器，自动处理下载和显存映射"""
    print(f"正在下载/加载模型：{model_id}")
    local_path = snapshot_download(model_id, cache_dir=cache_dir)
    
    tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(
        local_path, 
        trust_remote_code=True, 
        torch_dtype=TORCH_DTYPE, 
        device_map="auto"
    ).to(DEVICE)
    model.eval()
    return tokenizer, model
def mean_pooling(model_output, attention_mask):
    """均值池化：将Token级嵌入聚合为句子级嵌入并归一化"""
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).to(TORCH_DTYPE)
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, dim=1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    mean_embeddings = sum_embeddings / sum_mask
    return torch.nn.functional.normalize(mean_embeddings, p=2, dim=1)
def get_embeddings(texts: List[str], tokenizer: AutoTokenizer, model: AutoModel, batch_size: int = 32) -> np.ndarray:
    """批量获取文本嵌入向量（核心推理引擎）"""
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i:i+batch_size]
        encoded_input = tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(DEVICE)
        
        with torch.no_grad():
            model_output = model(**encoded_input)
        
        batch_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
        all_embeddings.extend(batch_embeddings.cpu().numpy())
    return np.array(all_embeddings)
def calculate_similarity(emb1: np.ndarray, emb2: np.ndarray) -> float:
    """计算余弦相似度 (1 - 余弦距离)"""
    return 1 - cosine(emb1, emb2)

import numpy as np
import time
from typing import List, Dict
# 假设第一段中的函数 (load_model..., get_embeddings, calculate_similarity) 已在此作用域可用
class SmartCustomerService:
    def __init__(self, model_id: str, cache_dir: str):
        """初始化：加载模型并准备存储结构"""
        self.tokenizer, self.model = load_model_and_tokenizer(model_id, cache_dir)
        self.knowledge_base = []  # 存储原始问答对
        self.knowledge_embeddings = None  # 存储预计算的向量矩阵
        self.query_times = []  # 性能监控列表
    
    def build_knowledge_base(self, knowledge_data: List[Dict]):
        """构建知识库：预计算所有问题的向量索引"""
        print(f"\n正在构建知识库，共{len(knowledge_data)}条问答对")
        self.knowledge_base = knowledge_data
        
        # 提取问题列表
        questions = [item["question"] for item in self.knowledge_base]
        
        # 【关键优化】预计算向量，避免每次查询都重复计算知识库
        start_time = time.time()
        self.knowledge_embeddings = get_embeddings(questions, self.tokenizer, self.model)
        end_time = time.time()
        
        print(f"知识库构建完成！耗时：{end_time - start_time:.2f}秒")
        print(f"向量维度：{self.knowledge_embeddings.shape}")
    def search_similar_questions(self, user_query: str, top_k: int = 3) -> List[Dict]:
        """核心检索：计算用户查询与知识库的相似度并排序"""
        start_time = time.time()
        
        # 1. 实时计算用户查询的向量
        query_embedding = get_embeddings([user_query], self.tokenizer, self.model)[0]
        
        # 2. 遍历计算相似度 (实际生产环境可用矩阵运算加速)
        similarities = []
        for idx, kb_embedding in enumerate(self.knowledge_embeddings):
            sim = calculate_similarity(query_embedding, kb_embedding)
            similarities.append({
                "index": idx,
                "similarity": sim,
                "question": self.knowledge_base[idx]["question"],
                "answer": self.knowledge_base[idx]["answer"],
                "language": self.knowledge_base[idx]["language"]
            })
        
        # 3. 按相似度降序排序
        similarities.sort(key=lambda x: x["similarity"], reverse=True)
        
        # 记录性能
        self.query_times.append(time.time() - start_time)
        return similarities[:top_k]
    def answer_query(self, user_query: str, top_k: int = 3) -> Dict:
        """对外接口：执行检索并格式化输出"""
        results = self.search_similar_questions(user_query, top_k)
        avg_time = round(np.mean(self.query_times) * 1000, 2) if self.query_times else 0
        
        print(f"\n=== 用户问题：{user_query} ===")
        print(f"平均响应时间：{avg_time}ms")
        for i, res in enumerate(results, 1):
            print(f"{i}. [相似度:{res['similarity']:.4f}] {res['question']} -> {res['answer']}")
        
        return {"query": user_query, "results": results, "time_ms": avg_time}

import matplotlib.pyplot as plt
import matplotlib
import seaborn as sns
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import pandas as pd
import numpy as np
# 设置中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
matplotlib.rcParams['axes.unicode_minus'] = False
# 假设 SmartCustomerService 类已定义
# 以下方法添加到 SmartCustomerService 类中，或作为独立函数调用类实例
def visualize_embeddings(self, save_path: str = "pca_distribution.png"):
    """PCA降维可视化：观察知识库问题的宏观分布"""
    if self.knowledge_embeddings is None: return
    pca = PCA(n_components=2)
    embeddings_2d = pca.fit_transform(self.knowledge_embeddings)
    
    df = pd.DataFrame({
        'x': embeddings_2d[:, 0], 'y': embeddings_2d[:, 1],
        'language': [item['language'] for item in self.knowledge_base],
        'text': [item['question'][:15] + '..' for item in self.knowledge_base]
    })
    
    plt.figure(figsize=(12, 8))
    for lang in df['language'].unique():
        subset = df[df['language'] == lang]
        plt.scatter(subset['x'], subset['y'], label=lang, alpha=0.7, s=100)
        # 添加文字标注
        for _, row in subset.iterrows():
            plt.text(row['x'], row['y'], row['text'], fontsize=8, alpha=0.6)
            
    plt.title(f'知识库语义分布 (PCA)\n方差解释率: {pca.explained_variance_ratio_.sum()*100:.2f}%')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig(save_path, dpi=300)
    print(f"✓ PCA图已保存: {save_path}")
    plt.close()
def visualize_tsne(self, save_path: str = "tsne_clusters.png"):
    """t-SNE降维可视化：观察局部聚类效果"""
    if self.knowledge_embeddings is None: return
    # perplexity需小于样本数
    perplexity = min(30, len(self.knowledge_embeddings) - 1)
    tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity)
    embeddings_2d = tsne.fit_transform(self.knowledge_embeddings)
    
    df = pd.DataFrame({
        'x': embeddings_2d[:, 0], 'y': embeddings_2d[:, 1],
        'language': [item['language'] for item in self.knowledge_base]
    })
    
    plt.figure(figsize=(12, 8))
    sns.scatterplot(data=df, x='x', y='y', hue='language', palette='Set2', s=150, edgecolor='white', linewidth=1.5)
    plt.title('知识库语义聚类 (t-SNE)')
    plt.legend(title='Language')
    plt.grid(True, alpha=0.3)
    plt.savefig(save_path, dpi=300)
    print(f"✓ t-SNE图已保存: {save_path}")
    plt.close()
# 动态绑定方法到类
SmartCustomerService.visualize_embeddings = visualize_embeddings
SmartCustomerService.visualize_tsne = visualize_tsne

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
# 以下方法继续扩展 SmartCustomerService 类
def visualize_similarity_heatmap(self, save_path: str = "similarity_heatmap.png"):
    """相似度热力图：检测知识库内部的冗余或冲突"""
    if self.knowledge_embeddings is None: return
    n = len(self.knowledge_embeddings)
    matrix = np.zeros((n, n))
    
    # 计算两两相似度矩阵
    for i in range(n):
        for j in range(n):
            matrix[i, j] = calculate_similarity(self.knowledge_embeddings[i], self.knowledge_embeddings[j])
    
    # 生成短标签
    labels = [q[:12] + '..' for q in [item['question'] for item in self.knowledge_base]]
    
    plt.figure(figsize=(14, 12))
    sns.heatmap(matrix, xticklabels=labels, yticklabels=labels, cmap='YlOrRd', annot=True, fmt='.2f')
    plt.title('知识库问题内部相似度热力图')
    plt.xticks(rotation=45)
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    print(f"✓ 热力图已保存: {save_path}")
    plt.close()
def visualize_query_performance(self, save_path: str = "performance_stats.png"):
    """性能监控：响应时间趋势与分布"""
    if not self.query_times: return
    
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    times_ms = [t * 1000 for t in self.query_times]
    
    # 1. 趋势图
    axes[0, 0].plot(times_ms, marker='o', color='#1890ff')
    axes[0, 0].set_title('查询响应时间趋势 (ms)')
    axes[0, 0].grid(True, alpha=0.3)
    
    # 2. 直方图
    axes[0, 1].hist(times_ms, bins=15, color='#52c41a', edgecolor='black', alpha=0.7)
    axes[0, 1].axvline(np.mean(times_ms), color='red', linestyle='--', label='平均值')
    axes[0, 1].set_title('响应时间分布')
    axes[0, 1].legend()
    
    # 3. 统计指标
    stats = [np.mean(times_ms), np.min(times_ms), np.max(times_ms)]
    axes[1, 0].bar(['平均', '最快', '最慢'], stats, color=['#1890ff', '#52c41a', '#ff4d4f'])
    axes[1, 0].set_title('关键性能指标 (ms)')
    
    # 4. 饼图 (模拟效率占比)
    avg = np.mean(times_ms)
    axes[1, 1].pie([avg, max(0, 200-avg)], labels=['消耗时间', '剩余预算(200ms)'], colors=['#faad14', '#e8e8e8'], autopct='%1.1f%%')
    axes[1, 1].set_title('延迟预算分析')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    print(f"✓ 性能图已保存: {save_path}")
    plt.close()
def visualize_search_results(self, user_query: str, results: List[Dict], save_path: str = "search_result.png"):
    """单次搜索结果可视化"""
    if not results: return
    sims = [r['similarity']*100 for r in results]
    labels = [f"{r['question'][:20]}..." for r in results]
    
    plt.figure(figsize=(10, 6))
    plt.barh(labels, sims, color=plt.cm.RdYlGn_r(np.linspace(0.3, 0.9, len(results))))
    plt.xlabel('相似度 (%)')
    plt.title(f'查询: "{user_query}" 的匹配结果')
    plt.xlim(0, 100)
    for i, v in enumerate(sims):
        plt.text(v + 1, i, f'{v:.1f}%', va='center')
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()
# 绑定方法
SmartCustomerService.visualize_similarity_heatmap = visualize_similarity_heatmap
SmartCustomerService.visualize_query_performance = visualize_query_performance
SmartCustomerService.visualize_search_results = visualize_search_results

if __name__ == "__main__":
    # 1. 准备多语言知识库数据
    knowledge_data = [
        {"question": "如何申请退款？", "answer": "订单页点击退款按钮，1-3天处理。", "language": "Chinese"},
        {"question": "物流信息多久更新？", "answer": "通常4-6小时更新一次。", "language": "Chinese"},
        {"question": "How to track my order?", "answer": "Use tracking number on website.", "language": "English"},
        {"question": "What is the return policy?", "answer": "30-day return policy.", "language": "English"},
        {"question": "配送料金はいくらですか？", "answer": "10000円以上送料無料。", "language": "Japanese"},
        # 同义测试
        {"question": "我想退钱", "answer": "订单页点击退款按钮。", "language": "Chinese"},
        {"question": "Can I get my money back?", "answer": "Contact service within 30 days.", "language": "English"}
    ]
    
    # 2. 初始化系统
    # 注意：确保前四段代码已执行，类已定义
    cs = SmartCustomerService(MODEL_ID, CACHE_DIR)
    
    # 3. 构建知识库
    cs.build_knowledge_base(knowledge_data)
    
    # 4. 定义测试用例
    test_queries = [
        "怎么申请退款？",          # 中文标准问
        "退款流程是什么？",         # 中文同义问
        "How can I get a refund?", # 英文跨语言问
        "我买的东西想退货",         # 中文口语化
        "返金を申請するには？",     # 日语跨语言问
        "你们公司地址在哪？"        # 无关问题（测试鲁棒性）
    ]
    
    print("\n" + "="*50)
    print("开始全场景测试...")
    print("="*50)
    
    # 5. 执行测试循环
    for idx, query in enumerate(test_queries, 1):
        result = cs.answer_query(query, top_k=2)
        # 为每个查询生成结果图
        cs.visualize_search_results(query, result['results'], f"result_plot_{idx}.png")
    
    # 6. 生成全局分析报告
    print("\n生成全局可视化报告...")
    cs.visualize_embeddings("final_pca_analysis.png")
    cs.visualize_tsne("final_tsne_clusters.png")
    cs.visualize_similarity_heatmap("final_heatmap.png")
    cs.visualize_query_performance("final_performance.png")
    
    # 7. 输出最终统计
    stats = cs.get_performance_stats() if hasattr(cs, 'get_performance_stats') else {"msg": "Stats collected"}
    print("\n测试完成！所有图表已保存至当前目录。")