"""
Clustering and embedding analysis tools.
"""

from dataclasses import dataclass
from typing import Any, Literal

import numpy as np


@dataclass
class ClusteringResult:
    """Results from clustering analysis."""

    embeddings: np.ndarray  # (N, D) original embeddings
    projections: np.ndarray  # (N, 2) or (N, 3) reduced projections
    cluster_ids: np.ndarray  # (N,) cluster assignments (-1 = noise)
    sample_ids: list[str]  # Sample identifiers
    labels: list[str] | None = None  # Optional ground truth labels

    @property
    def num_clusters(self) -> int:
        """Number of clusters (excluding noise points)."""
        unique = set(self.cluster_ids)
        return len(unique) - (1 if -1 in unique else 0)

    @property
    def num_samples(self) -> int:
        return len(self.sample_ids)

    def get_cluster(self, cluster_id: int) -> tuple[np.ndarray, list[str]]:
        """Get embeddings and IDs for a specific cluster."""
        mask = self.cluster_ids == cluster_id
        return (
            self.embeddings[mask],
            [sid for sid, m in zip(self.sample_ids, mask) if m],
        )

    def get_noise_points(self) -> tuple[np.ndarray, list[str]]:
        """Get embeddings and IDs for noise points (cluster_id = -1)."""
        return self.get_cluster(-1)

    def cluster_sizes(self) -> dict[int, int]:
        """Get size of each cluster."""
        unique, counts = np.unique(self.cluster_ids, return_counts=True)
        return dict(zip(unique.tolist(), counts.tolist()))

    def to_dataframe(self):
        """Convert to pandas DataFrame."""
        try:
            import pandas as pd
        except ImportError:
            raise ImportError("pandas required: pip install pandas")

        data = {
            "sample_id": self.sample_ids,
            "cluster_id": self.cluster_ids.tolist(),
        }

        # Add projection coordinates
        for i in range(self.projections.shape[1]):
            data[f"proj_{i}"] = self.projections[:, i].tolist()

        if self.labels is not None:
            data["label"] = self.labels

        return pd.DataFrame(data)


def cluster_embeddings(
    embeddings: np.ndarray,
    sample_ids: list[str],
    method: Literal["hdbscan", "kmeans", "agglomerative"] = "hdbscan",
    n_clusters: int | None = None,
    projection: Literal["umap", "tsne", "pca"] = "umap",
    projection_dim: int = 2,
    labels: list[str] | None = None,
    **kwargs,
) -> ClusteringResult:
    """
    Cluster embeddings and compute 2D/3D projections for visualization.

    Args:
        embeddings: (N, D) array of embeddings
        sample_ids: List of sample identifiers
        method: Clustering method ('hdbscan', 'kmeans', 'agglomerative')
        n_clusters: Number of clusters (required for kmeans/agglomerative)
        projection: Dimensionality reduction method
        projection_dim: Output dimensions (2 or 3)
        labels: Optional ground truth labels
        **kwargs: Additional arguments for clustering/projection

    Returns:
        ClusteringResult with projections and cluster assignments

    Example:
        >>> embeddings = np.random.randn(100, 256)
        >>> ids = [f"sample_{i}" for i in range(100)]
        >>> result = cluster_embeddings(embeddings, ids)
        >>> print(f"Found {result.num_clusters} clusters")
    """
    # Compute projections
    projections = reduce_dimensions(
        embeddings,
        method=projection,
        n_components=projection_dim,
        **kwargs.get("projection_params", {}),
    )

    # Clustering
    cluster_ids = _cluster(
        embeddings,
        method=method,
        n_clusters=n_clusters,
        **kwargs.get("cluster_params", {}),
    )

    return ClusteringResult(
        embeddings=embeddings,
        projections=projections,
        cluster_ids=cluster_ids,
        sample_ids=sample_ids,
        labels=labels,
    )


