# main.py
import os
import numpy as np
import pandas as pd
from bertopic import BERTopic
from umap import UMAP
from hdbscan import HDBSCAN
from sklearn.feature_extraction.text import CountVectorizer
import time

def run_pipeline():
    # --- 1. 路径配置 ---
    RAW_TEXT_PATH = './data/文本.txt'      # 原始文本 (用于展示)
    CUT_DATA_PATH = './data/切词.txt'      # 分词文本 (用于特征提取)
    EMBEDDING_PATH = './data/embedding.npy' # 词向量
    RESULT_DIR = './result'
    
    if not os.path.exists(RESULT_DIR):
        os.makedirs(RESULT_DIR)

    # --- 2. 加载并执行三重对齐 ---
    print(f"[{time.strftime('%H:%M:%S')}] 正在加载并校验数据一致性...")
    
    # 加载原始文本 (必须与 embedding.py 的过滤逻辑完全一致)
    with open(RAW_TEXT_PATH, 'r', encoding='utf-8') as f:
        full_texts = [line.strip() for line in f.readlines() if line.strip()]
        
    # 加载分词文本
    with open(CUT_DATA_PATH, 'r', encoding='utf-8') as f:
        segmented_docs = [line.strip() for line in f.readlines() if line.strip()]
        
    # 加载向量
    embeddings = np.load(EMBEDDING_PATH)

    # 【关键步骤】自动对齐：取三者最小长度，防止任何一方因为运行中断或文件更新导致的索引偏离
    min_len = min(len(full_texts), len(segmented_docs), len(embeddings))
    
    if len(full_texts) != len(segmented_docs) or len(full_texts) != len(embeddings):
        print(f"⚠️ 警告: 数据长度不一致! 文本:{len(full_texts)}, 切词:{len(segmented_docs)}, 向量:{len(embeddings)}")
        print(f"⚠️ 系统将自动按最小长度 {min_len} 进行截断对齐。")
        full_texts = full_texts[:min_len]
        segmented_docs = segmented_docs[:min_len]
        embeddings = embeddings[:min_len]
    else:
        print(f"✅ 三方数据完美对齐: {min_len} 条记录")

    # --- 3. 初始化模型组件 ---
    umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine', random_state=30)
    hdbscan_model = HDBSCAN(min_cluster_size=10, min_samples=5, metric='euclidean', prediction_data=True)
    
    # 既然是中文，token_pattern 必须处理好空格分词
    vectorizer_model = CountVectorizer(
        token_pattern=r"(?u)[^ ]+", 
        stop_words=['洛阳', '旅游', '文化', '我们', '一个'] 
    )

    # --- 4. 训练 BERTopic ---
    print(f"[{time.strftime('%H:%M:%S')}] 开始训练主题模型...")
    topic_model = BERTopic(
        embedding_model=None, # 强制离线，等价于导入模型计算结果
        umap_model=umap_model,
        hdbscan_model=hdbscan_model,
        vectorizer_model=vectorizer_model,
        verbose=True
    )
    
    # 使用分词后的 segmented_docs 训练模型，使用由原文生成的 embeddings
    topics, probs = topic_model.fit_transform(segmented_docs, embeddings=embeddings)

    # --- 5. 结果导出 (将原始文本与分类结果合并) ---
    print(f"[{time.strftime('%H:%M:%S')}] 正在生成导出结果...")
    
    # 主题概览
    topic_model.get_topic_info().to_csv(os.path.join(RESULT_DIR, '主题概览.csv'), index=False, encoding='utf-8-sig')
    
    # 核心修改：生成详情表时，将 segmented_docs 换成 full_texts
    # 这样生成的“聚类结果.csv”中 Document 列就是可读的原文，而不是碎词
    doc_info = topic_model.get_document_info(full_texts) 
    doc_info.to_csv(os.path.join(RESULT_DIR, '聚类结果.csv'), index=False, encoding='utf-8-sig')

    # --- 6. 生成可视化 ---
    print(f"[{time.strftime('%H:%M:%S')}] 正在生成交互式图表...")
    try:
        topic_model.visualize_topics().write_html(os.path.join(RESULT_DIR, "topics_distance.html"))
        topic_model.visualize_barchart(top_n_topics=15).write_html(os.path.join(RESULT_DIR, "topics_barchart.html"))
        print(f"✅ 可视化 HTML 已保存至 {RESULT_DIR}")
    except Exception as e:
        print(f"❌ 可视化失败: {e}")

    print(f"\n🎉 运行结束！结果保存在 ./result 文件夹中。")

if __name__ == "__main__":
    run_pipeline()