"""
i2i - Session Word2Vec算法实现
基于用户会话序列训练Word2Vec模型，获取物品向量相似度
"""
import pandas as pd
import json
import os
import argparse
from datetime import datetime
from collections import defaultdict
from gensim.models import Word2Vec
import numpy as np
from db_service import create_db_connection
from config.offline_config import (
    DB_CONFIG, OUTPUT_DIR, I2I_CONFIG, get_time_range,
    DEFAULT_LOOKBACK_DAYS, DEFAULT_I2I_TOP_N
)
from scripts.debug_utils import (
    setup_debug_logger, log_dataframe_info, log_dict_stats,
    save_readable_index, fetch_name_mappings, log_algorithm_params,
    log_processing_step
)


def prepare_session_data(df, max_session_length=50, min_session_length=2, logger=None):
    """
    准备会话数据 - 基于固定长度分块，适合B2B低频场景
    
    Args:
        df: DataFrame with columns: user_id, item_id, create_time
        session_gap_minutes: 会话间隔时间（分钟）
        logger: Logger instance for debugging
    
    Returns:
        List of sessions, each session is a list of item_ids
    """
    sessions = []
    
    if logger:
        logger.debug(f"开始准备会话数据（固定长度分块）：max_length={max_session_length}, min_length={min_session_length}")
    
    # 按用户和时间排序
    df = df.sort_values(['user_id', 'create_time'])
    
    # 按用户分组，获取每个用户的行为序列
    for user_id, user_df in df.groupby('user_id'):
        # 获取用户的item序列
        item_sequence = user_df['item_id'].astype(str).tolist()
        
        # 如果序列太短，跳过
        if len(item_sequence) < min_session_length:
            continue
        
        # 按最大长度分块（不重叠）
        user_sessions = [
            item_sequence[i:i + max_session_length] 
            for i in range(0, len(item_sequence), max_session_length)
        ]
        
        # 过滤掉长度不足的最后一块
        user_sessions = [s for s in user_sessions if len(s) >= min_session_length]
        
        sessions.extend(user_sessions)
    
    if logger:
        if sessions:
            session_lengths = [len(s) for s in sessions]
            logger.debug(f"生成 {len(sessions)} 个会话")
            logger.debug(f"会话长度统计：最小={min(session_lengths)}, 最大={max(session_lengths)}, "
                        f"平均={sum(session_lengths)/len(session_lengths):.2f}")
        else:
            logger.warning("未生成任何会话！")
    
    return sessions


def train_word2vec(sessions, config, logger=None):
    """
    训练Word2Vec模型
    
    Args:
        sessions: List of sessions
        config: Word2Vec配置
        logger: Logger instance for debugging
    
    Returns:
        Word2Vec模型
    """
    if logger:
        logger.info(f"训练Word2Vec模型，共 {len(sessions)} 个会话")
        logger.debug(f"模型参数：vector_size={config['vector_size']}, window={config['window_size']}, "
                    f"min_count={config['min_count']}, epochs={config['epochs']}")
    else:
        print(f"Training Word2Vec with {len(sessions)} sessions...")
    
    model = Word2Vec(
        sentences=sessions,
        vector_size=config['vector_size'],
        window=config['window_size'],
        min_count=config['min_count'],
        workers=config['workers'],
        sg=config['sg'],
        epochs=config['epochs'],
        seed=42
    )
    
    if logger:
        logger.info(f"训练完成。词汇表大小：{len(model.wv)}")
    else:
        print(f"Training completed. Vocabulary size: {len(model.wv)}")
    return model


def generate_similarities(model, top_n=50, logger=None):
    """
    生成物品相似度
    
    Args:
        model: Word2Vec模型
        top_n: Top N similar items
        logger: Logger instance for debugging
    
    Returns:
        Dict[item_id, List[Tuple(similar_item_id, score)]]
    """
    result = {}
    
    if logger:
        logger.info(f"生成Top {top_n} 相似物品")
    
    for item_id in model.wv.index_to_key:
        try:
            similar_items = model.wv.most_similar(item_id, topn=top_n)
            result[item_id] = [(sim_id, float(score)) for sim_id, score in similar_items]
        except KeyError:
            continue
    
    if logger:
        logger.info(f"生成了 {len(result)} 个物品的相似度")
    
    return result