def find_similar(
    query: np.ndarray,
    database: np.ndarray,
    sample_ids: list[str],
    k: int = 5,
    metric: str = "cosine",
) -> list[tuple[str, float]]:
    """
    Find k most similar samples to query.

    Args:
        query: (D,) or (1, D) query embedding
        database: (N, D) database embeddings
        sample_ids: List of database sample IDs
        k: Number of neighbors to return
        metric: Distance metric ('cosine', 'euclidean')

    Returns:
        List of (sample_id, distance) tuples, sorted by distance

    Example:
        >>> query = embeddings[0]
        >>> similar = find_similar(query, embeddings, ids, k=5)
        >>> for sample_id, distance in similar:
        ...     print(f"{sample_id}: {distance:.3f}")
    """
    try:
        from sklearn.neighbors import NearestNeighbors
    except ImportError:
        raise ImportError("scikit-learn required: pip install sltk[analysis]")

    query = np.atleast_2d(query)

    nn = NearestNeighbors(n_neighbors=min(k, len(database)), metric=metric)
    nn.fit(database)

    distances, indices = nn.kneighbors(query)

    results = []
    for dist, idx in zip(distances[0], indices[0]):
        results.append((sample_ids[idx], float(dist)))

    return results


def reduce_dimensions(
    embeddings: np.ndarray,
    method: Literal["umap", "tsne", "pca"] = "umap",
    n_components: int = 2,
    **kwargs,
) -> np.ndarray:
    """
    Reduce embedding dimensions for visualization.

    Args:
        embeddings: (N, D) embeddings
        method: Reduction method
        n_components: Output dimensions
        **kwargs: Method-specific parameters

    Returns:
        (N, n_components) projected embeddings
    """
    if method == "umap":
        try:
            import umap
        except ImportError:
            raise ImportError("umap-learn required: pip install sltk[analysis]")

        defaults = {"n_neighbors": 15, "min_dist": 0.1, "metric": "cosine"}
        defaults.update(kwargs)

        reducer = umap.UMAP(n_components=n_components, **defaults)
        return reducer.fit_transform(embeddings)

    elif method == "tsne":
        try:
            from sklearn.manifold import TSNE
        except ImportError:
            raise ImportError("scikit-learn required: pip install sltk[analysis]")

        defaults = {"perplexity": min(30, len(embeddings) - 1), "n_iter": 1000}
        defaults.update(kwargs)

        reducer = TSNE(n_components=n_components, **defaults)
        return reducer.fit_transform(embeddings)

    elif method == "pca":
        try:
            from sklearn.decomposition import PCA
        except ImportError:
            raise ImportError("scikit-learn required: pip install sltk[analysis]")

        reducer = PCA(n_components=n_components, **kwargs)
        return reducer.fit_transform(embeddings)

    else:
        raise ValueError(f"Unknown method: {method}")


