#!/usr/bin/env python
# coding=utf-8
"""
标题相关性分析（v2 简化版）
批量处理多天数据，计算相似度
"""
import pandas as pd
import numpy as np
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
from lib.text_embedding_api import compare_phrases_batch

task_dir = Path(__file__).parent
output_dir = task_dir / "output"
data_dir = output_dir / "标题相关性分析"

# 找到所有需要处理的数据文件
if data_dir.exists():
    csv_files = sorted(data_dir.glob("*.csv"))
else:
    csv_files = sorted(output_dir.glob("*.csv"))
    csv_files = [f for f in csv_files if '_含相似度' not in f.name]

if not csv_files:
    print("没有找到数据文件")
    exit(1)

print(f"找到 {len(csv_files)} 个数据文件")

BATCH_SIZE = 500

def calc_similarity(df):
    """计算标题相似度"""
    similarity_configs = [
        ('文章标题', 'title', '文章标题_视频标题_相似度'),
        ('分享标题', 'title', '分享标题_视频标题_相似度'),
    ]

    for col1, col2, result_col in similarity_configs:
        if result_col in df.columns and df[result_col].notna().any():
            # 已有相似度数据，跳过
            continue

        df[result_col] = np.nan

        pairs = []
        valid_indices = []

        for idx, row in df.iterrows():
            text1 = str(row[col1]) if pd.notna(row[col1]) and row[col1] != '' else ''
            text2 = str(row[col2]) if pd.notna(row[col2]) and row[col2] != '' else ''

            if text1 and text2:
                pairs.append((text1, text2))
                valid_indices.append(idx)

        if not pairs:
            continue

        print(f"  计算 {result_col}: {len(pairs)} 对")

        scores = []
        for i in range(0, len(pairs), BATCH_SIZE):
            batch = pairs[i:i+BATCH_SIZE]
            results = compare_phrases_batch(batch)
            scores.extend([r['相似度'] for r in results])
            if (i + BATCH_SIZE) % 10000 == 0:
                print(f"    已处理 {min(i+BATCH_SIZE, len(pairs))}/{len(pairs)}")

        for idx, score in zip(valid_indices, scores):
            df.at[idx, result_col] = score

    return df


# 处理每个文件
for csv_file in csv_files:
    dt = csv_file.stem  # 日期如 20260107
    output_file = output_dir / f"{dt}_含相似度.csv"

    # 检查是否已处理
    if output_file.exists():
        print(f"[{dt}] 已存在，跳过")
        continue

    print(f"[{dt}] 处理中...")
    df = pd.read_csv(csv_file)
    print(f"  记录数: {len(df):,}")

    df = calc_similarity(df)

    df.to_csv(output_file, index=False)
    print(f"  已保存: {output_file.name}")

print("\n全部完成!")
