AiDA_Python/app/service/recommendation_system/precompute.py

"""
预计算模块
包含：数据库表结构优化、Milvus集合创建、系统图向量预计算、初始用户偏好向量生成
"""
import logging
import math
import pymysql
import numpy as np
from datetime import datetime
from typing import List, Dict, Tuple, Optional
from collections import defaultdict

from app.service.recommendation_system.config import (
    MYSQL_CONFIG, TABLE_USER_PREFERENCE_LOG, TABLE_SYS_FILE,
    RECOMMENDATION_CONFIG, REDIS_KEY_USER_PREF_PREFIX
)
from app.service.recommendation_system.vector_utils import extract_feature_vector, normalize_vector, compute_weighted_average
from app.service.recommendation_system.milvus_client import (
    create_collection, insert_vectors, query_vectors_by_paths
)
from app.service.utils.redis_utils import Redis
import json

logger = logging.getLogger(__name__)


def optimize_database_table():
    """
    优化 user_preference 表结构
    添加冗余字段和索引
    """
    conn = None
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 1. 添加冗余字段
        logger.info("添加冗余字段...")
        alter_sqls = [
            f"ALTER TABLE {TABLE_USER_PREFERENCE_LOG} ADD COLUMN category VARCHAR(100) COMMENT '类别：lower(level3_type + \"_\" + level2_type)'",
            f"ALTER TABLE {TABLE_USER_PREFERENCE_LOG} ADD COLUMN style VARCHAR(50) COMMENT '风格样式'",
            f"ALTER TABLE {TABLE_USER_PREFERENCE_LOG} ADD COLUMN is_system_sketch TINYINT(1) DEFAULT 1 COMMENT '是否为系统图（1-是，0-用户图）'",
            f"ALTER TABLE {TABLE_USER_PREFERENCE_LOG} ADD COLUMN sys_file_id BIGINT NULL COMMENT '系统文件ID'",
        ]
        
        for sql in alter_sqls:
            try:
                cursor.execute(sql)
                logger.info(f"执行成功: {sql[:50]}...")
            except Exception as e:
                if "Duplicate column name" in str(e):
                    logger.info(f"字段已存在，跳过: {sql[:50]}...")
                else:
                    logger.warning(f"执行失败: {sql[:50]}... 错误: {e}")
        
        # 2. 创建索引（MySQL 不支持 IF NOT EXISTS，需要先检查）
        logger.info("创建索引...")
        index_definitions = [
            ("idx_account_category_time", ["account_id", "category", "data_time"]),
            ("idx_account_path", ["account_id", "path"]),
        ]
        
        for index_name, columns in index_definitions:
            try:
                # 检查索引是否已存在
                cursor.execute(f"""
                    SELECT COUNT(*) 
                    FROM information_schema.statistics 
                    WHERE table_schema = DATABASE() 
                      AND table_name = '{TABLE_USER_PREFERENCE_LOG}' 
                      AND index_name = '{index_name}'
                """)
                exists = cursor.fetchone()[0] > 0
                
                if exists:
                    logger.info(f"索引已存在，跳过: {index_name}")
                else:
                    # 创建索引
                    columns_str = ', '.join(columns)
                    create_sql = f"CREATE INDEX {index_name} ON {TABLE_USER_PREFERENCE_LOG}({columns_str})"
                    cursor.execute(create_sql)
                    logger.info(f"索引创建成功: {index_name}")
            except Exception as e:
                logger.warning(f"索引创建失败: {index_name} 错误: {e}")
        
        conn.commit()
        logger.info("数据库表结构优化完成")
        
    except Exception as e:
        logger.error(f"数据库表结构优化失败: {e}", exc_info=True)
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()