class PoseClusterer:
    """
    Cluster poses to identify common postures/configurations.

    Useful for analyzing sign patterns, finding representative poses,
    and reducing pose sequence complexity.

    Example:
        >>> clusterer = PoseClusterer(n_clusters=10, method='kmeans')
        >>> clusterer.fit(pose_sequence)
        >>> labels = clusterer.predict(pose_sequence)
        >>> centers = clusterer.get_cluster_centers()
    """

    def __init__(
        self,
        n_clusters: int = 10,
        method: Literal["kmeans", "agglomerative", "hdbscan"] = "kmeans",
        **kwargs,
    ):
        """
        Initialize pose clusterer.

        Args:
            n_clusters: Number of clusters (ignored for hdbscan)
            method: Clustering method
            **kwargs: Additional parameters passed to clustering algorithm
        """
        self.n_clusters = n_clusters
        self.method = method
        self.kwargs = kwargs
        self._clusterer: Any = None
        self._fitted = False

    def fit(self, poses) -> "PoseClusterer":
        """
        Fit clusterer to pose data.

        Args:
            poses: PoseSequence or numpy array of shape (T, N, C)

        Returns:
            self for chaining
        """
        from sltk.data.core import PoseSequence

        if isinstance(poses, PoseSequence):
            data = poses.data
        else:
            data = poses

        # Flatten each frame to 1D: (T, N*C)
        T = data.shape[0]
        flat_data = data.reshape(T, -1)

        # Remove invalid frames (all zeros)
        valid_mask = np.any(flat_data != 0, axis=1)
        valid_data = flat_data[valid_mask]

        if len(valid_data) == 0:
            raise ValueError("No valid pose data to cluster")

        # Create and fit clusterer
        if self.method == "kmeans":
            try:
                from sklearn.cluster import KMeans
            except ImportError:
                raise ImportError("scikit-learn required: pip install scikit-learn")

            self._clusterer = KMeans(
                n_clusters=min(self.n_clusters, len(valid_data)), random_state=42, **self.kwargs
            )
            self._clusterer.fit(valid_data)

        elif self.method == "agglomerative":
            try:
                from sklearn.cluster import AgglomerativeClustering
            except ImportError:
                raise ImportError("scikit-learn required: pip install scikit-learn")

            self._clusterer = AgglomerativeClustering(
                n_clusters=min(self.n_clusters, len(valid_data)), **self.kwargs
            )
            self._clusterer.fit(valid_data)
            self._cluster_labels = self._clusterer.labels_

        elif self.method == "hdbscan":
            try:
                import hdbscan
            except ImportError:
                raise ImportError("hdbscan required: pip install hdbscan")

            defaults = {"min_cluster_size": 5}
            defaults.update(self.kwargs)
            self._clusterer = hdbscan.HDBSCAN(**defaults)
            self._clusterer.fit(valid_data)

        else:
            raise ValueError(f"Unknown method: {self.method}")

        self._valid_data = valid_data
        self._valid_mask = valid_mask
        self._fitted = True
        return self

    def predict(self, poses) -> np.ndarray:
        """
        Predict cluster labels for pose data.

        Args:
            poses: PoseSequence or numpy array

        Returns:
            Array of cluster labels per frame (-1 for invalid frames)
        """
        if not self._fitted:
            raise RuntimeError("Must call fit() before predict()")

        from sltk.data.core import PoseSequence

        if isinstance(poses, PoseSequence):
            data = poses.data
        else:
            data = poses

        T = data.shape[0]
        flat_data = data.reshape(T, -1)

        labels = np.full(T, -1, dtype=np.int32)
        valid_mask = np.any(flat_data != 0, axis=1)
        valid_data = flat_data[valid_mask]

        if len(valid_data) == 0:
            return labels

        if self.method == "kmeans":
            valid_labels = self._clusterer.predict(valid_data)
        elif self.method == "agglomerative":
            # Agglomerative doesn't have predict, use nearest centroid
            try:
                from sklearn.metrics import pairwise_distances
            except ImportError:
                raise ImportError("scikit-learn required: pip install scikit-learn")

            centers = self.get_cluster_centers()
            dists = pairwise_distances(valid_data, centers)
            valid_labels = np.argmin(dists, axis=1)
        elif self.method == "hdbscan":
            try:
                import hdbscan
            except ImportError:
                raise ImportError("hdbscan required: pip install hdbscan")

            valid_labels, _ = hdbscan.approximate_predict(self._clusterer, valid_data)
        else:
            raise ValueError(f"Unknown method: {self.method}")

        labels[valid_mask] = valid_labels
        return labels

    def get_cluster_centers(self) -> np.ndarray:
        """
        Get cluster center poses.

        Returns:
            Array of shape (n_clusters, D) with cluster centers
        """
        if not self._fitted:
            raise RuntimeError("Must call fit() first")

        if self.method == "kmeans":
            return self._clusterer.cluster_centers_
        else:
            # Compute centers from data
            labels = self._clusterer.labels_
            n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
            centers = []
            for i in range(n_clusters):
                mask = labels == i
                if mask.any():
                    centers.append(self._valid_data[mask].mean(axis=0))
            return np.array(centers)


def _cluster(
    embeddings: np.ndarray,
    method: str,
    n_clusters: int | None,
    **kwargs,
) -> np.ndarray:
    """Perform clustering."""
    if method == "hdbscan":
        try:
            import hdbscan
        except ImportError:
            raise ImportError("hdbscan required: pip install sltk[analysis]")

        defaults = {"min_cluster_size": 5}
        defaults.update(kwargs)

        clusterer = hdbscan.HDBSCAN(**defaults)
        return clusterer.fit_predict(embeddings)

    elif method == "kmeans":
        try:
            from sklearn.cluster import KMeans
        except ImportError:
            raise ImportError("scikit-learn required: pip install sltk[analysis]")

        if n_clusters is None:
            raise ValueError("n_clusters required for kmeans")

        clusterer = KMeans(n_clusters=n_clusters, random_state=42, **kwargs)
        return clusterer.fit_predict(embeddings)

    elif method == "agglomerative":
        try:
            from sklearn.cluster import AgglomerativeClustering
        except ImportError:
            raise ImportError("scikit-learn required: pip install sltk[analysis]")

        if n_clusters is None:
            raise ValueError("n_clusters required for agglomerative")

        clusterer = AgglomerativeClustering(n_clusters=n_clusters, **kwargs)
        return clusterer.fit_predict(embeddings)

    else:
        raise ValueError(f"Unknown clustering method: {method}")
