"""
  i2i - DeepWalk算法实现
  基于用户-物品图结构训练DeepWalk模型，获取物品向量相似度

  复用 graphembedding/deepwalk/ 的高效实现

  """

  import pandas as pd
  import argparse

  import os

  from datetime import datetime
  from collections import defaultdict
  from gensim.models import Word2Vec

  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,

      save_readable_index, fetch_name_mappings, log_algorithm_params,
      log_processing_step
  )

  # 导入 DeepWalk 实现

  from deepwalk.deepwalk import DeepWalk

  def build_edge_file_from_db(df, behavior_weights, output_path, logger):

      """

      从数据库数据构建边文件
      边文件格式: item_id \t neighbor_id1:weight1,neighbor_id2:weight2,...

      
      Args:
          df: DataFrame with columns: user_id, item_id, event_type
          behavior_weights: 行为权重字典

          output_path: 边文件输出路径
          logger: 日志对象

      """

      logger.info("开始构建物品图...")
      

      # 构建用户-物品列表
      user_items = defaultdict(list)
      
      for _, row in df.iterrows():
          user_id = row['user_id']
          item_id = str(row['item_id'])
          event_type = row['event_type']
          weight = behavior_weights.get(event_type, 1.0)

          user_items[user_id].append((item_id, weight))
      

      logger.info(f"共有 {len(user_items)} 个用户")
      

      # 构建物品图边
      edge_dict = defaultdict(lambda: defaultdict(float))
      
      for user_id, items in user_items.items():

          # 限制每个用户的物品数量，避免内存爆炸
          if len(items) > 100:
              # 按权重排序，只保留前100个
              items = sorted(items, key=lambda x: -x[1])[:100]
          

          # 物品两两组合，构建边
          for i in range(len(items)):
              item_i, weight_i = items[i]
              for j in range(i + 1, len(items)):
                  item_j, weight_j = items[j]
                  
                  # 边的权重为两个物品权重的平均值
                  edge_weight = (weight_i + weight_j) / 2.0
                  edge_dict[item_i][item_j] += edge_weight
                  edge_dict[item_j][item_i] += edge_weight
      

      logger.info(f"构建物品图完成，共 {len(edge_dict)} 个节点")

      # 保存边文件
      logger.info(f"保存边文件到 {output_path}")

      with open(output_path, 'w', encoding='utf-8') as f:
          for item_id, neighbors in edge_dict.items():

              neighbor_str = ','.join([f'{nbr}:{weight:.4f}' for nbr, weight in neighbors.items()])
              f.write(f'{item_id}\t{neighbor_str}\n')
      

      logger.info(f"边文件保存完成")
      return len(edge_dict)

  def train_word2vec_from_walks(walks_file, config, logger):

      """

      从游走文件训练Word2Vec模型

      
      Args:

          walks_file: 游走序列文件路径
          config: Word2Vec配置
          logger: 日志对象

      
      Returns:

          Word2Vec模型

      """

      logger.info(f"从 {walks_file} 读取游走序列...")

      # 读取游走序列
      sentences = []
      with open(walks_file, 'r', encoding='utf-8') as f:
          for line in f:
              walk = line.strip().split()

              if len(walk) >= 2:

                  sentences.append(walk)

      logger.info(f"共读取 {len(sentences)} 条游走序列")

      # 训练Word2Vec
      logger.info("开始训练Word2Vec模型...")

      model = Word2Vec(

          sentences=sentences,

          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
      )
      

      logger.info(f"训练完成。词汇表大小：{len(model.wv)}")

      return model
  
  

  def generate_similarities(model, top_n, logger):

      """

      从Word2Vec模型生成物品相似度

      
      Args:
          model: Word2Vec模型
          top_n: Top N similar items

          logger: 日志对象

      
      Returns:
          Dict[item_id, List[Tuple(similar_item_id, score)]]
      """

      logger.info("生成相似度...")

      result = {}
      
      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
      

      logger.info(f"为 {len(result)} 个物品生成了相似度")

      return result
  
  

  def save_results(result, output_file, name_mappings, logger):
      """
      保存相似度结果到文件
      
      Args:
          result: 相似度字典
          output_file: 输出文件路径
          name_mappings: ID到名称的映射
          logger: 日志对象
      """
      logger.info(f"保存结果到 {output_file}...")
      
      with open(output_file, 'w', encoding='utf-8') as f:
          for item_id, sims in result.items():
              # 获取物品名称
              item_name = name_mappings.get(int(item_id), 'Unknown') if item_id.isdigit() 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"结果保存完成")
  
  

  def main():
      parser = argparse.ArgumentParser(description='Run DeepWalk for i2i similarity')
      parser.add_argument('--num_walks', type=int, default=I2I_CONFIG['deepwalk']['num_walks'],
                         help='Number of walks per node')
      parser.add_argument('--walk_length', type=int, default=I2I_CONFIG['deepwalk']['walk_length'],
                         help='Walk length')
      parser.add_argument('--window_size', type=int, default=I2I_CONFIG['deepwalk']['window_size'],
                         help='Window size for Word2Vec')
      parser.add_argument('--vector_size', type=int, default=I2I_CONFIG['deepwalk']['vector_size'],
                         help='Vector size for Word2Vec')
      parser.add_argument('--min_count', type=int, default=I2I_CONFIG['deepwalk']['min_count'],
                         help='Minimum word count')
      parser.add_argument('--workers', type=int, default=I2I_CONFIG['deepwalk']['workers'],
                         help='Number of workers')
      parser.add_argument('--epochs', type=int, default=I2I_CONFIG['deepwalk']['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('--output', type=str, default=None,
                         help='Output file path')
      parser.add_argument('--save_model', action='store_true',
                         help='Save Word2Vec model')
      parser.add_argument('--save_graph', action='store_true',
                         help='Save graph edge file')

      parser.add_argument('--debug', action='store_true',
                         help='Enable debug mode with detailed logging and readable output')

      parser.add_argument('--use_softmax', action='store_true',
                         help='Use softmax-based alias sampling (default: False)')
      parser.add_argument('--temperature', type=float, default=1.0,
                         help='Temperature for softmax (default: 1.0)')

      
      args = parser.parse_args()
      

      # 设置logger
      logger = setup_debug_logger('i2i_deepwalk', debug=args.debug)
      
      # 记录算法参数
      params = {
          'num_walks': args.num_walks,
          'walk_length': args.walk_length,
          '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,

          'debug': args.debug,
          'use_softmax': args.use_softmax,
          'temperature': args.temperature

      }
      log_algorithm_params(logger, params)
      

      # 创建临时目录
      temp_dir = os.path.join(OUTPUT_DIR, 'temp')
      os.makedirs(temp_dir, exist_ok=True)
      
      date_str = datetime.now().strftime('%Y%m%d')
      edge_file = os.path.join(temp_dir, f'item_graph_{date_str}.txt')
      walks_file = os.path.join(temp_dir, f'walks_{date_str}.txt')
      
      # ============================================================
      # 步骤1: 从数据库获取数据并构建边文件
      # ============================================================
      log_processing_step(logger, "从数据库获取数据")
      

      # 创建数据库连接

      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.event AS event_type,
          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
      """
      

      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, "用户行为数据")

      
      # 定义行为权重
      behavior_weights = {
          'click': 1.0,
          'contactFactory': 5.0,
          'addToPool': 2.0,
          'addToCart': 3.0,
          'purchase': 10.0
      }

      logger.debug(f"行为权重: {behavior_weights}")

      # 构建边文件
      log_processing_step(logger, "构建边文件")
      num_nodes = build_edge_file_from_db(df, behavior_weights, edge_file, logger)
      
      # ============================================================
      # 步骤2: 使用DeepWalk进行随机游走
      # ============================================================
      log_processing_step(logger, "执行DeepWalk随机游走")
      
      logger.info("初始化DeepWalk...")
      deepwalk = DeepWalk(
          edge_file=edge_file,
          node_tag_file=None,  # 不使用标签游走
          use_softmax=args.use_softmax,
          temperature=args.temperature,
          p_tag_walk=0.0  # 不使用标签游走
      )

      logger.info("开始随机游走...")
      deepwalk.simulate_walks(
          num_walks=args.num_walks,
          walk_length=args.walk_length,
          workers=args.workers,
          output_file=walks_file
      )

      # ============================================================
      # 步骤3: 训练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  # Skip-gram

      }

      logger.debug(f"Word2Vec配置: {w2v_config}")

      model = train_word2vec_from_walks(walks_file, w2v_config, logger)

      
      # 保存模型（可选）
      if args.save_model:

          model_path = os.path.join(OUTPUT_DIR, f'deepwalk_model_{date_str}.model')

          model.save(model_path)

          logger.info(f"模型已保存到 {model_path}")

      # ============================================================
      # 步骤4: 生成相似度
      # ============================================================

      log_processing_step(logger, "生成相似度")

      result = generate_similarities(model, args.top_n, logger)

      # ============================================================
      # 步骤5: 保存结果
      # ============================================================

      log_processing_step(logger, "保存结果")

      
      output_file = args.output or os.path.join(OUTPUT_DIR, f'i2i_deepwalk_{date_str}.txt')

      # 获取name mappings
      name_mappings = {}
      if args.debug:
          logger.info("获取物品名称映射...")
          name_mappings = fetch_name_mappings(engine, debug=True)
      

      save_results(result, output_file, name_mappings, logger)

      logger.info(f"✓ DeepWalk完成！")
      logger.info(f"  - 输出文件: {output_file}")
      logger.info(f"  - 商品数: {len(result)}")
      if result:
          avg_sims = sum(len(sims) for sims in result.values()) / len(result)
          logger.info(f"  - 平均相似商品数: {avg_sims:.1f}")

      
      # 如果启用debug模式，保存可读格式
      if args.debug:
          log_processing_step(logger, "保存Debug可读格式")
          save_readable_index(
              output_file,
              result,
              name_mappings,

              description='i2i:deepwalk'

          )

      
      # 清理临时文件（可选）
      if not args.save_graph:
          if os.path.exists(edge_file):
              os.remove(edge_file)
              logger.debug(f"已删除临时文件: {edge_file}")
      if os.path.exists(walks_file):
          os.remove(walks_file)
          logger.debug(f"已删除临时文件: {walks_file}")
      
      print(f"✓ DeepWalk相似度计算完成")
      print(f"  - 输出文件: {output_file}")
      print(f"  - 商品数: {len(result)}")

  
  
  if __name__ == '__main__':
      main()