def migrate_historical_data(batch_size: int = 1000):
    """
    历史数据迁移：批量更新冗余字段
    
    Args:
        batch_size: 每批处理数量
    """
    conn = None
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 查询需要更新的记录数
        cursor.execute(f"""
            SELECT COUNT(*) 
            FROM {TABLE_USER_PREFERENCE_LOG} u
            WHERE u.category IS NULL
        """)
        total_count = cursor.fetchone()[0]
        logger.info(f"需要迁移的记录数: {total_count}")
        
        if total_count == 0:
            logger.info("无需迁移数据")
            return
        
        # 分批处理
        offset = 0
        processed = 0
        
        while offset < total_count:
            # 查询一批记录
            cursor.execute(f"""
                SELECT u.id, u.path
                FROM {TABLE_USER_PREFERENCE_LOG} u
                WHERE u.category IS NULL
                LIMIT {batch_size} OFFSET {offset}
            """)
            records = cursor.fetchall()
            
            if not records:
                break
            
            # 批量更新
            for record_id, path in records:
                # 查询 t_sys_file 表
                cursor.execute(f"""
                    SELECT id, url, style, level3_type, level2_type, deprecated 
                    FROM {TABLE_SYS_FILE} 
                    WHERE url = %s 
                    LIMIT 1
                """, (path,))
                
                sys_file = cursor.fetchone()
                
                if sys_file:
                    # 系统图
                    sys_file_id, url, style, level3_type, level2_type, deprecated = sys_file
                    category = f"{level3_type.lower()}_{level2_type.lower()}"
                    
                    cursor.execute(f"""
                        UPDATE {TABLE_USER_PREFERENCE_LOG}
                        SET category = %s,
                            style = %s,
                            is_system_sketch = 1,
                            sys_file_id = %s
                        WHERE id = %s
                    """, (category, style, sys_file_id, record_id))
                else:
                    # 用户图
                    cursor.execute(f"""
                        UPDATE {TABLE_USER_PREFERENCE_LOG}
                        SET is_system_sketch = 0,
                            category = NULL,
                            style = NULL,
                            sys_file_id = NULL
                        WHERE id = %s
                    """, (record_id,))
            
            conn.commit()
            processed += len(records)
            offset += batch_size
            logger.info(f"已迁移 {processed}/{total_count} 条记录")
        
        logger.info("历史数据迁移完成")
        
    except Exception as e:
        logger.error(f"历史数据迁移失败: {e}", exc_info=True)
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()


def precompute_system_sketch_vectors(batch_size: int = 1000, retry_times: int = 3):
    """
    系统图向量预计算与导入
    
    Args:
        batch_size: 每批处理数量
        retry_times: 失败重试次数
    """
    conn = None
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 1. 数据筛选
        logger.info("查询系统图数据...")
        cursor.execute(f"""
            SELECT id, url, style, level3_type, level2_type, deprecated 
            FROM {TABLE_SYS_FILE} 
            WHERE level1_type = 'Images' 
              AND style IS NOT NULL 
              AND style != '' 
              AND deprecated != 1
        """)
        records = cursor.fetchall()
        logger.info(f"找到 {len(records)} 条系统图记录")
        
        if not records:
            logger.warning("没有找到系统图数据")
            return
        
        # 2. 批量处理
        failed_records = []
        batch_data = []
        
        for idx, (sys_file_id, url, style, level3_type, level2_type, deprecated) in enumerate(records, 1):
            try:
                # 计算 category
                category = f"{level3_type.lower()}_{level2_type.lower()}"
                
                # 提取特征向量
                feature_vector = extract_feature_vector(url)
                
                # 检查向量是否有效
                if np.all(feature_vector == 0):
                    logger.warning(f"向量提取失败，跳过: {url}")
                    failed_records.append((sys_file_id, url))
                    continue
                
                # 准备数据
                data_item = {
                    "path": url,
                    "sys_file_id": sys_file_id,
                    "style": style,
                    "category": category,
                    "is_system_sketch": 1,
                    "deprecated": deprecated if deprecated else 0,
                    "feature_vector": feature_vector.tolist()
                }
                
                batch_data.append(data_item)
                
                # 批量写入
                if len(batch_data) >= batch_size:
                    try:
                        insert_vectors(batch_data)
                        batch_data = []
                        logger.info(f"已处理 {idx}/{len(records)} 条记录")
                    except Exception as e:
                        logger.error(f"批量写入失败: {e}")
                        failed_records.extend([(item["sys_file_id"], item["path"]) for item in batch_data])
                        batch_data = []
                
            except Exception as e:
                logger.error(f"处理记录失败 [{url}]: {e}")
                failed_records.append((sys_file_id, url))
        
        # 写入剩余数据
        if batch_data:
            try:
                insert_vectors(batch_data)
            except Exception as e:
                logger.error(f"写入剩余数据失败: {e}")
                failed_records.extend([(item["sys_file_id"], item["path"]) for item in batch_data])
        
        # 3. 重试失败记录
        if failed_records and retry_times > 0:
            logger.info(f"重试 {len(failed_records)} 条失败记录...")
            for retry in range(retry_times):
                retry_failed = []
                for sys_file_id, url in failed_records:
                    try:
                        category = f"{level3_type.lower()}_{level2_type.lower()}"
                        feature_vector = extract_feature_vector(url)
                        if not np.all(feature_vector == 0):
                            data_item = {
                                "path": url,
                                "sys_file_id": sys_file_id,
                                "style": style,
                                "category": category,
                                "is_system_sketch": 1,
                                "deprecated": 0,
                                "feature_vector": feature_vector.tolist()
                            }
                            insert_vectors([data_item])
                        else:
                            retry_failed.append((sys_file_id, url))
                    except Exception as e:
                        logger.error(f"重试失败 [{url}]: {e}")
                        retry_failed.append((sys_file_id, url))
                
                failed_records = retry_failed
                if not failed_records:
                    break
        
        if failed_records:
            logger.warning(f"仍有 {len(failed_records)} 条记录处理失败")
        
        logger.info("系统图向量预计算完成")
        
    except Exception as e:
        logger.error(f"系统图向量预计算失败: {e}", exc_info=True)
    finally:
        if conn:
            conn.close()


