import logging
import argparse
from gensim.models import Word2Vec
import os
import multiprocessing

"""
说明

输入文件格式：
--input-file：输入的文件路径，文件内容每一行是一个空格分隔的 bid 序列（即一个随机游走结果，或者直接是一个用户行为session）。
输出文件格式：
--output-dir：输出的目录路径，保存嵌入向量和相似 bid 的文件。
生成两个文件：
bid_embeddings.txt：每个 bid 的嵌入向量，格式为 bid embedding_vector，例如：
123 0.12 0.34 0.56 ... 0.78
456 0.23 0.45 0.67 ... 0.89
bid_top_similar.txt：每个 bid 最相似的 top K 个 bid，格式为 bid similar_bid1:similarity1 similar_bid2:similarity2 ...，例如：
123 456:0.89 789:0.88 101:0.87 ...
456 123:0.89 678:0.85 234:0.84 ...
命令行参数：
--embedding-size：指定 Word2Vec 嵌入向量的维度，默认为 128。
--window：Word2Vec 模型的窗口大小，控制词的上下文范围，默认为 5。
--min-count：忽略词频低于该值的 bid，默认是 1，即不忽略任何 bid。
--workers：并行计算的线程数量，默认为 4。
--top-k：每个 bid 输出的最相似的 top K bid，默认是 200。


执行示例

假设：
输入文件路径是 input_sentences.txt
输出目录是 output/

那么可以使用以下命令：
bash
python word2vec_bid_similarity.py --input-file input_sentences.txt --output-dir output/ --embedding-size 128 --top-k 200


依赖项

请确保安装了以下依赖项：
bash
pip install gensim
"""

def train_word2vec(input_file, output_dir, embedding_size=128, window=5, min_count=1, workers=None, top_k=200, epochs=5):
    """
    训练Word2Vec模型，并保存每个bid的embedding及top K相似的bid。
    
    :param input_file: 句子文件路径
    :param output_dir: 输出文件的目录路径
    :param embedding_size: 嵌入维度大小
    :param window: Word2Vec中的窗口大小
    :param min_count: Word2Vec中忽略频次低于min_count的词
    :param workers: 使用的线程数，如果为None，则设置为cpu_count-2
    :param top_k: 每个bid的最相似bid的数量
    :param epochs: 训练的epoch数量
    """
    # 如果未设置workers，默认使用CPU核心数-2
    if workers is None:
        workers = max(1, multiprocessing.cpu_count() - 2)
    
    # 检查输出目录是否存在，不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    logging.info(f"Loading sentences from {input_file}")
    # 读取输入文件，格式是每行一个bid序列
    sentences = []
    with open(input_file, 'r') as f:
        for line in f:
            sentences.append(line.strip().split())

    # 训练Word2Vec模型
    logging.info(f"Training Word2Vec model with embedding size {embedding_size}, window {window}, epochs {epochs}, workers {workers}")
    model = Word2Vec(sentences, vector_size=embedding_size, window=window, min_count=min_count, workers=workers, epochs=epochs)

    # 保存每个bid的embedding
    embedding_file = os.path.join(output_dir, "bid_embeddings.txt")
    logging.info(f"Saving embeddings to {embedding_file}")
    with open(embedding_file, 'w') as f_out:
        for bid in model.wv.index_to_key:
            vector = model.wv[bid]
            f_out.write(f"{bid}\t{','.join(map(str, vector))}\n")

    # 保存每个bid的top K相似bid
    similar_file = os.path.join(output_dir, "bid_top_similar.txt")
    logging.info(f"Saving top {top_k} similar bids for each bid to {similar_file}")
    with open(similar_file, 'w') as f_out:
        for bid in model.wv.index_to_key:
            similar_bids = model.wv.most_similar(bid, topn=top_k)
            similar_bids_str = ','.join([f"{similar_bid[0]}:{round(similar_bid[1], 4)}" for similar_bid in similar_bids])
            f_out.write(f"{bid}\t{similar_bids_str}\n")
    
    logging.info("Process completed successfully.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train Word2Vec model and calculate bid similarity.")
    parser.add_argument('--input-file', type=str, required=True, help="Path to the input sentence file")
    parser.add_argument('--output-dir', type=str, required=True, help="Directory to save output embeddings and similarity results")
    parser.add_argument('--embedding-size', type=int, default=128, help="Size of the bid embedding vectors (default: 128)")
    parser.add_argument('--window', type=int, default=5, help="Window size for Word2Vec (default: 5)")
    parser.add_argument('--min-count', type=int, default=1, help="Minimum frequency of bids to be considered (default: 1)")
    parser.add_argument('--workers', type=int, default=None, help="Number of workers (default: cpu_count-2)")
    parser.add_argument('--top-k', type=int, default=200, help="Number of top similar bids to output (default: 200)")
    parser.add_argument('--epochs', type=int, default=5, help="Number of epochs to train the model (default: 5)")

    args = parser.parse_args()

    # 初始化日志
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)

    # 执行训练和保存结果
    train_word2vec(input_file=args.input_file, output_dir=args.output_dir, 
                   embedding_size=args.embedding_size, window=args.window, 
                   min_count=args.min_count, workers=args.workers, top_k=args.top_k, epochs=args.epochs)


# python w2v.py --input-file input_sentences.txt --output-dir output/ --embedding-size 128 --top-k 200 --epochs 10