# word2vec_helper.py
# Word2Vec模型加载与使用的辅助模块

import os
import numpy as np
from gensim.models import KeyedVectors, Word2Vec
import jieba
import logging
import time

# 设置日志
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)

class Word2VecHelper:
    def __init__(self, model_path=None):
        """
        初始化Word2Vec辅助类
        
        参数:
            model_path: 预训练模型路径，支持word2vec格式和二进制格式
                        如果为None，将使用默认路径或尝试下载小型模型
        """
        self.model = None
        
        # 更改默认模型路径和备用选项
        if model_path:
            self.model_path = model_path
        else:
            # 首选路径 - 大型腾讯模型
            primary_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 
                                      "models", "chinese_word2vec.bin")
            
            # 备用路径 - 小型模型
            backup_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 
                                     "models", "chinese_word2vec_small.bin")
            
            if os.path.exists(primary_path):
                self.model_path = primary_path
            elif os.path.exists(backup_path):
                self.model_path = backup_path
            else:
                # 如果都不存在，可以尝试自动下载小模型
                self.model_path = primary_path
                self._try_download_small_model()
        
        self.initialized = False
        # 缓存查询结果，提高性能
        self.similarity_cache = {}
        self.similar_words_cache = {}
    
    def _try_download_small_model(self):
        """尝试下载小型词向量模型作为备用选项"""
        try:
            import gensim.downloader as api
            logging.info("尝试下载小型中文词向量模型...")
            
            # 创建模型目录
            os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
            
            # 尝试下载fastText的小型中文模型
            small_model = api.load("fasttext-wiki-news-subwords-300")
            small_model.save(self.model_path.replace(".bin", "_small.bin"))
            logging.info(f"小型模型已下载并保存到 {self.model_path}")
        except Exception as e:
            logging.error(f"无法下载备用模型: {e}")

    def load_model(self):
        """加载Word2Vec模型"""
        try:
            start_time = time.time()
            logging.info(f"开始加载Word2Vec模型: {self.model_path}")
            
            # 判断文件扩展名，选择合适的加载方式
            if self.model_path.endswith('.bin'):
                # 加载二进制格式的模型
                self.model = KeyedVectors.load_word2vec_format(self.model_path, binary=True)
            else:
                # 加载文本格式的模型或gensim模型
                self.model = Word2Vec.load(self.model_path).wv
                
            self.initialized = True
            logging.info(f"Word2Vec模型加载完成，耗时 {time.time() - start_time:.2f} 秒")
            logging.info(f"词向量维度: {self.model.vector_size}")
            logging.info(f"词汇表大小: {len(self.model.index_to_key)}")
            return True
        except Exception as e:
            logging.error(f"加载Word2Vec模型失败: {e}")
            self.initialized = False
            return False
    
    def ensure_initialized(self):
        """确保模型已初始化"""
        if not self.initialized:
            return self.load_model()
        return True
    
    def get_similar_words(self, word, topn=10, min_similarity=0.5):
        """
        获取与给定词语最相似的词语列表
        
        参数:
            word: 输入词语
            topn: 返回相似词的数量
            min_similarity: 最小相似度阈值
        返回:
            相似词列表，如果词不存在或模型未加载则返回空列表
        """
        if not self.ensure_initialized():
            return []
            
        # 检查缓存
        cache_key = f"{word}_{topn}_{min_similarity}"
        if cache_key in self.similar_words_cache:
            return self.similar_words_cache[cache_key]
        
        try:
            # 如果词不在词汇表中，进行分词处理
            if word not in self.model.key_to_index:
                # 对中文词进行分词，然后查找每个子词的相似词
                word_parts = list(jieba.cut(word))
                
                if not word_parts:
                    return []
                
                # 如果存在多个子词，找到存在于模型中的子词
                valid_parts = [w for w in word_parts if w in self.model.key_to_index]
                
                if not valid_parts:
                    return []
                
                # 使用最长的有效子词或第一个有效子词
                valid_parts.sort(key=len, reverse=True)
                word = valid_parts[0]
                
                # 如果替换后的词仍不在词汇表中，返回空列表
                if word not in self.model.key_to_index:
                    return []
            
            # 获取相似词
            similar_words = self.model.most_similar(word, topn=topn*2)  # 多获取一些，后续过滤
            
            # 过滤低于阈值的结果，并只返回词语(不返回相似度)
            filtered_words = [w for w, sim in similar_words if sim >= min_similarity][:topn]
            
            # 缓存结果
            self.similar_words_cache[cache_key] = filtered_words
            return filtered_words
            
        except Exception as e:
            logging.error(f"获取相似词失败: {e}, 词语: {word}")
            return []
    
    def calculate_similarity(self, word1, word2):
        """
        计算两个词的相似度
        
        参数:
            word1, word2: 输入词语
        返回:
            相似度分数(0-1)，如果任意词不存在则返回0
        """
        if not self.ensure_initialized():
            return 0
            
        # 检查缓存
        cache_key = f"{word1}_{word2}"
        reverse_key = f"{word2}_{word1}"
        
        if cache_key in self.similarity_cache:
            return self.similarity_cache[cache_key]
        if reverse_key in self.similarity_cache:
            return self.similarity_cache[reverse_key]
        
        try:
            # 检查词是否在词汇表中
            if word1 not in self.model.key_to_index or word2 not in self.model.key_to_index:
                return 0
            
            similarity = self.model.similarity(word1, word2)
            
            # 缓存结果
            self.similarity_cache[cache_key] = similarity
            return similarity
            
        except Exception as e:
            logging.error(f"计算相似度失败: {e}, 词语: {word1}, {word2}")
            return 0
    
    def expand_query(self, query, topn=5, min_similarity=0.6):
        """
        扩展查询词，返回相关词汇
        
        参数:
            query: 查询词
            topn: 每个词扩展的相似词数量
            min_similarity: 最小相似度阈值
        返回:
            扩展后的词语列表
        """
        if not self.ensure_initialized():
            return [query]
            
        expanded_terms = [query]
        
        # 对查询进行分词
        words = list(jieba.cut(query))
        
        # 为每个词找相似词
        for word in words:
            if len(word) <= 1:  # 忽略单字，减少噪音
                continue
                
            similar_words = self.get_similar_words(word, topn=topn, min_similarity=min_similarity)
            expanded_terms.extend(similar_words)
        
        # 确保唯一性
        return list(set(expanded_terms))

