"""
  i2i - Session Word2Vec算法实现
  基于用户会话序列训练Word2Vec模型，获取物品向量相似度
  """

  import pandas as pd
  import json
  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 offline_tasks.config.offline_config import (
      DB_CONFIG, OUTPUT_DIR, I2I_CONFIG, get_time_range,
      DEFAULT_LOOKBACK_DAYS, DEFAULT_I2I_TOP_N
  )

  from offline_tasks.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)} 条记录")
      
      # 记录数据信息
      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()