"""
QKMeans: Fast k-means++ seeding using rejection sampling.

This module provides a scikit-learn compatible interface to the QKMeans
algorithm for fast k-means++ seeding under the manifold hypothesis.
"""

import numpy as np
from sklearn.base import BaseEstimator, ClusterMixin, TransformerMixin
from sklearn.utils.validation import check_array, check_is_fitted

from . import _qkmeans_cpp


class QKMeansSeeding(BaseEstimator, ClusterMixin, TransformerMixin):
    """Fast k-means++ seeding using rejection sampling.

    This implements the QKMeans algorithm from "Faster k-means Seeding
    Under The Manifold Hypothesis" (ICML 2026).

    The algorithm uses rejection sampling to approximate the D² sampling
    distribution of k-means++, achieving O(nD) + O(D(k log k)^(1+δ)) runtime
    under the manifold hypothesis, compared to O(nkD) for standard k-means++.

    Parameters
    ----------
    n_clusters : int, default=8

    
        Number of clusters (k).

    anns : {'lsh', 'faiss', 'hnsw', 'brute_force'}, default='lsh'
        Approximate nearest neighbor search backend:
        - 'lsh': Custom LSH implementation (always available)
        - 'faiss': Facebook FAISS library (requires faiss-cpu/faiss-gpu)
        - 'hnsw': hnswlib HNSW graphs (requires hnswlib)
        - 'brute_force': Exact NN search (for debugging/small k)

    anns_params : dict or None, default=None
        Backend-specific parameters. If None, uses sensible defaults.

        LSH parameters:
            - num_tables : int, default=10
                Number of hash tables.
            - num_hashes : int, default=8
                Number of hash functions per table.
            - rho : float, default=0.5
                Approximation parameter (lower = more accurate but slower).

        HNSW parameters:
            - M : int, default=16
                Number of bi-directional links per node.
            - ef_construction : int, default=200
                Size of the dynamic list during construction.
            - ef_search : int, default=50
                Size of the dynamic list during search.

        FAISS parameters:
            - index_type : str, default='flat'
                FAISS index type: 'flat' (exact), 'hnsw' (graph-based),
                'lsh' (locality-sensitive hashing), 'ivf' (inverted file).
            - hnsw_M : int, default=32
                HNSW: number of links per node.
            - hnsw_efConstruction : int, default=40
                HNSW: construction-time search depth.
            - hnsw_efSearch : int, default=16
                HNSW: search-time depth.
            - lsh_nbits : int, default=0
                LSH: number of bits (0 = auto: 2*dim).
            - nlist : int, default=100
                IVF: number of inverted lists.
            - nprobe : int, default=10
                IVF: number of lists to probe during search.

    max_rejections : int or None, default=None
        Maximum rejection sampling iterations per center.
        If None, automatically set to m·ln(k) where m is the oversampling factor.

    random_state : int, RandomState, or None, default=None
        Random seed for reproducibility.

    n_jobs : int or None, default=None
        Number of parallel threads. None means all available.

    Attributes
    ----------
    cluster_centers_ : ndarray of shape (n_clusters, n_features)
        Coordinates of cluster centers.

    center_indices_ : ndarray of shape (n_clusters,)
        Indices of the data points chosen as centers.

    n_features_in_ : int
        Number of features seen during fit.

    seeding_time_ : float
        Time taken for seeding in milliseconds.

    preprocess_time_ : float
        Time taken for preprocessing in milliseconds.

    n_rejections_ : int
        Total number of rejected samples during seeding.

    n_fallbacks_ : int
        Number of times uniform fallback was used.

    Examples
    --------
    >>> import numpy as np
    >>> from qkmeans import QKMeansSeeding
    >>> X = np.random.randn(10000, 128)
    >>> seeder = QKMeansSeeding(n_clusters=100)
    >>> seeder.fit(X)
    QKMeansSeeding(n_clusters=100)
    >>> seeder.cluster_centers_.shape
    (100, 128)

    Notes
    -----
    The algorithm achieves O(ρ⁻² log k) approximation ratio in expected time
    O(nD) + O(D(k log k)^(1+δ)) where δ = ε + ρ and ε = 2/d is the
    quantization exponent under the manifold hypothesis.

    References
    ----------
    Anonymous (2026). Faster k-means Seeding Under The Manifold Hypothesis.
    ICML 2026 (Under Review).
    """

    _ANNS_MAP = {
        'lsh': _qkmeans_cpp.ANNSType.LSH,
        'faiss': _qkmeans_cpp.ANNSType.FAISS,
        'hnsw': _qkmeans_cpp.ANNSType.HNSW,
        'brute_force': _qkmeans_cpp.ANNSType.BRUTE_FORCE,
    }

    def __init__(
        self,
        n_clusters=8,
        *,
        anns='lsh',
        anns_params=None,
        max_rejections=None,
        random_state=None,
        n_jobs=None,
    ):
        self.n_clusters = n_clusters
        self.anns = anns
        self.anns_params = anns_params
        self.max_rejections = max_rejections
        self.random_state = random_state
        self.n_jobs = n_jobs

    def _validate_params(self):
        """Validate parameters."""
        if self.n_clusters < 1:
            raise ValueError(f"n_clusters must be >= 1, got {self.n_clusters}")

        if self.anns not in self._ANNS_MAP:
            raise ValueError(
                f"anns must be one of {list(self._ANNS_MAP.keys())}, "
                f"got '{self.anns}'"
            )

        # Check if backend is available
        anns_type = self._ANNS_MAP[self.anns]
        if not _qkmeans_cpp.is_anns_available(anns_type):
            raise RuntimeError(
                f"ANNS backend '{self.anns}' is not available. "
                "Please install the required dependencies or use 'lsh' or 'brute_force'."
            )

    def _build_config(self):
        """Build C++ configuration from Python parameters."""
        config = _qkmeans_cpp.QKMeansConfig()
        config.k = self.n_clusters
        config.anns_type = self._ANNS_MAP[self.anns]
        config.max_rejections = self.max_rejections or 0
        config.random_seed = self._get_seed()
        config.n_threads = self.n_jobs if self.n_jobs is not None else -1

        # Set ANNS-specific parameters
        anns_params = self.anns_params or {}

        if self.anns == 'lsh':
            config.lsh_num_tables = anns_params.get('num_tables', 10)
            config.lsh_num_hashes = anns_params.get('num_hashes', 8)
            config.lsh_rho = anns_params.get('rho', 0.5)
        elif self.anns == 'hnsw':
            config.hnsw_M = anns_params.get('M', 16)
            config.hnsw_ef_construction = anns_params.get('ef_construction', 200)
            config.hnsw_ef_search = anns_params.get('ef_search', 50)
        elif self.anns == 'faiss':
            # Map index type string to integer
            faiss_index_map = {'flat': 0, 'hnsw': 1, 'lsh': 2, 'ivf': 3}
            index_type_str = anns_params.get('index_type', 'flat').lower()
            if index_type_str not in faiss_index_map:
                raise ValueError(
                    f"FAISS index_type must be one of {list(faiss_index_map.keys())}, "
                    f"got '{index_type_str}'"
                )
            config.faiss_index_type = faiss_index_map[index_type_str]
            config.faiss_hnsw_M = anns_params.get('hnsw_M', 32)
            config.faiss_hnsw_efConstruction = anns_params.get('hnsw_efConstruction', 40)
            config.faiss_hnsw_efSearch = anns_params.get('hnsw_efSearch', 16)
            config.faiss_lsh_nbits = anns_params.get('lsh_nbits', 0)
            config.faiss_nlist = anns_params.get('nlist', 100)
            config.faiss_nprobe = anns_params.get('nprobe', 10)

        return config

    def _get_seed(self):
        """Get random seed from random_state parameter."""
        if self.random_state is None:
            return 0  # Let C++ use random device
        if isinstance(self.random_state, np.random.RandomState):
            return int(self.random_state.randint(0, 2**32))
        if hasattr(self.random_state, 'integers'):
            # numpy Generator
            return int(self.random_state.integers(0, 2**32))
        return int(self.random_state)

    def fit(self, X, y=None):
        """Compute k-means++ seeding.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : Ignored
            Not used, present for API consistency.

        Returns
        -------
        self : object
            Fitted estimator.
        """
        self._validate_params()

        # Validate and convert input
        X = check_array(X, dtype=np.float64, order='C', ensure_2d=True)

        n_samples, n_features = X.shape
        if self.n_clusters > n_samples:
            raise ValueError(
                f"n_clusters ({self.n_clusters}) must be <= n_samples ({n_samples})"
            )

        # Build configuration and run algorithm
        config = self._build_config()
        result = _qkmeans_cpp.qkmeans_seed(X, config)

        # Store results
        self.cluster_centers_ = np.array(result.centers).reshape(
            self.n_clusters, n_features
        )
        self.center_indices_ = np.array(result.center_indices, dtype=np.intp)
        self.n_features_in_ = n_features
        self.seeding_time_ = result.seeding_time_ms
        self.preprocess_time_ = result.preprocess_time_ms
        self.n_rejections_ = result.total_rejections
        self.n_fallbacks_ = result.fallback_samples
        self._total_cost = result.total_cost

        return self

    def predict(self, X):
        """Predict the closest cluster for each sample.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            New data to predict.

        Returns
        -------
        labels : ndarray of shape (n_samples,)
            Index of the cluster each sample belongs to.
        """
        check_is_fitted(self)
        X = check_array(X, dtype=np.float64)

        if X.shape[1] != self.n_features_in_:
            raise ValueError(
                f"X has {X.shape[1]} features, expected {self.n_features_in_}"
            )

        # Compute squared distances to all centers
        # Using broadcasting: (n, 1, d) - (1, k, d) -> (n, k, d) -> sum -> (n, k)
        diff = X[:, np.newaxis, :] - self.cluster_centers_[np.newaxis, :, :]
        distances_sq = np.sum(diff ** 2, axis=2)

        return np.argmin(distances_sq, axis=1)

    def fit_predict(self, X, y=None):
        """Compute seeding and predict cluster labels.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : Ignored
            Not used, present for API consistency.

        Returns
        -------
        labels : ndarray of shape (n_samples,)
            Cluster labels.
        """
        return self.fit(X, y).predict(X)

    def transform(self, X):
        """Transform X to cluster-distance space.

        Each dimension is the distance to a cluster center.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data to transform.

        Returns
        -------
        X_new : ndarray of shape (n_samples, n_clusters)
            X transformed to cluster-distance space.
        """
        check_is_fitted(self)
        X = check_array(X, dtype=np.float64)

        if X.shape[1] != self.n_features_in_:
            raise ValueError(
                f"X has {X.shape[1]} features, expected {self.n_features_in_}"
            )

        # Compute distances to all centers
        diff = X[:, np.newaxis, :] - self.cluster_centers_[np.newaxis, :, :]
        return np.sqrt(np.sum(diff ** 2, axis=2))

    def fit_transform(self, X, y=None):
        """Fit and transform X to cluster-distance space.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : Ignored
            Not used, present for API consistency.

        Returns
        -------
        X_new : ndarray of shape (n_samples, n_clusters)
            X transformed to cluster-distance space.
        """
        return self.fit(X, y).transform(X)

    def score(self, X, y=None):
        """Return the negative inertia (for sklearn compatibility).

        The k-means objective is to minimize inertia, so we return
        the negative for compatibility with sklearn's scoring convention
        (higher is better).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data to score.
        y : Ignored
            Not used, present for API consistency.

        Returns
        -------
        score : float
            Negative of the inertia (sum of squared distances to nearest center).
        """
        check_is_fitted(self)
        X = check_array(X, dtype=np.float64)

        # Compute inertia
        labels = self.predict(X)
        inertia = 0.0
        for i in range(len(X)):
            diff = X[i] - self.cluster_centers_[labels[i]]
            inertia += np.dot(diff, diff)

        return -inertia

    @property
    def inertia_(self):
        """Sum of squared distances to closest center (k-means cost)."""
        check_is_fitted(self)
        return self._total_cost

    def __repr__(self):
        params = [f"n_clusters={self.n_clusters}"]
        if self.anns != 'lsh':
            params.append(f"anns='{self.anns}'")
        if self.anns_params is not None:
            params.append(f"anns_params={self.anns_params}")
        if self.max_rejections is not None:
            params.append(f"max_rejections={self.max_rejections}")
        if self.random_state is not None:
            params.append(f"random_state={self.random_state}")
        if self.n_jobs is not None:
            params.append(f"n_jobs={self.n_jobs}")
        return f"QKMeansSeeding({', '.join(params)})"