# 单例模式，全局使用一个模型实例
_word2vec_helper = None

def get_word2vec_helper(model_path=None):
    """获取Word2Vec辅助类的全局单例"""
    global _word2vec_helper
    if _word2vec_helper is None:
        _word2vec_helper = Word2VecHelper(model_path)
        _word2vec_helper.ensure_initialized()
    return _word2vec_helper

# 便捷函数，方便直接调用
def get_similar_words(word, topn=10, min_similarity=0.5):
    """获取相似词的便捷函数"""
    helper = get_word2vec_helper()
    return helper.get_similar_words(word, topn, min_similarity)

def calculate_similarity(word1, word2):
    """计算相似度的便捷函数"""
    helper = get_word2vec_helper()
    return helper.calculate_similarity(word1, word2)

def expand_query(query, topn=5, min_similarity=0.6):
    """扩展查询的便捷函数"""
    helper = get_word2vec_helper()
    return helper.expand_query(query, topn, min_similarity)

# 使用示例
if __name__ == "__main__":
    # 测试模型加载和词语相似度
    helper = get_word2vec_helper()
    
    # 测试词
    test_words = ["电影", "功夫", "熊猫", "科幻", "漫威"]
    
    for word in test_words:
        print(f"\n{word} 的相似词:")
        similar = helper.get_similar_words(word, topn=5)
        for sim_word in similar:
            print(f"  - {sim_word}")
    
    # 测试相似度计算
    word_pairs = [
        ("电影", "电视"),
        ("功夫", "武术"),
        ("科幻", "未来"),
        ("漫威", "超级英雄")
    ]
    
    print("\n词语相似度:")
    for w1, w2 in word_pairs:
        sim = helper.calculate_similarity(w1, w2)
        print(f"  {w1} <-> {w2}: {sim:.4f}")
    
    # 测试查询扩展
    test_queries = ["功夫熊猫", "科幻电影", "漫威英雄"]
    
    print("\n查询扩展:")
    for query in test_queries:
        expanded = helper.expand_query(query)
        print(f"  {query} -> {expanded}")