# topic_cluster.py
from sklearn.cluster import KMeans, DBSCAN
from sklearn.decomposition import PCA
import numpy as np
import plotly.graph_objects as go
from typing import List, Dict
import pandas as pd


class TopicCluster:
    """主题聚类分析"""

    def __init__(self, corpus):
        self.corpus = corpus

    def cluster_documents(self, n_clusters=5, use_pca=True):
        """聚类文档"""
        # 获取所有嵌入向量
        all_docs = self.corpus.collection.get(include=["embeddings", "metadatas"])

        if not all_docs['embeddings']:
            return None

        embeddings = np.array(all_docs['embeddings'])

        # 降维（可选）
        if use_pca and embeddings.shape[1] > 2:
            pca = PCA(n_components=2)
            reduced_embeddings = pca.fit_transform(embeddings)
        else:
            reduced_embeddings = embeddings[:, :2]

        # 聚类
        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        clusters = kmeans.fit_predict(embeddings)

        # 分析每个簇的主题
        cluster_topics = self._analyze_cluster_topics(
            clusters,
            all_docs['documents']
        )

        return {
            'clusters': clusters.tolist(),
            'reduced_embeddings': reduced_embeddings.tolist(),
            'cluster_topics': cluster_topics,
            'documents': all_docs['documents'],
            'metadatas': all_docs['metadatas']
        }

    def _analyze_cluster_topics(self, clusters, documents):
        """分析每个簇的主题"""
        from collections import Counter
        import re

        unique_clusters = np.unique(clusters)
        topics = {}

        for cluster_id in unique_clusters:
            # 获取该簇的所有文档
            cluster_docs = [
                doc for doc, c in zip(documents, clusters)
                if c == cluster_id
            ]

            # 提取高频词
            all_text = " ".join(cluster_docs)
            words = re.findall(r'\b[a-zA-Z]{4,}\b', all_text.lower())

            # 移除停用词
            stop_words = set(['this', 'that', 'these', 'those', 'with', 'from', 'have', 'were', 'which'])
            filtered_words = [w for w in words if w not in stop_words]

            word_counts = Counter(filtered_words).most_common(5)
            topics[cluster_id] = [word for word, count in word_counts]

        return topics

    def visualize_clusters(self, cluster_data):
        """可视化聚类结果"""
        fig = go.Figure()

        clusters = cluster_data['clusters']
        embeddings = cluster_data['reduced_embeddings']
        topics = cluster_data['cluster_topics']

        # 创建散点图
        for cluster_id in np.unique(clusters):
            mask = np.array(clusters) == cluster_id
            cluster_points = np.array(embeddings)[mask]

            fig.add_trace(go.Scatter(
                x=cluster_points[:, 0],
                y=cluster_points[:, 1],
                mode='markers',
                name=f'Cluster {cluster_id}: {", ".join(topics[cluster_id][:3])}',
                text=[f"主题: {topics[cluster_id]}" for _ in range(len(cluster_points))],
                hovertemplate='%{text}<extra></extra>'
            ))

        fig.update_layout(
            title='文献主题聚类',
            xaxis_title='Component 1',
            yaxis_title='Component 2',
            showlegend=True
        )

        return fig