import numpy as np
from sklearn.cluster import KMeans
from dataclasses import dataclass
from typing import List, Dict, Any
from .embedder import Embedder

@dataclass
class Index:
    embedder: Embedder
    docs: List[str]
    doc_vecs: Any
    topic_labels: np.ndarray
    topics: Dict[int, list]

def build_index(docs: List[str], n_topics: int = 5) -> Index:
    emb = Embedder().fit(docs)
    X = emb.doc_mat
    km = KMeans(n_clusters=n_topics, n_init=10, random_state=0)
    labels = km.fit_predict(X)
    topics = {}
    for i, t in enumerate(labels):
        topics.setdefault(t, []).append(i)
    return Index(emb, docs, X, labels, topics)