def compute_user_preference_vector(
    account_id: int,
    category: str,
    conn: Optional[pymysql.connections.Connection] = None
    # max_date: Optional[datetime] = None
) -> Optional[np.ndarray]:
    """
    计算用户偏好向量
    
    Args:
        account_id: 用户ID
        category: 类别
        conn: 数据库连接（可选）
        max_date: 最大日期（可选，用于评估时只使用训练集数据）
    
    Returns:
        用户偏好向量（2048维），失败返回 None
    """
    from datetime import datetime
    
    should_close = False
    if conn is None:
        conn = pymysql.connect(**MYSQL_CONFIG)
        should_close = True
    
    try:
        cursor = conn.cursor()
        
        # 1. 获取点赞记录（如果指定了max_date，只查询该日期之前的数据）
        if max_date:
            cursor.execute(f"""
                SELECT path, data_time 
                FROM {TABLE_USER_PREFERENCE_LOG} 
                WHERE account_id = %s AND category = %s AND style is not null
                  AND data_time < %s
                ORDER BY data_time DESC
            """, (account_id, category, max_date))
        else:
            cursor.execute(f"""
                SELECT path, data_time 
                FROM {TABLE_USER_PREFERENCE_LOG} 
                WHERE account_id = %s AND category = %s  AND style is not null
                ORDER BY data_time DESC
            """, (account_id, category))
        
        like_records = cursor.fetchall()
        
        if not like_records:
            return None
        
        # 2. 批量查询点赞次数（如果指定了max_date，只统计该日期之前的点赞）
        paths = [r[0] for r in like_records]
        if not paths:
            return None
            
        placeholders = ','.join(['%s'] * len(paths))
        if max_date:
            cursor.execute(f"""
                SELECT path, COUNT(*) as like_count 
                FROM {TABLE_USER_PREFERENCE_LOG} 
                WHERE account_id = %s AND category = %s AND path IN ({placeholders})
                  AND data_time < %s
                GROUP BY path
            """, (account_id, category) + tuple(paths) + (max_date,))
        else:
            cursor.execute(f"""
                SELECT path, COUNT(*) as like_count 
                FROM {TABLE_USER_PREFERENCE_LOG} 
                WHERE account_id = %s AND category = %s AND path IN ({placeholders})
                GROUP BY path
            """, (account_id, category) + tuple(paths))
        
        like_counts = {row[0]: row[1] for row in cursor.fetchall()}
        
        # 3. 批量获取向量
        vectors_dict = query_vectors_by_paths(paths)
        
        # 处理查询不到的 path（用户图或异常情况）
        missing_paths = [p for p in paths if p not in vectors_dict]
        if missing_paths:
            logger.info(f"用户 {account_id} 类别 {category} 有 {len(missing_paths)} 个 path 需要实时计算向量")
            # 目前未有非系统图向量，跳过
            # 这里可以实时计算并写入 Milvus，但为了简化，先跳过
            # 实际实现中应该调用 vector_utils.extract_feature_vector 并写入 Milvus
        
        # 4. 计算权重并加权平均
        vectors = []
        weights = []
        K_half = RECOMMENDATION_CONFIG["K_half"]
        
        for k, (path, data_time) in enumerate(like_records, 1):
            if path not in vectors_dict:
                continue
            
            vector_data = vectors_dict[path]
            feature_vector = np.array(vector_data["feature_vector"])
            
            # 时间衰减权重
            d_k = 0.5 ** (k / K_half)
            
            # 点赞次数权重
            like_count = like_counts.get(path, 1)
            p_i = 1 + math.log(1 + like_count)
            
            # 综合权重
            w_i = d_k * p_i
            # w_i = p_i

            vectors.append(feature_vector)
            weights.append(w_i)
        
        if not vectors:
            return None
        
        # 5. 计算加权平均并做 L2 归一化，IP≈cosine
        preference_vector = compute_weighted_average(vectors, weights)
        preference_vector = normalize_vector(preference_vector)
        
        return preference_vector
        
    except Exception as e:
        logger.error(f"计算用户偏好向量失败 [user={account_id}, category={category}]: {e}", exc_info=True)
        return None
    finally:
        if should_close and conn:
            conn.close()