def main():
    parser = argparse.ArgumentParser(description='Run Session Word2Vec for i2i similarity')
    parser.add_argument('--window_size', type=int, default=I2I_CONFIG['session_w2v']['window_size'],
                       help='Window size for Word2Vec')
    parser.add_argument('--vector_size', type=int, default=I2I_CONFIG['session_w2v']['vector_size'],
                       help='Vector size for Word2Vec')
    parser.add_argument('--min_count', type=int, default=I2I_CONFIG['session_w2v']['min_count'],
                       help='Minimum word count')
    parser.add_argument('--workers', type=int, default=I2I_CONFIG['session_w2v']['workers'],
                       help='Number of workers')
    parser.add_argument('--epochs', type=int, default=I2I_CONFIG['session_w2v']['epochs'],
                       help='Number of epochs')
    parser.add_argument('--top_n', type=int, default=DEFAULT_I2I_TOP_N,
                       help=f'Top N similar items to output (default: {DEFAULT_I2I_TOP_N})')
    parser.add_argument('--lookback_days', type=int, default=DEFAULT_LOOKBACK_DAYS,
                       help=f'Number of days to look back (default: {DEFAULT_LOOKBACK_DAYS})')
    parser.add_argument('--max_session_length', type=int, default=50,
                       help='Maximum session length for chunking (default: 50)')
    parser.add_argument('--min_session_length', type=int, default=2,
                       help='Minimum session length to keep (default: 2)')
    parser.add_argument('--output', type=str, default=None,
                       help='Output file path')
    parser.add_argument('--save_model', action='store_true',
                       help='Save Word2Vec model')
    parser.add_argument('--debug', action='store_true',
                       help='Enable debug mode with detailed logging and readable output')
    
    args = parser.parse_args()
    
    # 设置logger
    logger = setup_debug_logger('i2i_session_w2v', debug=args.debug)
    
    # 记录算法参数
    params = {
        'window_size': args.window_size,
        'vector_size': args.vector_size,
        'min_count': args.min_count,
        'workers': args.workers,
        'epochs': args.epochs,
        'top_n': args.top_n,
        'lookback_days': args.lookback_days,
        'max_session_length': args.max_session_length,
        'min_session_length': args.min_session_length,
        'debug': args.debug
    }
    log_algorithm_params(logger, params)
    
    # 创建数据库连接
    logger.info("连接数据库...")
    engine = create_db_connection(
        DB_CONFIG['host'],
        DB_CONFIG['port'],
        DB_CONFIG['database'],
        DB_CONFIG['username'],
        DB_CONFIG['password']
    )
    
    # 获取时间范围
    start_date, end_date = get_time_range(args.lookback_days)
    logger.info(f"获取数据范围：{start_date} 到 {end_date}")
    
    # SQL查询 - 获取用户行为序列
    sql_query = f"""
    SELECT 
        se.anonymous_id AS user_id,
        se.item_id,
        se.create_time,
        pgs.name AS item_name
    FROM 
        sensors_events se
    LEFT JOIN prd_goods_sku pgs ON se.item_id = pgs.id
    WHERE 
        se.event IN ('click', 'contactFactory', 'addToPool', 'addToCart', 'purchase')
        AND se.create_time >= '{start_date}'
        AND se.create_time <= '{end_date}'
        AND se.item_id IS NOT NULL
        AND se.anonymous_id IS NOT NULL
    ORDER BY 
        se.anonymous_id,
        se.create_time
    """
    
    logger.info("执行SQL查询...")
    df = pd.read_sql(sql_query, engine)
    logger.info(f"获取到 {len(df)} 条记录")
    
    # 确保ID为整数类型
    df['item_id'] = df['item_id'].astype(int)
    df['user_id'] = df['user_id'].astype(str)
    
    # 记录数据信息
    log_dataframe_info(logger, df, "用户行为数据")
    
    # 转换create_time为datetime
    df['create_time'] = pd.to_datetime(df['create_time'])
    
    # 准备会话数据
    log_processing_step(logger, "准备会话数据")
    sessions = prepare_session_data(
        df, 
        max_session_length=args.max_session_length,
        min_session_length=args.min_session_length,
        logger=logger
    )
    logger.info(f"生成 {len(sessions)} 个会话")
    
    # 训练Word2Vec模型
    log_processing_step(logger, "训练Word2Vec模型")
    w2v_config = {
        'vector_size': args.vector_size,
        'window_size': args.window_size,
        'min_count': args.min_count,
        'workers': args.workers,
        'epochs': args.epochs,
        'sg': 1
    }
    
    model = train_word2vec(sessions, w2v_config, logger=logger)
    
    # 保存模型（可选）
    if args.save_model:
        model_path = os.path.join(OUTPUT_DIR, f'session_w2v_model_{datetime.now().strftime("%Y%m%d")}.model')
        model.save(model_path)
        logger.info(f"模型已保存到 {model_path}")
    
    # 生成相似度
    log_processing_step(logger, "生成相似度")
    result = generate_similarities(model, top_n=args.top_n, logger=logger)
    
    # 输出结果
    log_processing_step(logger, "保存结果")
    output_file = args.output or os.path.join(OUTPUT_DIR, f'i2i_session_w2v_{datetime.now().strftime("%Y%m%d")}.txt')
    
    # 获取name mappings用于标准输出格式
    name_mappings = {}
    if args.debug:
        logger.info("获取物品名称映射...")
        name_mappings = fetch_name_mappings(engine, debug=True)
    
    logger.info(f"写入结果到 {output_file}...")
    with open(output_file, 'w', encoding='utf-8') as f:
        for item_id, sims in result.items():
            # 使用name_mappings获取名称，如果没有则从df中获取
            item_name = name_mappings.get(int(item_id), 'Unknown') if item_id.isdigit() else 'Unknown'
            if item_name == 'Unknown' and 'item_name' in df.columns:
                item_name = df[df['item_id'].astype(str) == item_id]['item_name'].iloc[0] if len(df[df['item_id'].astype(str) == item_id]) > 0 else 'Unknown'
            
            if not sims:
                continue
            
            # 格式：item_id \t item_name \t similar_item_id1:score1,similar_item_id2:score2,...
            sim_str = ','.join([f'{sim_id}:{score:.4f}' for sim_id, score in sims])
            f.write(f'{item_id}\t{item_name}\t{sim_str}\n')
    
    logger.info(f"完成！为 {len(result)} 个物品生成了相似度")
    logger.info(f"输出保存到：{output_file}")
    
    # 如果启用debug模式，保存可读格式
    if args.debug:
        log_processing_step(logger, "保存Debug可读格式")
        save_readable_index(
            output_file,
            result,
            name_mappings,
            description='i2i:session_w2v'
        )


if __name__ == '__main__':
    main()