"""Protocol for fitted Agglomerative-like models."""

from __future__ import annotations

from typing import Literal, Protocol, runtime_checkable

import numpy as np

AgglomerativeMetric = Literal[
    "braycurtis",
    "canberra",
    "chebychev",
    "cityblock",
    "cosine",
    "dice",
    "euclidean",
    "hamming",
    "jaccard",
    "mahalanobis",
    "minkowski",
    "russellrao",
    "seuclidean",
    "sokalsneath",
    "sqeuclidean",
]


@runtime_checkable
class AgglomerativeModel(Protocol):
    """Protocol for fitted Agglomerative-like models.

    Both CPU (sklearn) and CUDA (custom) implementations must satisfy this.
    """

    @property
    def cluster_centers_(self) -> np.ndarray:
        """Cluster centers of shape (n_clusters, n_features)."""
        ...

    @property
    def labels_(self) -> np.ndarray:
        """Labels assigned during fit of shape (n_samples,)."""
        ...

    @property
    def n_clusters_(self) -> int:
        """Number of clusters."""
        ...

    def predict(self, embeddings: np.ndarray) -> np.ndarray:
        """Predict cluster assignments."""
        ...
