# main_enhanced.py
import argparse
from batch_processor import BatchProcessor
from topic_cluster import TopicCluster
from corpus_builder import E


def main_enhanced():
    parser = argparse.ArgumentParser(description="增强版文献语料库")
    parser.add_argument('--batch', action='store_true', help='批量处理模式')
    parser.add_argument('--cluster', type=int, help='聚类分析，指定簇数量')
    parser.add_argument('--llm', action='store_true', help='启用LLM功能')
    parser.add_argument('--visualize', action='store_true', help='可视化聚类')

    args = parser.parse_args()

    # 初始化增强版语料库
    corpus = EnhancedLiteratureCorpus(use_llm=args.llm)

    if args.batch:
        print("启动批量处理...")
        from batch_processor import batch_process_main
        batch_process_main()

    if args.cluster:
        print("执行主题聚类...")
        cluster_analyzer = TopicCluster(corpus)
        cluster_data = cluster_analyzer.cluster_documents(n_clusters=args.cluster)

        if cluster_data:
            # 保存聚类结果
            with open('clusters.json', 'w') as f:
                json.dump(cluster_data, f)

            if args.visualize:
                fig = cluster_analyzer.visualize_clusters(cluster_data)
                fig.write_html('clusters_visualization.html')
                print("聚类可视化已保存")

    # 启动增强版应用
    if args.llm:
        from app import EnhancedLiteratureSearchApp
        app = EnhancedLiteratureSearchApp(corpus)
        app.run()