def generate_initial_user_preference_vectors(batch_size: int = 100):
    """
    初始用户偏好向量生成
    
    Args:
        batch_size: 每批处理用户数
    """
    conn = None
    try:
        conn = pymysql.connect(**MYSQL_CONFIG)
        cursor = conn.cursor()
        
        # 1. 扫描历史数据
        logger.info("扫描用户和类别组合...")
        cursor.execute(f"""
            SELECT DISTINCT account_id, category 
            FROM {TABLE_USER_PREFERENCE_LOG} 
            WHERE category IS NOT NULL
            AND style IS NOT NULL
        """)
        
        user_categories = cursor.fetchall()
        logger.info(f"找到 {len(user_categories)} 个用户-类别组合")
        
        if not user_categories:
            logger.warning("没有找到用户-类别组合")
            return
        
        # 2. 批量处理
        processed = 0
        failed = 0
        
        for account_id, category in user_categories:
            try:
                # 计算偏好向量
                preference_vector = compute_user_preference_vector(account_id, category, conn)
                
                if preference_vector is not None:
                    # 写入 Redis
                    key = f"{REDIS_KEY_USER_PREF_PREFIX}:{account_id}:{category}"
                    # 序列化向量（使用 JSON）
                    vector_json = json.dumps(preference_vector.tolist())
                    Redis.write(
                        key=key,
                        value=vector_json,
                        expire=RECOMMENDATION_CONFIG["redis_expire_seconds"]
                    )
                    processed += 1
                else:
                    failed += 1
                
                if (processed + failed) % batch_size == 0:
                    logger.info(f"已处理 {processed + failed}/{len(user_categories)} 个组合，成功: {processed}, 失败: {failed}")
                    
            except Exception as e:
                logger.error(f"处理失败 [user={account_id}, category={category}]: {e}")
                failed += 1
        
        logger.info(f"初始用户偏好向量生成完成，成功: {processed}, 失败: {failed}")
        
    except Exception as e:
        logger.error(f"初始用户偏好向量生成失败: {e}", exc_info=True)
    finally:
        if conn:
            conn.close()


def run_precompute():
    """
    运行所有预计算任务
    """
    logger.info("=" * 50)
    logger.info("开始预计算任务")
    logger.info("=" * 50)
    
    # 1. 优化数据库表结构
    logger.info("\n[1/5] 优化数据库表结构...")
    optimize_database_table()
    
    # # 2. 创建 Milvus 集合
    # logger.info("\n[2/5] 创建 Milvus 集合...")
    # create_collection()
    
    # 3. 历史数据迁移
    logger.info("\n[3/5] 历史数据迁移...")
    migrate_historical_data()
    
    # # 4. 系统图向量预计算
    # logger.info("\n[4/5] 系统图向量预计算...")
    # precompute_system_sketch_vectors()
    
    # 5. 初始用户偏好向量生成
    logger.info("\n[5/5] 初始用户偏好向量生成...")
    generate_initial_user_preference_vectors()
    
    logger.info("=" * 50)
    logger.info("预计算任务完成")
    logger.info("=" * 50)


if __name__ == "__main__":
    # 1. 优化数据库表结构
    logger.info("\n[1/5] 优化数据库表结构...")
    optimize_database_table()

    # 3. 历史数据迁移
    logger.info("\n[3/5] 历史数据迁移...")
    migrate_historical_data()

    # 5. 初始用户偏好向量生成
    logger.info("\n[5/5] 初始用户偏好向量生成...")
    generate_initial_user_preference_vectors()