"""Кластеризация персонажей в стилометрическом пространстве."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Sequence

import numpy as np
import pandas as pd
from scipy.sparse import issparse
from sklearn.cluster import AgglomerativeClustering, KMeans
from sklearn.metrics import (
    adjusted_rand_score,
    calinski_harabasz_score,
    davies_bouldin_score,
    normalized_mutual_info_score,
    silhouette_score,
)
from sklearn.preprocessing import normalize

from .pipelines import FeatureConfig, build_feature_union


@dataclass
class ClusteringResult:
    metrics: pd.DataFrame
    centroids: pd.DataFrame


def _centroid_matrix(matrix: np.ndarray, labels: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    classes = np.unique(labels)
    centroids = np.vstack([matrix[labels == label].mean(axis=0) for label in classes])
    return classes, centroids


def run_clustering(
    X: pd.DataFrame,
    y: Sequence,
    feature_config: FeatureConfig,
    n_clusters: int | None = None,
    seed: int = 42,
) -> ClusteringResult:
    """Кластеризует центроиды персонажей и считает внешние метрики."""
    y = np.asarray(y)
    features = build_feature_union(feature_config)
    matrix = features.fit_transform(X)
    matrix = matrix.toarray() if issparse(matrix) else np.asarray(matrix)
    matrix = normalize(matrix, norm="l2", copy=False)
    classes, centroids = _centroid_matrix(matrix, y)
    # Если число кластеров равно числу персонажей, кластеризация тривиальна,
    # а силуэт не определён. По умолчанию берём примерно корень из числа
    # классов: для 34 персонажей это 6 групп.
    n_clusters = int(n_clusters or max(2, round(len(classes) ** 0.5)))
    n_clusters = min(n_clusters, len(classes) - 1)

    rows: list[dict[str, Any]] = []
    for name, model in (
        (
            "agglomerative_cosine",
            AgglomerativeClustering(
                n_clusters=n_clusters, metric="cosine", linkage="average"
            ),
        ),
        (
            "kmeans",
            KMeans(n_clusters=n_clusters, random_state=seed, n_init=20),
        ),
    ):
        predicted = model.fit_predict(centroids)
        n_predicted = len(set(predicted))
        valid_partition = 2 <= n_predicted < len(centroids)
        rows.append(
            {
                "method": name,
                "n_clusters": n_clusters,
                "adjusted_rand_index": float(adjusted_rand_score(classes, predicted)),
                "normalized_mutual_info": float(normalized_mutual_info_score(classes, predicted)),
                "silhouette_cosine": float(
                    silhouette_score(centroids, predicted, metric="cosine")
                ) if valid_partition else 0.0,
                "calinski_harabasz": float(calinski_harabasz_score(centroids, predicted))
                if valid_partition else 0.0,
                "davies_bouldin": float(davies_bouldin_score(centroids, predicted))
                if valid_partition else 0.0,
            }
        )

    centroid_frame = pd.DataFrame(centroids, index=classes)
    centroid_frame.index.name = "character"
    return ClusteringResult(metrics=pd.DataFrame(rows), centroids=centroid_frame)